Pin toolchains and split CLI orchestration

.github/workflows/ci.yml

@@ -13,9 +13,8 @@ jobs:
     timeout-minutes: 30
     steps:
       - uses: actions/checkout@v6
-      - uses: dtolnay/rust-toolchain@stable
-        with:
-          components: clippy,rustfmt
+      - name: Install pinned Rust
+        run: rustup toolchain install
       - uses: taiki-e/install-action@just
       - name: Install native build tools
         run: sudo apt-get update && sudo apt-get install -y autopoint clang cmake ninja-build
@@ -84,7 +83,8 @@ jobs:
             clang curl flex gawk libtool pkg-config python3
       - name: Install Rust
         run: |
-          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable
+          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none
+          "$HOME/.cargo/bin/rustup" toolchain install
           echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
       - uses: actions/download-artifact@v8
         with:

.github/workflows/release.yml

@@ -69,7 +69,8 @@ jobs:
             clang curl flex gawk libtool pkg-config python3
       - name: Install Rust
         run: |
-          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable
+          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none
+          "$HOME/.cargo/bin/rustup" toolchain install
           echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
       - uses: actions/download-artifact@v8
         with:

AGENTS.md

@@ -75,9 +75,16 @@
   digest, testing installation, and inspecting all shipped ELF compatibility.
 - Retry CI only when logs identify a transient infrastructure failure. Preserve
   real compiler, test, packaging, and compatibility failures for diagnosis.
+- Pin Rust in `rust-toolchain.toml`. Change the pin only in a dedicated change
+  and rerun lint and tests under the new compiler.
 
 ## Verification
 
+- Run `just` recipes inside the `xprobe-dev` Mamba environment from
+  `environment.yml`, for example `mamba run -n xprobe-dev just test`. A bare
+  host shell may lack Clang, Ninja, and autotools.
+- Set `PYTORCH_ENV` to an environment created from `environment-pytorch.yml`
+  for PyTorch live gates and `just test-release-live`.
 - Run `just fmt-check`, `just lint`, and `just test` for Rust changes.
 - Run `just test-bpf-live` for BPF attachment changes.
 - Run `just test-cupti-live` for CUPTI ABI or callback changes.

docs/development.md

@@ -2,7 +2,9 @@
 
 ## Native toolchain
 
-Rust is managed with `rustup`. Native build tools are isolated in the checked-in
+Rust is managed with `rustup` and pinned by `rust-toolchain.toml`; local
+builds, CI, and release packaging use the same compiler and Clippy. Update the
+pin in a dedicated change. Native build tools are isolated in the checked-in
 Mamba environment:
 
 ```bash
@@ -55,9 +57,14 @@ creating a tag, run the complete live suite on the local NVIDIA development
 machine:
 
 ```bash
-PYTORCH_ENV=/path/to/pytorch-env just test-release-live
+mamba env create --file environment-pytorch.yml
+PYTORCH_ENV="$(mamba run -n xprobe-pytorch python -c 'import sys; print(sys.prefix)')" \
+  just test-release-live
 ```
 
+`environment-pytorch.yml` pins the PyTorch fixture: Python 3.11, because
+torch 2.3 Dynamo does not support 3.12, plus Triton and NumPy 1.x.
+
 This local gate covers BPF and perf attachment, CUDA 12 and 13 behavior,
 injection, NVTX, mixed host/GPU collection, PyTorch, concurrent workers, and
 the release benchmarks. The tag workflow independently rebuilds both CUPTI
@@ -127,9 +134,10 @@ just test-pytorch-live
 
 By default the live recipes use a pinned NVIDIA PyTorch image with Ubuntu 24.04
 and CUDA 12.9, so live checks do not depend on the host Mamba environment.
-During local development, set `PYTORCH_ENV=/path/to/env` to mount
-an existing environment into the already-pinned CUDA fixtures instead of
-pulling the PyTorch image. Run eager matrix multiplication, convolution,
+During local development, set `PYTORCH_ENV` to the `xprobe-pytorch`
+environment to mount it into the already-pinned CUDA fixtures instead of
+pulling the PyTorch image. The compiled-Triton case checks Dynamo, Triton, and
+NumPy support before starting and fails with an explicit message otherwise. Run eager matrix multiplication, convolution,
 compiled Triton, bidirectional transfer, selected-kernel, and
 stream-synchronization profiling on the local GPU. The test also wraps an eager
 operation in an NVTX range and verifies exact range-ID matching:

environment-pytorch.yml

@@ -0,0 +1,17 @@
+name: xprobe-pytorch
+
+# Local PyTorch fixture for PYTORCH_ENV live tests and benchmarks. Dynamo
+# (torch.compile) in torch 2.3 supports Python 3.11 but not 3.12.
+channels:
+  - https://conda.anaconda.org/conda-forge
+  - nodefaults
+
+dependencies:
+  - python=3.11
+  - pip
+  - pip:
+      - --index-url https://download.pytorch.org/whl/cu121
+      - --extra-index-url https://pypi.org/simple
+      - torch==2.3.0
+      - triton==2.3.0
+      - numpy<2

rust-toolchain.toml

@@ -0,0 +1,4 @@
+[toolchain]
+channel = "1.99.1"
+profile = "minimal"
+components = ["clippy", "rustfmt"]

tests/integration/test_pytorch_cuda.py

@@ -73,10 +73,10 @@ def run_inner(xprobe: pathlib.Path) -> None:
         temp = pathlib.Path(temp_dir)
         agent = temp / "libxprobe-cupti.so"
         mode = temp / "mode"
+        require_compile_support()
         build_agent(workspace, agent)
         set_mode(mode, "mm")
-        environment = os.environ.copy()
-        environment["NVTX_INJECTION64_PATH"] = str(agent)
+        environment = workload_environment(agent)
         workload = subprocess.Popen(
             [
                 sys.executable,
@@ -188,6 +188,49 @@ def run_inner(xprobe: pathlib.Path) -> None:
     )
 
 
+def require_compile_support() -> None:
+    check = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import numpy, triton, torch._dynamo as dynamo; "
+            "raise SystemExit(0 if dynamo.is_dynamo_supported() else 3)",
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if check.returncode != 0:
+        raise SystemExit(
+            "PyTorch environment cannot run torch.compile; it needs numpy, triton, "
+            "and a Python version supported by Dynamo (see environment-pytorch.yml)\n"
+            + check.stderr
+        )
+
+
+def workload_environment(agent: pathlib.Path) -> dict[str, str]:
+    environment = os.environ.copy()
+    environment["NVTX_INJECTION64_PATH"] = str(agent)
+    # Triton links its launcher against an unversioned libcuda.so. Some container
+    # runtimes (for example WSL) register only libcuda.so.1, so fall back to the
+    # toolkit link stub; the real driver is still loaded through its SONAME.
+    stub = pathlib.Path("/usr/local/cuda/lib64/stubs/libcuda.so")
+    if (
+        "TRITON_LIBCUDA_PATH" not in environment
+        and not unversioned_libcuda_registered()
+        and stub.is_file()
+    ):
+        environment["TRITON_LIBCUDA_PATH"] = str(stub.parent)
+    return environment
+
+
+def unversioned_libcuda_registered() -> bool:
+    cache = subprocess.run(
+        ["/sbin/ldconfig", "-p"], check=True, capture_output=True, text=True
+    ).stdout
+    return any(line.split()[:1] == ["libcuda.so"] for line in cache.splitlines())
+
+
 def build_agent(workspace: pathlib.Path, output: pathlib.Path) -> None:
     subprocess.run(
         [