Disable XLA command buffers in JAX examples

[rostan-t]

Category:

Other (e.g. Documentation, Tests, Configuration)

Description:

PR #6286 changed the JAX training example to compile its jax.jit function before DALI runs. This doesn't actually fix the issue of XLA doing CUDA graph captures while DALI is running since even after the function is compiled, it's possible that XLA needs to capture new CUDA graphs.

This PR recommends setting the environment variable XLA_FLAGS to disable XLA command buffer.

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

• Existing tests apply
• New tests added
  • Python tests
  • GTests
  • Benchmark
  • Other
• N/A

Checklist

Documentation

• Existing documentation applies
• Documentation updated
  • Docstring
  • Doxygen
  • RST
  • Jupyter
  • Other
• N/A

DALI team only

Requirements

• Implements new requirements
• Affects existing requirements
• N/A

REQ IDs: N/A

JIRA TASK: N/A

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[rostan-t]

!build

[greptile-apps]

Greptile Summary

This PR updates the JAX examples and their notebook QA environments to disable XLA command buffers before JAX initializes, avoiding CUDA graph invalidation while DALI is active.

• Adds an explanation and early XLA_FLAGS assignment to seven JAX notebooks.
• Replaces the jax.jit warmup workaround in the basic example.
• Documents the command-buffer limitation on the JAX examples index.
• Exports the same setting from both JAX notebook test suites.

Confidence Score: 5/5

The PR appears safe to merge with no concrete blocking or independently actionable non-blocking issue identified.

The notebooks set the XLA option before importing JAX, spawned helper processes inherit it, and the QA harness starts fresh notebook kernels with the option already exported.

Important Files Changed

Filename	Overview
docs/examples/frameworks/jax/jax-basic_example.ipynb	Replaces the compilation warmup with an early command-buffer disable and initializes the model immediately before training.
docs/examples/custom_operations/jax_operator_multi_gpu.ipynb	Sets the XLA option before spawning the second process and before the notebook imports JAX.
docs/examples/frameworks/jax/index.rst	Documents why DALI buffer lifetime synchronization conflicts with XLA command buffers.
qa/TL0_jupyter/test_jax.sh	Exports the command-buffer setting before notebook execution.
qa/TL1_jupyter_conda/test_jax.sh	Exports the command-buffer setting before notebook execution in the Conda suite.

_{Reviews (1): Last reviewed commit: "Disable XLA command buffers in JAX examp..." | Re-trigger Greptile}

[dali-automaton]

CI MESSAGE: [62169436]: BUILD STARTED