Merge branch 'main' into aybchan/load-balance

Author: aybchan

.github/container/Dockerfile.base

@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04
+ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.06-cuda13.3-devel-ubuntu24.04
 ARG GIT_USER_NAME="JAX Toolbox"
 ARG GIT_USER_EMAIL=jax@nvidia.com
 ARG JAX_TOOLBOX_REF=main
@@ -133,3 +133,7 @@ ENV MANIFEST_FILE="/opt/manifest.d/manifest.yaml"
 ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/
 
 COPY patches/ /opt/manifest.d/patches/
+
+# Temporary workaround for TransformerEngine, which added a >=1.25.0 requirement
+# when no public cuda-dl-base container had that CUDNN_FRONTEND_VERSION.
+ENV CUDNN_FRONTEND_VERSION=1.25.0

.github/container/Dockerfile.jax

@@ -132,6 +132,7 @@ COPY --from=builder ${SRC_PATH_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE
 RUN <<"EOF" bash -ex
 ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl
 echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl)" > /opt/pip-tools.d/requirements-te.in
+echo "nvidia-cudnn-frontend==${CUDNN_FRONTEND_VERSION}" >> /opt/pip-tools.d/requirements-te.in
 EOF
 
 ## nvidia-cutlass-dsl[cu13]

.github/container/build-te.sh

@@ -138,13 +138,15 @@ pushd ${SRC_PATH_TE}
 # JAX, or the wheel-based installation of CUDA. Note that when we build TE as
 # part of building the JAX containers, JAX and XLA are not yet installed.
 python - << EOF
-import subprocess, sys, tomllib
+import os, subprocess, sys, tomllib
 with open("pyproject.toml", "rb") as ifile:
     data = tomllib.load(ifile)
 subprocess.run(
     [sys.executable, "-m", "pip", "install"]
     + [r for r in data["build-system"]["requires"]
-       if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")])
+       if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")]
+    + [f"nvidia-cudnn-frontend=={os.environ['CUDNN_FRONTEND_VERSION']}"]
+)
 EOF
 if [[ "${CCACHE}" == "1" ]]; then
     # Install ccache if not present (needs >= 4.1 for Redis remote storage support)

.github/triage/jax_toolbox_triage/args.py

@@ -98,6 +98,16 @@ def parse_args(args=None) -> argparse.Namespace:
             `date` will be substituted, e.g. ghcr.io/nvidia/jax:{container}-{date} for
             the JAX-Toolbox public nightlies.""",
     )
+    parser.add_argument("--metric-name", type=str, help="Example: tflops_per_sec")
+    parser.add_argument(
+        "--passing-metric",
+        action="append",
+        type=float,
+        help="Good value for the metric",
+    )
+    parser.add_argument(
+        "--failing-metric", action="append", type=float, help="Bad value for the metric"
+    )
     parser.add_argument(
         "--output-prefix",
         default=None,
@@ -316,6 +326,14 @@ def parse_args(args=None) -> argparse.Namespace:
         default="main",
         help="The name of the main branch (e.g. main) to derive cherry-picks from",
     )
+    parser.add_argument(
+        "--container-registry",
+        type=str,
+        help="""
+            Remote registry for use with the plugin backend. This can optionally
+            include a tag prefix, i.e. gitlab.com/USER/containers and
+            gitlab.com/USER/containers:PREFIX- are both valid.""",
+    )
     args = parser.parse_args(args=args)
     if args.restart:
         if args.output_prefix is None:
@@ -331,12 +349,38 @@ def parse_args(args=None) -> argparse.Namespace:
             args.output_prefix = pathlib.Path(
                 datetime.datetime.now().strftime("triage-%Y-%m-%d-%H-%M-%S")
             )
-
+    args.optional_software = optional_software.copy()
+    if args.exclude_transformer_engine:
+        args.optional_software.remove("transformer-engine")
     assert args.container_runtime in {
         "docker",
         "pyxis",
         "local",
+        "plugin",
     }, args.container_runtime
+
+    # Metric-based triage checks
+    container_search_options_passed = (
+        args.container is not None
+        or args.start_date is not None
+        or args.end_date is not None
+    )
+    if args.metric_name is not None:
+        if container_search_options_passed:
+            raise Exception(
+                "--metric-name is only supported in version-level search; use "
+                "--{passing,failing}-{container,versions}."
+            )
+        if not args.passing_metric or not args.failing_metric:
+            raise Exception(
+                "You should pass seed metric values via --passing-metric and "
+                "--failing-metric if --metric-name is passed."
+            )
+    if (args.passing_metric or args.failing_metric) and args.metric_name is None:
+        raise Exception(
+            "--metric-name must be passed if --passing-metric or --failing-metric is."
+        )
+
     args.workaround_buggy_container = set(args.workaround_buggy_container)
     # --{passing,failing}-commits are deprecated aliases for --{passing,failing}-versions.
     for prefix in ["passing", "failing"]:
@@ -393,9 +437,7 @@ def parse_args(args=None) -> argparse.Namespace:
         # If the container-level search is being skipped, because a valid combination
         # of --{passing,failing}-{versions,container} is passed, then no container-level
         # search options should be passed.
-        assert (
-            args.container is None and args.start_date is None and args.end_date is None
-        ), (
+        assert not container_search_options_passed, (
             "No container-level search options should be passed if the passing/failing"
             " containers/versions have been passed explicitly."
         )
@@ -417,12 +459,7 @@ def parse_args(args=None) -> argparse.Namespace:
     else:
         # None of --{passing,failing}-{versions,container} were passed, make sure the
         # compulsory arguments for the container-level search were passed
-        assert (
-            args.container is not None
-        ), "--container must be passed for the container-level search"
-
-    args.optional_software = optional_software.copy()
-    if args.exclude_transformer_engine:
-        args.optional_software.remove("transformer-engine")
-
+        assert args.container is not None, (
+            "--container must be passed for the container-level search"
+        )
     return args

.github/triage/jax_toolbox_triage/data_files/Dockerfile.triage-tool

@@ -0,0 +1,9 @@
+# syntax=docker/dockerfile:1
+ARG BASE_IMAGE
+ARG BUILD_CMD
+ARG GIT_CMD
+FROM ${BASE_IMAGE}
+ARG BUILD_CMD
+ARG GIT_CMD
+RUN sh -c "${GIT_CMD}"
+RUN --mount=type=cache,target=/root/.cache sh -c "${BUILD_CMD}"

.github/triage/jax_toolbox_triage/docker.py

@@ -24,10 +24,15 @@ def __init__(
 
     def __enter__(self):
         self._logger.debug(f"Launching {self}")
-        have_gpus = shutil.which("nvidia-smi") is not None
+        have_gpus = (
+            shutil.which("nvidia-smi") is not None
+            and subprocess.run(["nvidia-smi"]).returncode == 0
+        )
         if not have_gpus:
             self._logger.warning("No GPUs detected!")
-        gpu_args = ["--gpus=all"] if have_gpus else []
+        gpu_args = (
+            ["--gpus=all"] if have_gpus else ["-e", "NVIDIA_VISIBLE_DEVICES=void"]
+        )
         result = run_and_log(
             [
                 "docker",
@@ -93,6 +98,7 @@ def exists(self) -> bool:
         """
         Check if the given container exists.
         """
+        # TODO: check without pulling; docker manifest inspect
         result = run_and_log(
             ["docker", "pull", self._url],
             logger=self._logger,