Cherrypick: Add the RunLLM widget to the website (#159592)

Add the RunLLM widget to the website (#152055)

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152055
Approved by: https://github.com/albanD

(cherry picked from commit e4e2701429c17078c3c475382a8b1fa4c8a8cefc)

.bazelrc

@@ -2,7 +2,7 @@ build --cxxopt=--std=c++17
 build --copt=-I.
 # Bazel does not support including its cc_library targets as system
 # headers. We work around this for generated code
-# (e.g. torch/headeronly/macros/cmake_macros.h) by making the generated directory a
+# (e.g. c10/macros/cmake_macros.h) by making the generated directory a
 # system include path.
 build --copt=-isystem --copt bazel-out/k8-fastbuild/bin
 build --copt=-isystem --copt bazel-out/darwin-fastbuild/bin

.bc-linter.yml

@@ -1,15 +0,0 @@
-version: 1
-paths:
-include:
-  - "**/*.py"
-exclude:
-  - ".*"
-  - ".*/**"
-  - "**/.*/**"
-  - "**/.*"
-  - "**/_*/**"
-  - "**/_*.py"
-  - "**/test/**"
-  - "**/benchmarks/**"
-  - "**/test_*.py"
-  - "**/*_test.py"

.ci/aarch64_linux/aarch64_ci_build.sh

@@ -3,20 +3,8 @@ set -eux -o pipefail
 
 GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-}
 
-# Set CUDA architecture lists to match x86 build_cuda.sh
-if [[ "$GPU_ARCH_VERSION" == *"12.6"* ]]; then
-    export TORCH_CUDA_ARCH_LIST="8.0;9.0"
-elif [[ "$GPU_ARCH_VERSION" == *"12.8"* ]]; then
+if [[ "$GPU_ARCH_VERSION" == *"12.9"* ]]; then
     export TORCH_CUDA_ARCH_LIST="8.0;9.0;10.0;12.0"
-elif [[ "$GPU_ARCH_VERSION" == *"13.0"* ]]; then
-    export TORCH_CUDA_ARCH_LIST="8.0;9.0;10.0;11.0;12.0+PTX"
-fi
-
-# Compress the fatbin with -compress-mode=size for CUDA 13
-if [[ "$DESIRED_CUDA" == *"13"* ]]; then
-    export TORCH_NVCC_FLAGS="-compress-mode=size"
-    # Bundle ptxas into the cu13 wheel, see https://github.com/pytorch/pytorch/issues/163801
-    export BUILD_BUNDLE_PTXAS=1
 fi
 
 SCRIPTPATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
@@ -30,22 +18,14 @@ cd /
 # on the mounted pytorch repo
 git config --global --add safe.directory /pytorch
 pip install -r /pytorch/requirements.txt
-pip install auditwheel==6.2.0 wheel
+pip install auditwheel==6.2.0
 if [ "$DESIRED_CUDA" = "cpu" ]; then
     echo "BASE_CUDA_VERSION is not set. Building cpu wheel."
-    python /pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn
+    #USE_PRIORITIZED_TEXT_FOR_LD for enable linker script optimization https://github.com/pytorch/pytorch/pull/121975/files
+    USE_PRIORITIZED_TEXT_FOR_LD=1 python /pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn
 else
     echo "BASE_CUDA_VERSION is set to: $DESIRED_CUDA"
     export USE_SYSTEM_NCCL=1
-
-    # Check if we should use NVIDIA libs from PyPI (similar to x86 build_cuda.sh logic)
-    if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
-        echo "Bundling CUDA libraries with wheel for aarch64."
-    else
-        echo "Using nvidia libs from pypi for aarch64."
-        echo "Updated PYTORCH_EXTRA_INSTALL_REQUIREMENTS for aarch64: $PYTORCH_EXTRA_INSTALL_REQUIREMENTS"
-        export USE_NVIDIA_PYPI_LIBS=1
-    fi
-
-    python /pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn --enable-cuda
+    #USE_PRIORITIZED_TEXT_FOR_LD for enable linker script optimization https://github.com/pytorch/pytorch/pull/121975/files
+    USE_PRIORITIZED_TEXT_FOR_LD=1 python /pytorch/.ci/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn --enable-cuda
 fi

.ci/aarch64_linux/aarch64_wheel_ci_build.py

@@ -69,186 +69,60 @@ def replace_tag(filename) -> None:
         f.writelines(lines)
 
 
-def patch_library_rpath(
-    folder: str,
-    lib_name: str,
-    use_nvidia_pypi_libs: bool = False,
-    desired_cuda: str = "",
-) -> None:
-    """Apply patchelf to set RPATH for a library in torch/lib"""
-    lib_path = f"{folder}/tmp/torch/lib/{lib_name}"
-
-    if use_nvidia_pypi_libs:
-        # For PyPI NVIDIA libraries, construct CUDA RPATH
-        cuda_rpaths = [
-            "$ORIGIN/../../nvidia/cudnn/lib",
-            "$ORIGIN/../../nvidia/nvshmem/lib",
-            "$ORIGIN/../../nvidia/nccl/lib",
-            "$ORIGIN/../../nvidia/cusparselt/lib",
-        ]
-
-        if "130" in desired_cuda:
-            cuda_rpaths.append("$ORIGIN/../../nvidia/cu13/lib")
-        else:
-            cuda_rpaths.extend(
-                [
-                    "$ORIGIN/../../nvidia/cublas/lib",
-                    "$ORIGIN/../../nvidia/cuda_cupti/lib",
-                    "$ORIGIN/../../nvidia/cuda_nvrtc/lib",
-                    "$ORIGIN/../../nvidia/cuda_runtime/lib",
-                    "$ORIGIN/../../nvidia/cufft/lib",
-                    "$ORIGIN/../../nvidia/curand/lib",
-                    "$ORIGIN/../../nvidia/cusolver/lib",
-                    "$ORIGIN/../../nvidia/cusparse/lib",
-                    "$ORIGIN/../../nvidia/nvtx/lib",
-                    "$ORIGIN/../../nvidia/cufile/lib",
-                ]
-            )
-
-        # Add $ORIGIN for local torch libs
-        rpath = ":".join(cuda_rpaths) + ":$ORIGIN"
-    else:
-        # For bundled libraries, just use $ORIGIN
-        rpath = "$ORIGIN"
-
-    if os.path.exists(lib_path):
-        os.system(
-            f"cd {folder}/tmp/torch/lib/; "
-            f"patchelf --set-rpath '{rpath}' --force-rpath {lib_name}"
-        )
-
-
-def copy_and_patch_library(
-    src_path: str,
-    folder: str,
-    use_nvidia_pypi_libs: bool = False,
-    desired_cuda: str = "",
-) -> None:
-    """Copy a library to torch/lib and patch its RPATH"""
-    if os.path.exists(src_path):
-        lib_name = os.path.basename(src_path)
-        shutil.copy2(src_path, f"{folder}/tmp/torch/lib/{lib_name}")
-        patch_library_rpath(folder, lib_name, use_nvidia_pypi_libs, desired_cuda)
-
-
 def package_cuda_wheel(wheel_path, desired_cuda) -> None:
     """
     Package the cuda wheel libraries
     """
     folder = os.path.dirname(wheel_path)
+    wheelname = os.path.basename(wheel_path)
     os.mkdir(f"{folder}/tmp")
     os.system(f"unzip {wheel_path} -d {folder}/tmp")
-    # Delete original wheel since it will be repackaged
-    os.system(f"rm {wheel_path}")
+    libs_to_copy = [
+        "/usr/local/cuda/extras/CUPTI/lib64/libcupti.so.12",
+        "/usr/local/cuda/lib64/libcudnn.so.9",
+        "/usr/local/cuda/lib64/libcublas.so.12",
+        "/usr/local/cuda/lib64/libcublasLt.so.12",
+        "/usr/local/cuda/lib64/libcudart.so.12",
+        "/usr/local/cuda/lib64/libcufft.so.11",
+        "/usr/local/cuda/lib64/libcusparse.so.12",
+        "/usr/local/cuda/lib64/libcusparseLt.so.0",
+        "/usr/local/cuda/lib64/libcusolver.so.11",
+        "/usr/local/cuda/lib64/libcurand.so.10",
+        "/usr/local/cuda/lib64/libnccl.so.2",
+        "/usr/local/cuda/lib64/libnvJitLink.so.12",
+        "/usr/local/cuda/lib64/libnvrtc.so.12",
+        "/usr/local/cuda/lib64/libcudnn_adv.so.9",
+        "/usr/local/cuda/lib64/libcudnn_cnn.so.9",
+        "/usr/local/cuda/lib64/libcudnn_graph.so.9",
+        "/usr/local/cuda/lib64/libcudnn_ops.so.9",
+        "/usr/local/cuda/lib64/libcudnn_engines_runtime_compiled.so.9",
+        "/usr/local/cuda/lib64/libcudnn_engines_precompiled.so.9",
+        "/usr/local/cuda/lib64/libcudnn_heuristic.so.9",
+        "/lib64/libgomp.so.1",
+        "/usr/lib64/libgfortran.so.5",
+        "/acl/build/libarm_compute.so",
+        "/acl/build/libarm_compute_graph.so",
+        "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0",
+        "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0",
+        "/usr/local/lib/libnvpl_lapack_core.so.0",
+        "/usr/local/lib/libnvpl_blas_core.so.0",
+    ]
 
-    # Check if we should use PyPI NVIDIA libraries or bundle system libraries
-    use_nvidia_pypi_libs = os.getenv("USE_NVIDIA_PYPI_LIBS", "0") == "1"
-
-    if use_nvidia_pypi_libs:
-        print("Using nvidia libs from pypi - skipping CUDA library bundling")
-        # For PyPI approach, we don't bundle CUDA libraries - they come from PyPI packages
-        # We only need to bundle non-NVIDIA libraries
-        minimal_libs_to_copy = [
-            "/lib64/libgomp.so.1",
-            "/usr/lib64/libgfortran.so.5",
-            "/acl/build/libarm_compute.so",
-            "/acl/build/libarm_compute_graph.so",
-            "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0",
-            "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0",
-            "/usr/local/lib/libnvpl_lapack_core.so.0",
-            "/usr/local/lib/libnvpl_blas_core.so.0",
-        ]
-
-        # Copy minimal libraries to unzipped_folder/torch/lib
-        for lib_path in minimal_libs_to_copy:
-            copy_and_patch_library(lib_path, folder, use_nvidia_pypi_libs, desired_cuda)
-
-        # Patch torch libraries used for searching libraries
-        torch_libs_to_patch = [
-            "libtorch.so",
-            "libtorch_cpu.so",
-            "libtorch_cuda.so",
-            "libtorch_cuda_linalg.so",
-            "libtorch_global_deps.so",
-            "libtorch_python.so",
-            "libtorch_nvshmem.so",
-            "libc10.so",
-            "libc10_cuda.so",
-            "libcaffe2_nvrtc.so",
-            "libshm.so",
-        ]
-        for lib_name in torch_libs_to_patch:
-            patch_library_rpath(folder, lib_name, use_nvidia_pypi_libs, desired_cuda)
-    else:
-        print("Bundling CUDA libraries with wheel")
-        # Original logic for bundling system CUDA libraries
-        # Common libraries for all CUDA versions
-        common_libs = [
-            # Non-NVIDIA system libraries
-            "/lib64/libgomp.so.1",
-            "/usr/lib64/libgfortran.so.5",
-            "/acl/build/libarm_compute.so",
-            "/acl/build/libarm_compute_graph.so",
-            # Common CUDA libraries (same for all versions)
-            "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0",
-            "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0",
-            "/usr/local/lib/libnvpl_lapack_core.so.0",
-            "/usr/local/lib/libnvpl_blas_core.so.0",
-            "/usr/local/cuda/extras/CUPTI/lib64/libnvperf_host.so",
-            "/usr/local/cuda/lib64/libcudnn.so.9",
-            "/usr/local/cuda/lib64/libcusparseLt.so.0",
-            "/usr/local/cuda/lib64/libcurand.so.10",
-            "/usr/local/cuda/lib64/libnccl.so.2",
-            "/usr/local/cuda/lib64/libnvshmem_host.so.3",
-            "/usr/local/cuda/lib64/libcudnn_adv.so.9",
-            "/usr/local/cuda/lib64/libcudnn_cnn.so.9",
-            "/usr/local/cuda/lib64/libcudnn_graph.so.9",
-            "/usr/local/cuda/lib64/libcudnn_ops.so.9",
-            "/usr/local/cuda/lib64/libcudnn_engines_runtime_compiled.so.9",
-            "/usr/local/cuda/lib64/libcudnn_engines_precompiled.so.9",
-            "/usr/local/cuda/lib64/libcudnn_heuristic.so.9",
+    if "129" in desired_cuda:
+        libs_to_copy += [
+            "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.9",
             "/usr/local/cuda/lib64/libcufile.so.0",
             "/usr/local/cuda/lib64/libcufile_rdma.so.1",
-            "/usr/local/cuda/lib64/libcusparse.so.12",
         ]
 
-        # CUDA version-specific libraries
-        if "13" in desired_cuda:
-            minor_version = desired_cuda[-1]
-            version_specific_libs = [
-                "/usr/local/cuda/extras/CUPTI/lib64/libcupti.so.13",
-                "/usr/local/cuda/lib64/libcublas.so.13",
-                "/usr/local/cuda/lib64/libcublasLt.so.13",
-                "/usr/local/cuda/lib64/libcudart.so.13",
-                "/usr/local/cuda/lib64/libcufft.so.12",
-                "/usr/local/cuda/lib64/libcusolver.so.12",
-                "/usr/local/cuda/lib64/libnvJitLink.so.13",
-                "/usr/local/cuda/lib64/libnvrtc.so.13",
-                f"/usr/local/cuda/lib64/libnvrtc-builtins.so.13.{minor_version}",
-            ]
-        elif "12" in desired_cuda:
-            # Get the last character for libnvrtc-builtins version (e.g., "129" -> "9")
-            minor_version = desired_cuda[-1]
-            version_specific_libs = [
-                "/usr/local/cuda/extras/CUPTI/lib64/libcupti.so.12",
-                "/usr/local/cuda/lib64/libcublas.so.12",
-                "/usr/local/cuda/lib64/libcublasLt.so.12",
-                "/usr/local/cuda/lib64/libcudart.so.12",
-                "/usr/local/cuda/lib64/libcufft.so.11",
-                "/usr/local/cuda/lib64/libcusolver.so.11",
-                "/usr/local/cuda/lib64/libnvJitLink.so.12",
-                "/usr/local/cuda/lib64/libnvrtc.so.12",
-                f"/usr/local/cuda/lib64/libnvrtc-builtins.so.12.{minor_version}",
-            ]
-        else:
-            raise ValueError(f"Unsupported CUDA version: {desired_cuda}.")
-
-        # Combine all libraries
-        libs_to_copy = common_libs + version_specific_libs
-
-        # Copy libraries to unzipped_folder/torch/lib
-        for lib_path in libs_to_copy:
-            copy_and_patch_library(lib_path, folder, use_nvidia_pypi_libs, desired_cuda)
+    # Copy libraries to unzipped_folder/a/lib
+    for lib_path in libs_to_copy:
+        lib_name = os.path.basename(lib_path)
+        shutil.copy2(lib_path, f"{folder}/tmp/torch/lib/{lib_name}")
+        os.system(
+            f"cd {folder}/tmp/torch/lib/; "
+            f"patchelf --set-rpath '$ORIGIN' --force-rpath {folder}/tmp/torch/lib/{lib_name}"
+        )
 
     # Make sure the wheel is tagged with manylinux_2_28
     for f in os.scandir(f"{folder}/tmp/"):
@@ -256,8 +130,14 @@ def package_cuda_wheel(wheel_path, desired_cuda) -> None:
             replace_tag(f"{f.path}/WHEEL")
             break
 
-    os.system(f"wheel pack {folder}/tmp/ -d {folder}")
-    os.system(f"rm -rf {folder}/tmp/")
+    os.mkdir(f"{folder}/cuda_wheel")
+    os.system(f"cd {folder}/tmp/; zip -r {folder}/cuda_wheel/{wheelname} *")
+    shutil.move(
+        f"{folder}/cuda_wheel/{wheelname}",
+        f"{folder}/{wheelname}",
+        copy_function=shutil.copy2,
+    )
+    os.system(f"rm -rf {folder}/tmp/ {folder}/cuda_wheel/")
 
 
 def complete_wheel(folder: str) -> str:
@@ -280,7 +160,14 @@ def complete_wheel(folder: str) -> str:
             f"/{folder}/dist/{repaired_wheel_name}",
         )
     else:
-        repaired_wheel_name = list_dir(f"/{folder}/dist")[0]
+        repaired_wheel_name = wheel_name.replace(
+            "linux_aarch64", "manylinux_2_28_aarch64"
+        )
+        print(f"Renaming {wheel_name} wheel to {repaired_wheel_name}")
+        os.rename(
+            f"/{folder}/dist/{wheel_name}",
+            f"/{folder}/dist/{repaired_wheel_name}",
+        )
 
     print(f"Copying {repaired_wheel_name} to artifacts")
     shutil.copy2(
@@ -317,20 +204,10 @@ if __name__ == "__main__":
     ).decode()
 
     print("Building PyTorch wheel")
-    build_vars = ""
+    build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
     # MAX_JOB=5 is not required for CPU backend (see commit 465d98b)
     if enable_cuda:
-        build_vars += "MAX_JOBS=5 "
-
-        # Handle PyPI NVIDIA libraries vs bundled libraries
-        use_nvidia_pypi_libs = os.getenv("USE_NVIDIA_PYPI_LIBS", "0") == "1"
-        if use_nvidia_pypi_libs:
-            print("Configuring build for PyPI NVIDIA libraries")
-            # Configure for dynamic linking (matching x86 logic)
-            build_vars += "ATEN_STATIC_CUDA=0 USE_CUDA_STATIC_LINK=0 USE_CUPTI_SO=1 "
-        else:
-            print("Configuring build for bundled NVIDIA libraries")
-            # Keep existing static linking approach - already configured above
+        build_vars = "MAX_JOBS=5 " + build_vars
 
     override_package_version = os.getenv("OVERRIDE_PACKAGE_VERSION")
     desired_cuda = os.getenv("DESIRED_CUDA")
@@ -372,7 +249,7 @@ if __name__ == "__main__":
     else:
         print("build pytorch without mkldnn backend")
 
-    os.system(f"cd /pytorch; {build_vars} python3 -m build --wheel --no-isolation")
+    os.system(f"cd /pytorch; {build_vars} python3 setup.py bdist_wheel")
     if enable_cuda:
         print("Updating Cuda Dependency")
         filename = os.listdir("/pytorch/dist/")

.ci/aarch64_linux/build_aarch64_wheel.py

@@ -241,7 +241,7 @@ def wait_for_connection(addr, port, timeout=15, attempt_cnt=5):
         try:
             with socket.create_connection((addr, port), timeout=timeout):
                 return
-        except (ConnectionRefusedError, TimeoutError):  # noqa: PERF203
+        except (ConnectionRefusedError, socket.timeout):  # noqa: PERF203
             if i == attempt_cnt - 1:
                 raise
             time.sleep(timeout)
@@ -438,11 +438,13 @@ def build_torchvision(
         )
         build_vars += f"BUILD_VERSION={version}.dev{build_date}"
     elif build_version is not None:
-        build_vars += f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-', maxsplit=1)[0]}"
+        build_vars += (
+            f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-')[0]}"
+        )
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
 
-    host.run_cmd(f"cd vision && {build_vars} python3 -m build --wheel --no-isolation")
+    host.run_cmd(f"cd vision && {build_vars} python3 setup.py bdist_wheel")
     vision_wheel_name = host.list_dir("vision/dist")[0]
     embed_libgomp(host, use_conda, os.path.join("vision", "dist", vision_wheel_name))
 
@@ -493,11 +495,13 @@ def build_torchdata(
         )
         build_vars += f"BUILD_VERSION={version}.dev{build_date}"
     elif build_version is not None:
-        build_vars += f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-', maxsplit=1)[0]}"
+        build_vars += (
+            f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-')[0]}"
+        )
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
 
-    host.run_cmd(f"cd data && {build_vars} python3 -m build --wheel --no-isolation")
+    host.run_cmd(f"cd data && {build_vars} python3 setup.py bdist_wheel")
     wheel_name = host.list_dir("data/dist")[0]
     embed_libgomp(host, use_conda, os.path.join("data", "dist", wheel_name))
 
@@ -549,11 +553,13 @@ def build_torchtext(
         )
         build_vars += f"BUILD_VERSION={version}.dev{build_date}"
     elif build_version is not None:
-        build_vars += f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-', maxsplit=1)[0]}"
+        build_vars += (
+            f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-')[0]}"
+        )
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
 
-    host.run_cmd(f"cd text && {build_vars} python3 -m build --wheel --no-isolation")
+    host.run_cmd(f"cd text && {build_vars} python3 setup.py bdist_wheel")
     wheel_name = host.list_dir("text/dist")[0]
     embed_libgomp(host, use_conda, os.path.join("text", "dist", wheel_name))
 
@@ -607,14 +613,16 @@ def build_torchaudio(
         )
         build_vars += f"BUILD_VERSION={version}.dev{build_date}"
     elif build_version is not None:
-        build_vars += f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-', maxsplit=1)[0]}"
+        build_vars += (
+            f"BUILD_VERSION={build_version} PYTORCH_VERSION={branch[1:].split('-')[0]}"
+        )
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
 
     host.run_cmd(
         f"cd audio && export FFMPEG_ROOT=$(pwd)/third_party/ffmpeg && export USE_FFMPEG=1 \
         && ./packaging/ffmpeg/build.sh \
-        && {build_vars} python3 -m build --wheel --no-isolation"
+        && {build_vars} python3 setup.py bdist_wheel"
     )
 
     wheel_name = host.list_dir("audio/dist")[0]
@@ -726,7 +734,7 @@ def start_build(
     print("Building PyTorch wheel")
     build_opts = ""
     if pytorch_build_number is not None:
-        build_opts += f" -C--build-option=--build-number={pytorch_build_number}"
+        build_opts += f" --build-number {pytorch_build_number}"
     # Breakpad build fails on aarch64
     build_vars = "USE_BREAKPAD=0 "
     if branch == "nightly":
@@ -747,8 +755,7 @@ def start_build(
         print("build pytorch with mkldnn+acl backend")
         build_vars += " USE_MKLDNN=ON USE_MKLDNN_ACL=ON"
         host.run_cmd(
-            f"cd $HOME/pytorch && export ACL_ROOT_DIR=$HOME/ComputeLibrary && "
-            f"{build_vars} python3 -m build --wheel --no-isolation{build_opts}"
+            f"cd $HOME/pytorch && export ACL_ROOT_DIR=$HOME/ComputeLibrary && {build_vars} python3 setup.py bdist_wheel{build_opts}"
         )
         print("Repair the wheel")
         pytorch_wheel_name = host.list_dir("pytorch/dist")[0]
@@ -764,7 +771,7 @@ def start_build(
     else:
         print("build pytorch without mkldnn backend")
         host.run_cmd(
-            f"cd pytorch && {build_vars} python3 -m build --wheel --no-isolation{build_opts}"
+            f"cd pytorch && {build_vars} python3 setup.py bdist_wheel{build_opts}"
         )
 
     print("Deleting build folder")
@@ -1005,7 +1012,7 @@ if __name__ == "__main__":
         install_condaforge_python(host, args.python_version)
         sys.exit(0)
 
-    python_version = args.python_version if args.python_version is not None else "3.10"
+    python_version = args.python_version if args.python_version is not None else "3.9"
 
     if args.use_torch_from_pypi:
         configure_system(host, compiler=args.compiler, python_version=python_version)

.ci/caffe2/test.sh

@@ -5,7 +5,7 @@ source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 
 if [[ ${BUILD_ENVIRONMENT} == *onnx* ]]; then
   pip install click mock tabulate networkx==2.0
-  pip -q install "file:///var/lib/jenkins/workspace/third_party/onnx#egg=onnx"
+  pip -q install --user "file:///var/lib/jenkins/workspace/third_party/onnx#egg=onnx"
 fi
 
 # Skip tests in environments where they are not built/applicable
@@ -147,8 +147,8 @@ export DNNL_MAX_CPU_ISA=AVX2
 if [[ "${SHARD_NUMBER:-1}" == "1" ]]; then
   # TODO(sdym@meta.com) remove this when the linked issue resolved.
   # py is temporary until https://github.com/Teemu/pytest-sugar/issues/241 is fixed
-  pip install py==1.11.0
-  pip install pytest-sugar
+  pip install --user py==1.11.0
+  pip install --user pytest-sugar
   # NB: Warnings are disabled because they make it harder to see what
   # the actual erroring test is
   "$PYTHON" \

.ci/docker/README.md

@@ -36,104 +36,3 @@ See `build.sh` for valid build environments (it's the giant switch).
 # Set flags (see build.sh) and build image
 sudo bash -c 'TRITON=1 ./build.sh pytorch-linux-bionic-py3.8-gcc9 -t myimage:latest
 ```
-
-## [Guidance] Adding a New Base Docker Image
-
-### Background
-
-The base Docker images in directory `.ci/docker/` are built by the `docker-builds.yml` workflow. Those images are used throughout the PyTorch CI/CD pipeline. You should only create or modify a base Docker image if you need specific environment changes or dependencies before building PyTorch on CI.
-
-1. **Automatic Rebuilding**:
-   - The Docker image building process is triggered automatically when changes are made to files in the `.ci/docker/*` directory
-   - This ensures all images stay up-to-date with the latest dependencies and configurations
-
-2. **Image Reuse in PyTorch Build Workflows** (example: linux-build):
-   - The images generated by `docker-builds.yml` are reused in `_linux-build.yml` through the `calculate-docker-image` step
-   - The `_linux-build.yml` workflow:
-     - Pulls the Docker image determined by the `calculate-docker-image` step
-     - Runs a Docker container with that image
-     - Executes `.ci/pytorch/build.sh` inside the container to build PyTorch
-
-3. **Usage in Test Workflows** (example: linux-test):
-   - The same Docker images are also used in `_linux-test.yml` for running tests
-   - The `_linux-test.yml` workflow follows a similar pattern:
-     - It uses the `calculate-docker-image` step to determine which Docker image to use
-     - It pulls the Docker image and runs a container with that image
-     - It installs the wheels from the artifacts generated by PyTorch build jobs
-     - It executes test scripts (like `.ci/pytorch/test.sh` or `.ci/pytorch/multigpu-test.sh`) inside the container
-
-### Understanding File Purposes
-
-#### `.ci/docker/build.sh` vs `.ci/pytorch/build.sh`
-- **`.ci/docker/build.sh`**:
-  - Used for building base Docker images
-  - Executed by the `docker-builds.yml` workflow to pre-build Docker images for CI
-  - Contains configurations for different Docker build environments
-
-- **`.ci/pytorch/build.sh`**:
-  - Used for building PyTorch inside a Docker container
-  - Called by workflows like `_linux-build.yml` after the Docker container is started
-  - Builds PyTorch wheels and other artifacts
-
-#### `.ci/docker/ci_commit_pins/` vs `.github/ci_commit_pins`
-- **`.ci/docker/ci_commit_pins/`**:
-  - Used for pinning dependency versions during base Docker image building
-  - Ensures consistent environments for building PyTorch
-  - Changes here trigger base Docker image rebuilds
-
-- **`.github/ci_commit_pins`**:
-  - Used for pinning dependency versions during PyTorch building and tests
-  - Ensures consistent dependencies for PyTorch across different builds
-  - Used by build scripts running inside Docker containers
-
-### Step-by-Step Guide for Adding a New Base Docker Image
-
-#### 1. Add Pinned Commits (If Applicable)
-
-We use pinned commits for build stability. The `nightly.yml` workflow checks and updates pinned commits for certain repository dependencies daily.
-
-If your new Docker image needs a library installed from a specific pinned commit or built from source:
-
-1. Add the repository you want to track in `nightly.yml` and `merge-rules.yml`
-2. Add the initial pinned commit in `.ci/docker/ci_commit_pins/`. The text filename should match the one defined in step 1
-
-#### 2. Configure the Base Docker Image
-1. **Add new Base Docker image configuration** (if applicable):
-
-   Add the configuration in `.ci/docker/build.sh`. For example:
-   ```bash
-   pytorch-linux-jammy-cuda12.8-cudnn9-py3.12-gcc11-new1)
-     CUDA_VERSION=12.8.1
-     ANACONDA_PYTHON_VERSION=3.12
-     GCC_VERSION=11
-     VISION=yes
-     KATEX=yes
-     UCX_COMMIT=${_UCX_COMMIT}
-     UCC_COMMIT=${_UCC_COMMIT}
-     TRITON=yes
-     NEW_ARG_1=yes
-     ;;
-   ```
-
-2. **Add build arguments to Docker build command**:
-
-   If you're introducing a new argument to the Docker build, make sure to add it in the Docker build step in `.ci/docker/build.sh`:
-   ```bash
-   docker build \
-     ....
-     --build-arg "NEW_ARG_1=${NEW_ARG_1}"
-   ```
-
-3. **Update Dockerfile logic**:
-
-   Update the Dockerfile to use the new argument. For example, in `ubuntu/Dockerfile`:
-   ```dockerfile
-   ARG NEW_ARG_1
-   # Set up environment for NEW_ARG_1
-   RUN if [ -n "${NEW_ARG_1}" ]; then bash ./do_something.sh; fi
-   ```
-
-4. **Add the Docker configuration** in `.github/workflows/docker-builds.yml`:
-
-   The `docker-builds.yml` workflow pre-builds the Docker images whenever changes occur in the `.ci/docker/` directory. This includes the
-   pinned commit updates.

.ci/docker/almalinux/Dockerfile

@@ -64,13 +64,8 @@ FROM cuda as cuda12.9
 RUN bash ./install_cuda.sh 12.9
 ENV DESIRED_CUDA=12.9
 
-FROM cuda as cuda13.0
-RUN bash ./install_cuda.sh 13.0
-ENV DESIRED_CUDA=13.0
-
 FROM ${ROCM_IMAGE} as rocm
-ARG PYTORCH_ROCM_ARCH
-ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
+ENV PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
 ADD ./common/install_mkl.sh install_mkl.sh
 RUN bash ./install_mkl.sh && rm install_mkl.sh
 ENV MKLROOT /opt/intel
@@ -81,10 +76,10 @@ ADD ./common/install_mnist.sh install_mnist.sh
 RUN bash ./install_mnist.sh
 
 FROM base as all_cuda
+COPY --from=cuda11.8  /usr/local/cuda-11.8 /usr/local/cuda-11.8
 COPY --from=cuda12.6  /usr/local/cuda-12.6 /usr/local/cuda-12.6
 COPY --from=cuda12.8  /usr/local/cuda-12.8 /usr/local/cuda-12.8
 COPY --from=cuda12.9  /usr/local/cuda-12.9 /usr/local/cuda-12.9
-COPY --from=cuda13.0  /usr/local/cuda-13.0 /usr/local/cuda-13.0
 
 # Final step
 FROM ${BASE_TARGET} as final

.ci/docker/almalinux/build.sh

@@ -36,12 +36,6 @@ case ${DOCKER_TAG_PREFIX} in
     ;;
   rocm*)
     BASE_TARGET=rocm
-    PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
-    # add gfx950 conditionally starting in ROCm 7.0
-    if [[ "$ROCM_VERSION" == *"7.0"* ]]; then
-        PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH};gfx950"
-    fi
-    EXTRA_BUILD_ARGS="${EXTRA_BUILD_ARGS} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}"
     ;;
   *)
     echo "ERROR: Unknown docker tag ${DOCKER_TAG_PREFIX}"

.ci/docker/build.sh

@@ -52,8 +52,6 @@ fi
 
 if [[ "$image" == *-jammy* ]]; then
   UBUNTU_VERSION=22.04
-elif [[ "$image" == *-noble* ]]; then
-  UBUNTU_VERSION=24.04
 elif [[ "$image" == *ubuntu* ]]; then
   extract_version_from_image_name ubuntu UBUNTU_VERSION
 fi
@@ -76,16 +74,13 @@ elif [[ "$image" == *cuda*linter* ]]; then
 elif [[ "$image" == *linter* ]]; then
   # Use a separate Dockerfile for linter to keep a small image size
   DOCKERFILE="linter/Dockerfile"
-elif [[ "$image" == *riscv* ]]; then
-  # Use RISC-V specific Dockerfile
-  DOCKERFILE="ubuntu-cross-riscv/Dockerfile"
 fi
 
-_UCX_COMMIT=7836b165abdbe468a2f607e7254011c07d788152
-_UCC_COMMIT=430e241bf5d38cbc73fc7a6b89155397232e3f96
+_UCX_COMMIT=7bb2722ff2187a0cad557ae4a6afa090569f83fb
+_UCC_COMMIT=20eae37090a4ce1b32bcce6144ccad0b49943e0b
 if [[ "$image" == *rocm* ]]; then
-  _UCX_COMMIT=29831d319e6be55cb8c768ca61de335c934ca39e
-  _UCC_COMMIT=9f4b242cbbd8b1462cbc732eb29316cdfa124b77
+  _UCX_COMMIT=cc312eaa4655c0cc5c2bcd796db938f90563bcf6
+  _UCC_COMMIT=0c0fc21559835044ab107199e334f7157d6a0d3d
 fi
 
 tag=$(echo $image | awk -F':' '{print $2}')
@@ -94,28 +89,9 @@ tag=$(echo $image | awk -F':' '{print $2}')
 # configuration, so we hardcode everything here rather than do it
 # from scratch
 case "$tag" in
-  pytorch-linux-jammy-cuda12.4-cudnn9-py3-gcc11)
-    CUDA_VERSION=12.4
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=11
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    TRITON=yes
-    ;;
   pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11)
     CUDA_VERSION=12.8.1
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=11
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    TRITON=yes
-    ;;
-  pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11)
-    CUDA_VERSION=13.0.0
+    CUDNN_VERSION=9
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
     VISION=yes
@@ -126,6 +102,7 @@ case "$tag" in
     ;;
   pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9-inductor-benchmarks)
     CUDA_VERSION=12.8.1
+    CUDNN_VERSION=9
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
     VISION=yes
@@ -135,18 +112,80 @@ case "$tag" in
     TRITON=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-jammy-cuda12.8-cudnn9-py3.12-gcc11-vllm)
+  pytorch-linux-jammy-cuda12.8-cudnn9-py3.12-gcc9-inductor-benchmarks)
     CUDA_VERSION=12.8.1
+    CUDNN_VERSION=9
     ANACONDA_PYTHON_VERSION=3.12
-    GCC_VERSION=11
+    GCC_VERSION=9
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-jammy-cuda12.8-cudnn9-py3.13-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.8.1
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.13
+    GCC_VERSION=9
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-jammy-cuda12.6-cudnn9-py3-gcc9)
+    CUDA_VERSION=12.6.3
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=9
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
     ;;
+  pytorch-linux-jammy-cuda12.6-cudnn9-py3-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=9
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-jammy-cuda12.6-cudnn9-py3.12-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.12
+    GCC_VERSION=9
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-jammy-cuda12.6-cudnn9-py3.13-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.13
+    GCC_VERSION=9
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
   pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc9)
     CUDA_VERSION=12.8.1
+    CUDNN_VERSION=9
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
     VISION=yes
@@ -156,54 +195,71 @@ case "$tag" in
     TRITON=yes
     ;;
   pytorch-linux-jammy-py3-clang12-onnx)
-    ANACONDA_PYTHON_VERSION=3.10
+    ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=12
     VISION=yes
     ONNX=yes
     ;;
-  pytorch-linux-jammy-py3.10-clang12)
-    ANACONDA_PYTHON_VERSION=3.10
+  pytorch-linux-jammy-py3.9-clang12)
+    ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=12
     VISION=yes
     TRITON=yes
     ;;
-  pytorch-linux-jammy-rocm-n-py3 | pytorch-linux-jammy-rocm-n-py3-benchmarks | pytorch-linux-noble-rocm-n-py3)
-    if [[ $tag =~ "jammy" ]]; then
-      ANACONDA_PYTHON_VERSION=3.10
-    else
-      ANACONDA_PYTHON_VERSION=3.12
-    fi
+  pytorch-linux-jammy-py3.11-clang12)
+    ANACONDA_PYTHON_VERSION=3.11
+    CLANG_VERSION=12
+    VISION=yes
+    TRITON=yes
+    ;;
+  pytorch-linux-jammy-py3.9-gcc9)
+    ANACONDA_PYTHON_VERSION=3.9
+    GCC_VERSION=9
+    VISION=yes
+    TRITON=yes
+    ;;
+  pytorch-linux-jammy-rocm-n-1-py3)
+    ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
     VISION=yes
-    ROCM_VERSION=7.0
+    ROCM_VERSION=6.3
     NINJA_VERSION=1.9.0
     TRITON=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
-    PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950"
-    if [[ $tag =~ "benchmarks" ]]; then
-      INDUCTOR_BENCHMARKS=yes
-    fi
+    INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-jammy-xpu-n-1-py3)
+  pytorch-linux-jammy-rocm-n-py3)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
     VISION=yes
+    ROCM_VERSION=6.4
+    NINJA_VERSION=1.9.0
+    TRITON=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-jammy-xpu-2025.0-py3)
+    ANACONDA_PYTHON_VERSION=3.9
+    GCC_VERSION=11
+    VISION=yes
+    XPU_VERSION=2025.0
+    NINJA_VERSION=1.9.0
+    TRITON=yes
+    ;;
+  pytorch-linux-jammy-xpu-2025.1-py3)
+    ANACONDA_PYTHON_VERSION=3.9
+    GCC_VERSION=11
+    VISION=yes
     XPU_VERSION=2025.1
     NINJA_VERSION=1.9.0
     TRITON=yes
     ;;
-  pytorch-linux-jammy-xpu-n-py3)
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=11
-    VISION=yes
-    XPU_VERSION=2025.2
-    NINJA_VERSION=1.9.0
-    TRITON=yes
-    ;;
-  pytorch-linux-jammy-py3-gcc11-inductor-benchmarks)
-    ANACONDA_PYTHON_VERSION=3.10
+    pytorch-linux-jammy-py3.9-gcc11-inductor-benchmarks)
+    ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     VISION=yes
     KATEX=yes
@@ -211,20 +267,32 @@ case "$tag" in
     DOCS=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang12)
-    ANACONDA_PYTHON_VERSION=3.10
+  pytorch-linux-jammy-cuda12.8-cudnn9-py3.9-clang12)
+    ANACONDA_PYTHON_VERSION=3.9
     CUDA_VERSION=12.8.1
+    CUDNN_VERSION=9
     CLANG_VERSION=12
     VISION=yes
     TRITON=yes
     ;;
+  pytorch-linux-jammy-py3-clang12-asan)
+    ANACONDA_PYTHON_VERSION=3.9
+    CLANG_VERSION=12
+    VISION=yes
+    TRITON=yes
+    ;;
+  pytorch-linux-jammy-py3-clang15-asan)
+    ANACONDA_PYTHON_VERSION=3.10
+    CLANG_VERSION=15
+    VISION=yes
+    ;;
   pytorch-linux-jammy-py3-clang18-asan)
     ANACONDA_PYTHON_VERSION=3.10
     CLANG_VERSION=18
     VISION=yes
     ;;
-  pytorch-linux-jammy-py3.10-gcc11)
-    ANACONDA_PYTHON_VERSION=3.10
+  pytorch-linux-jammy-py3.9-gcc11)
+    ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     VISION=yes
     KATEX=yes
@@ -251,10 +319,13 @@ case "$tag" in
     TRITON_CPU=yes
     ;;
   pytorch-linux-jammy-linter)
-    PYTHON_VERSION=3.10
+    # TODO: Use 3.9 here because of this issue https://github.com/python/mypy/issues/13627.
+    # We will need to update mypy version eventually, but that's for another day. The task
+    # would be to upgrade mypy to 1.0.0 with Python 3.11
+    PYTHON_VERSION=3.9
     ;;
-  pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-linter)
-    PYTHON_VERSION=3.10
+  pytorch-linux-jammy-cuda12.8-cudnn9-py3.9-linter)
+    PYTHON_VERSION=3.9
     CUDA_VERSION=12.8.1
     ;;
   pytorch-linux-jammy-aarch64-py3.10-gcc11)
@@ -262,6 +333,7 @@ case "$tag" in
     GCC_VERSION=11
     ACL=yes
     VISION=yes
+    CONDA_CMAKE=yes
     OPENBLAS=yes
     # snadampal: skipping llvm src build install because the current version
     # from pytorch/llvm:9.0.1 is x86 specific
@@ -272,15 +344,13 @@ case "$tag" in
     GCC_VERSION=11
     ACL=yes
     VISION=yes
+    CONDA_CMAKE=yes
     OPENBLAS=yes
     # snadampal: skipping llvm src build install because the current version
     # from pytorch/llvm:9.0.1 is x86 specific
     SKIP_LLVM_SRC_BUILD_INSTALL=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-noble-riscv64-py3.12-gcc14)
-    GCC_VERSION=14
-    ;;
   *)
     # Catch-all for builds that are not hardcoded.
     VISION=yes
@@ -290,6 +360,7 @@ case "$tag" in
     fi
     if [[ "$image" == *cuda* ]]; then
       extract_version_from_image_name cuda CUDA_VERSION
+      extract_version_from_image_name cudnn CUDNN_VERSION
     fi
     if [[ "$image" == *rocm* ]]; then
       extract_version_from_image_name rocm ROCM_VERSION
@@ -341,6 +412,9 @@ docker build \
        --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
        --build-arg "GCC_VERSION=${GCC_VERSION}" \
        --build-arg "CUDA_VERSION=${CUDA_VERSION}" \
+       --build-arg "CUDNN_VERSION=${CUDNN_VERSION}" \
+       --build-arg "TENSORRT_VERSION=${TENSORRT_VERSION}" \
+       --build-arg "GRADLE_VERSION=${GRADLE_VERSION}" \
        --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \
        --build-arg "KATEX=${KATEX:-}" \
        --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \
@@ -401,14 +475,7 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
 fi
 
 if [ -n "$GCC_VERSION" ]; then
-  if [[ "$image" == *riscv* ]]; then
-    # Check RISC-V cross-compilation toolchain version
-    if !(drun riscv64-linux-gnu-gcc-${GCC_VERSION} --version 2>&1 | grep -q " $GCC_VERSION\\W"); then
-      echo "RISC-V GCC_VERSION=$GCC_VERSION, but:"
-      drun riscv64-linux-gnu-gcc-${GCC_VERSION} --version
-      exit 1
-    fi
-  elif !(drun gcc --version 2>&1 | grep -q " $GCC_VERSION\\W"); then
+  if !(drun gcc --version 2>&1 | grep -q " $GCC_VERSION\\W"); then
     echo "GCC_VERSION=$GCC_VERSION, but:"
     drun gcc --version
     exit 1
@@ -441,3 +508,12 @@ elif [ "$HAS_TRITON" = "yes" ]; then
   echo "expecting triton to not be installed, but it is"
   exit 1
 fi
+
+# Sanity check cmake version.  Executorch reinstalls cmake and I'm not sure if
+# they support 4.0.0 yet, so exclude them from this check.
+CMAKE_VERSION=$(drun cmake --version)
+if [[ "$EXECUTORCH" != *yes* && "$CMAKE_VERSION" != *4.* ]]; then
+  echo "CMake version is not 4.0.0:"
+  drun cmake --version
+  exit 1
+fi

.ci/docker/centos-rocm/Dockerfile

@@ -56,13 +56,9 @@ ENV INSTALLED_VISION ${VISION}
 
 # Install rocm
 ARG ROCM_VERSION
-RUN mkdir ci_commit_pins
-COPY ./common/common_utils.sh common_utils.sh
-COPY ./ci_commit_pins/rocm-composable-kernel.txt ci_commit_pins/rocm-composable-kernel.txt
 COPY ./common/install_rocm.sh install_rocm.sh
 RUN bash ./install_rocm.sh
-RUN rm install_rocm.sh common_utils.sh
-RUN rm -r ci_commit_pins
+RUN rm install_rocm.sh
 COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
 RUN bash ./install_rocm_magma.sh ${ROCM_VERSION}
 RUN rm install_rocm_magma.sh

.ci/docker/ci_commit_pins/executorch.txt

@@ -1 +1 @@
-e0dda9059d082537cee36be6c5e4fe3b18c880c0
+56392aa978594cc155fa8af48cd949f5b5f1823a

.ci/docker/ci_commit_pins/huggingface-requirements.txt

@@ -1,2 +0,0 @@
-transformers==4.56.0
-soxr==0.5.0

.ci/docker/ci_commit_pins/huggingface.txt

@@ -0,0 +1 @@
+243e186efbf7fb93328dd6b34927a4e8c8f24395

.ci/docker/ci_commit_pins/nccl-cu12.txt

@@ -1 +1 @@
-v2.28.3-1
+v2.27.3-1

.ci/docker/ci_commit_pins/nccl-cu13.txt

@@ -1 +0,0 @@
-v2.28.3-1

.ci/docker/ci_commit_pins/rocm-composable-kernel.txt

@@ -1 +0,0 @@
-7fe50dc3da2069d6645d9deb8c017a876472a977

.ci/docker/ci_commit_pins/torchbench.txt

@@ -1 +0,0 @@
-74a23feff57432129df84d8099e622773cf77925

.ci/docker/ci_commit_pins/triton-xpu.txt

@@ -1 +1 @@
-1b0418a9a454b2b93ab8d71f40e59d2297157fae
+ae324eeac8e102a2b40370e341460f3791353398

.ci/docker/ci_commit_pins/triton.txt

@@ -1 +1 @@
-bbb06c0334a6772b92d24bde54956e675c8c6604
+11ec6354315768a85da41032535e3b7b99c5f706

.ci/docker/common/common_utils.sh

@@ -23,10 +23,6 @@ conda_install() {
   as_jenkins conda install -q -n py_$ANACONDA_PYTHON_VERSION -y python="$ANACONDA_PYTHON_VERSION" $*
 }
 
-conda_install_through_forge() {
-  as_jenkins conda install -c conda-forge -q -n py_$ANACONDA_PYTHON_VERSION -y python="$ANACONDA_PYTHON_VERSION" $*
-}
-
 conda_run() {
   as_jenkins conda run -n py_$ANACONDA_PYTHON_VERSION --no-capture-output $*
 }

.ci/docker/common/install_base.sh

@@ -15,9 +15,6 @@ install_ubuntu() {
   elif [[ "$UBUNTU_VERSION" == "22.04"* ]]; then
     cmake3="cmake=3.22*"
     maybe_libiomp_dev=""
-  elif [[ "$UBUNTU_VERSION" == "24.04"* ]]; then
-    cmake3="cmake=3.28*"
-    maybe_libiomp_dev=""
   else
     cmake3="cmake=3.5*"
     maybe_libiomp_dev="libiomp-dev"

.ci/docker/common/install_conda.sh

@@ -65,10 +65,10 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
   fi
 
   # Install PyTorch conda deps, as per https://github.com/pytorch/pytorch README
-  if [[ $(uname -m) != "aarch64" ]]; then
-    pip_install mkl==2024.2.0
-    pip_install mkl-static==2024.2.0
-    pip_install mkl-include==2024.2.0
+  if [[ $(uname -m) == "aarch64" ]]; then
+    conda_install "openblas==0.3.29=*openmp*"
+  else
+    conda_install "mkl=2021.4.0 mkl-include=2021.4.0"
   fi
 
   # Install llvm-8 as it is required to compile llvmlite-0.30.0 from source
@@ -82,10 +82,6 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
     conda_run ${SCRIPT_FOLDER}/install_magma_conda.sh $(cut -f1-2 -d'.' <<< ${CUDA_VERSION})
   fi
 
-  if [[ "$UBUNTU_VERSION" == "24.04"* ]] ; then
-    conda_install_through_forge libstdcxx-ng=14
-  fi
-
   # Install some other packages, including those needed for Python test reporting
   pip_install -r /opt/conda/requirements-ci.txt
 

.ci/docker/common/install_cpython.sh

@@ -3,10 +3,11 @@
 set -uex -o pipefail
 
 PYTHON_DOWNLOAD_URL=https://www.python.org/ftp/python
+PYTHON_DOWNLOAD_GITHUB_BRANCH=https://github.com/python/cpython/archive/refs/heads  # @lint-ignore
 GET_PIP_URL=https://bootstrap.pypa.io/get-pip.py
 
 # Python versions to be installed in /opt/$VERSION_NO
-CPYTHON_VERSIONS=${CPYTHON_VERSIONS:-"3.9.0 3.10.1 3.11.0 3.12.0 3.13.0 3.13.0t 3.14.0 3.14.0t"}
+CPYTHON_VERSIONS=${CPYTHON_VERSIONS:-"3.9.0 3.10.1 3.11.0 3.12.0 3.13.0 3.13.0t"}
 
 function check_var {
     if [ -z "$1" ]; then
@@ -23,8 +24,9 @@ function do_cpython_build {
     tar -xzf Python-$py_ver.tgz
 
     local additional_flags=""
-    if [[ "$py_ver" == *"t" ]]; then
+    if [ "$py_ver" == "3.13.0t" ]; then
         additional_flags=" --disable-gil"
+        mv cpython-3.13/ cpython-3.13t/
     fi
 
     pushd $py_folder
@@ -66,29 +68,32 @@ function do_cpython_build {
         ln -s pip3 ${prefix}/bin/pip
     fi
     # install setuptools since python 3.12 is required to use distutils
-    # packaging is needed to create symlink since wheel no longer provides needed information
-    ${prefix}/bin/pip install packaging==25.0 wheel==0.45.1 setuptools==80.9.0
-    local abi_tag=$(${prefix}/bin/python -c "from packaging.tags import interpreter_name, interpreter_version; import sysconfig ; from sysconfig import get_config_var; print('{0}{1}-{0}{1}{2}'.format(interpreter_name(), interpreter_version(), 't' if sysconfig.get_config_var('Py_GIL_DISABLED') else ''))")
+    ${prefix}/bin/pip install wheel==0.34.2 setuptools==68.2.2
+    local abi_tag=$(${prefix}/bin/python -c "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag; print('{0}{1}-{2}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag()))")
     ln -sf ${prefix} /opt/python/${abi_tag}
 }
 
 function build_cpython {
     local py_ver=$1
     check_var $py_ver
-    local py_suffix=$py_ver
-    local py_folder=$py_ver
+    check_var $PYTHON_DOWNLOAD_URL
+    local py_ver_folder=$py_ver
 
-    # Special handling for nogil
-    if [[ "${py_ver}" == *"t" ]]; then
-        py_suffix=${py_ver::-1}
-        py_folder=$py_suffix
+    if [ "$py_ver" = "3.13.0t" ]; then
+        PY_VER_SHORT="3.13"
+        PYT_VER_SHORT="3.13t"
+        check_var $PYTHON_DOWNLOAD_GITHUB_BRANCH
+        wget $PYTHON_DOWNLOAD_GITHUB_BRANCH/$PY_VER_SHORT.tar.gz -O Python-$py_ver.tgz
+        do_cpython_build $py_ver cpython-$PYT_VER_SHORT
+    elif [ "$py_ver" = "3.13.0" ]; then
+        PY_VER_SHORT="3.13"
+        check_var $PYTHON_DOWNLOAD_GITHUB_BRANCH
+        wget $PYTHON_DOWNLOAD_GITHUB_BRANCH/$PY_VER_SHORT.tar.gz -O Python-$py_ver.tgz
+        do_cpython_build $py_ver cpython-$PY_VER_SHORT
+    else
+        wget -q $PYTHON_DOWNLOAD_URL/$py_ver_folder/Python-$py_ver.tgz
+        do_cpython_build $py_ver Python-$py_ver
     fi
-    # Update to rc2 due to https://github.com/python/cpython/commit/c72699086fe4
-    if [ "$py_suffix" == "3.14.0" ]; then
-        py_suffix="3.14.0rc2"
-    fi
-    wget -q $PYTHON_DOWNLOAD_URL/$py_folder/Python-$py_suffix.tgz -O Python-$py_ver.tgz
-    do_cpython_build $py_ver Python-$py_suffix
 
     rm -f Python-$py_ver.tgz
 }

.ci/docker/common/install_cuda.sh

@@ -10,8 +10,6 @@ else
   arch_path='sbsa'
 fi
 
-NVSHMEM_VERSION=3.3.24
-
 function install_cuda {
   version=$1
   runfile=$2
@@ -42,67 +40,13 @@ function install_cudnn {
   rm -rf tmp_cudnn
 }
 
-function install_nvshmem {
-  cuda_major_version=$1      # e.g. "12"
-  nvshmem_version=$2         # e.g. "3.3.9"
-
-  case "${arch_path}" in
-    sbsa)
-      dl_arch="aarch64"
-      ;;
-    x86_64)
-      dl_arch="x64"
-      ;;
-    *)
-      dl_arch="${arch}"
-      ;;
-  esac
-
-  tmpdir="tmp_nvshmem"
-  mkdir -p "${tmpdir}" && cd "${tmpdir}"
-
-  # nvSHMEM license: https://docs.nvidia.com/nvshmem/api/sla.html
-  # This pattern is a lie as it is not consistent across versions, for 3.3.9 it was cuda_ver-arch-nvshhem-ver
-  filename="libnvshmem-linux-${arch_path}-${nvshmem_version}_cuda${cuda_major_version}-archive"
-  suffix=".tar.xz"
-  url="https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/linux-${arch_path}/${filename}${suffix}"
-
-  # download, unpack, install
-  wget -q "${url}"
-  tar xf "${filename}${suffix}"
-  cp -a "${filename}/include/"* /usr/local/cuda/include/
-  cp -a "${filename}/lib/"*     /usr/local/cuda/lib64/
-
-  # cleanup
-  cd ..
-  rm -rf "${tmpdir}"
-
-  echo "nvSHMEM ${nvshmem_version} for CUDA ${cuda_major_version} (${arch_path}) installed."
-}
-
-function install_124 {
-  CUDNN_VERSION=9.1.0.70
-  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.2"
-  install_cuda 12.4.1 cuda_12.4.1_550.54.15_linux
-
-  install_cudnn 12 $CUDNN_VERSION
-
-  CUDA_VERSION=12.4 bash install_nccl.sh
-
-  CUDA_VERSION=12.4 bash install_cusparselt.sh
-
-  ldconfig
-}
-
 function install_126 {
   CUDNN_VERSION=9.10.2.21
-  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.7.1"
   install_cuda 12.6.3 cuda_12.6.3_560.35.05_linux
 
   install_cudnn 12 $CUDNN_VERSION
 
-  install_nvshmem 12 $NVSHMEM_VERSION
-
   CUDA_VERSION=12.6 bash install_nccl.sh
 
   CUDA_VERSION=12.6 bash install_cusparselt.sh
@@ -112,15 +56,13 @@ function install_126 {
 
 function install_129 {
   CUDNN_VERSION=9.10.2.21
-  echo "Installing CUDA 12.9.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  echo "Installing CUDA 12.9.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.7.1"
   # install CUDA 12.9.1 in the same container
   install_cuda 12.9.1 cuda_12.9.1_575.57.08_linux
 
   # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
   install_cudnn 12 $CUDNN_VERSION
 
-  install_nvshmem 12 $NVSHMEM_VERSION
-
   CUDA_VERSION=12.9 bash install_nccl.sh
 
   CUDA_VERSION=12.9 bash install_cusparselt.sh
@@ -128,17 +70,49 @@ function install_129 {
   ldconfig
 }
 
+function prune_126 {
+  echo "Pruning CUDA 12.6"
+  #####################################################################################
+  # CUDA 12.6 prune static libs
+  #####################################################################################
+  export NVPRUNE="/usr/local/cuda-12.6/bin/nvprune"
+  export CUDA_LIB_DIR="/usr/local/cuda-12.6/lib64"
+
+  export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+  export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+
+  if [[ -n "$OVERRIDE_GENCODE" ]]; then
+      export GENCODE=$OVERRIDE_GENCODE
+  fi
+  if [[ -n "$OVERRIDE_GENCODE_CUDNN" ]]; then
+      export GENCODE_CUDNN=$OVERRIDE_GENCODE_CUDNN
+  fi
+
+  # all CUDA libs except CuDNN and CuBLAS
+  ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis"  \
+      | xargs -I {} bash -c \
+                "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
+
+  # prune CuDNN and CuBLAS
+  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
+  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
+
+  #####################################################################################
+  # CUDA 12.6 prune visual tools
+  #####################################################################################
+  export CUDA_BASE="/usr/local/cuda-12.6/"
+  rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2024.3.2 $CUDA_BASE/nsight-systems-2024.5.1/
+}
+
 function install_128 {
   CUDNN_VERSION=9.8.0.87
-  echo "Installing CUDA 12.8.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
+  echo "Installing CUDA 12.8.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.7.1"
   # install CUDA 12.8.1 in the same container
   install_cuda 12.8.1 cuda_12.8.1_570.124.06_linux
 
   # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
   install_cudnn 12 $CUDNN_VERSION
 
-  install_nvshmem 12 $NVSHMEM_VERSION
-
   CUDA_VERSION=12.8 bash install_nccl.sh
 
   CUDA_VERSION=12.8 bash install_cusparselt.sh
@@ -146,38 +120,16 @@ function install_128 {
   ldconfig
 }
 
-function install_130 {
-  CUDNN_VERSION=9.13.0.50
-  echo "Installing CUDA 13.0 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
-  # install CUDA 13.0 in the same container
-  install_cuda 13.0.0 cuda_13.0.0_580.65.06_linux
-
-  # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
-  install_cudnn 13 $CUDNN_VERSION
-
-  install_nvshmem 13 $NVSHMEM_VERSION
-
-  CUDA_VERSION=13.0 bash install_nccl.sh
-
-  CUDA_VERSION=13.0 bash install_cusparselt.sh
-
-  ldconfig
-}
-
 # idiomatic parameter and option handling in sh
 while test $# -gt 0
 do
     case "$1" in
-    12.4) install_124;
-        ;;
-    12.6|12.6.*) install_126;
+    12.6|12.6.*) install_126; prune_126
         ;;
     12.8|12.8.*) install_128;
         ;;
     12.9|12.9.*) install_129;
         ;;
-    13.0|13.0.*) install_130;
-        ;;
     *) echo "bad argument $1"; exit 1
         ;;
     esac

.ci/docker/common/install_cudnn.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+if [[ -n "${CUDNN_VERSION}" ]]; then
+    # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+    mkdir tmp_cudnn
+    pushd tmp_cudnn
+    if [[ ${CUDA_VERSION:0:4} == "12.9" || ${CUDA_VERSION:0:4} == "12.8" ]]; then
+        CUDNN_NAME="cudnn-linux-x86_64-9.10.2.21_cuda12-archive"
+    elif [[ ${CUDA_VERSION:0:4} == "12.6" ]]; then
+        CUDNN_NAME="cudnn-linux-x86_64-9.10.2.21_cuda12-archive"
+    elif [[ ${CUDA_VERSION:0:2} == "11" ]]; then
+        CUDNN_NAME="cudnn-linux-x86_64-9.1.0.70_cuda11-archive"
+    else
+        print "Unsupported CUDA version ${CUDA_VERSION}"
+        exit 1
+    fi
+    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/${CUDNN_NAME}.tar.xz
+    tar xf ${CUDNN_NAME}.tar.xz
+    cp -a ${CUDNN_NAME}/include/* /usr/local/cuda/include/
+    cp -a ${CUDNN_NAME}/lib/* /usr/local/cuda/lib64/
+    popd
+    rm -rf tmp_cudnn
+    ldconfig
+fi

.ci/docker/common/install_cusparselt.sh

@@ -5,15 +5,7 @@ set -ex
 # cuSPARSELt license: https://docs.nvidia.com/cuda/cusparselt/license.html
 mkdir tmp_cusparselt && cd tmp_cusparselt
 
-if [[ ${CUDA_VERSION:0:4} =~ "13" ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.8.0.4_cuda13-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
-elif [[ ${CUDA_VERSION:0:4} =~ ^12\.[5-9]$ ]]; then
+if [[ ${CUDA_VERSION:0:4} =~ ^12\.[5-9]$ ]]; then
     arch_path='sbsa'
     export TARGETARCH=${TARGETARCH:-$(uname -m)}
     if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
@@ -21,14 +13,6 @@ elif [[ ${CUDA_VERSION:0:4} =~ ^12\.[5-9]$ ]]; then
     fi
     CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.7.1.0-archive"
     curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
-elif [[ ${CUDA_VERSION:0:4} == "12.4" ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.6.2.3-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
 else
     echo "Not sure which libcusparselt version to install for this ${CUDA_VERSION}"
 fi

.ci/docker/common/install_executorch.sh

@@ -42,27 +42,22 @@ install_pip_dependencies() {
   # A workaround, ExecuTorch has moved to numpy 2.0 which is not compatible with the current
   # numba and scipy version used in PyTorch CI
   conda_run pip uninstall -y numba scipy
-  # Yaspin is needed for running CI test (get_benchmark_analysis_data.py)
-  pip_install yaspin==3.1.0
 
   popd
 }
 
 setup_executorch() {
+  pushd executorch
+
   export PYTHON_EXECUTABLE=python
-  export CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON -DEXECUTORCH_BUILD_TESTS=ON"
+  export CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON"
 
   as_jenkins .ci/scripts/setup-linux.sh --build-tool cmake || true
+  popd
 }
 
-if [ $# -eq 0 ]; then
-  clone_executorch
-  install_buck2
-  install_conda_dependencies
-  install_pip_dependencies
-  pushd executorch
-  setup_executorch
-  popd
-else
-  "$@"
-fi
+clone_executorch
+install_buck2
+install_conda_dependencies
+install_pip_dependencies
+setup_executorch

.ci/docker/common/install_inductor_benchmark_deps.sh

@@ -5,7 +5,9 @@ set -ex
 source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
 
 function install_huggingface() {
-  pip_install -r huggingface-requirements.txt
+  local version
+  commit=$(get_pinned_commit huggingface)
+  pip_install "git+https://github.com/huggingface/transformers@${commit}"
 }
 
 function install_timm() {
@@ -13,34 +15,11 @@ function install_timm() {
   commit=$(get_pinned_commit timm)
 
   pip_install "git+https://github.com/huggingface/pytorch-image-models@${commit}"
-}
-
-function install_torchbench() {
-  local commit
-  commit=$(get_pinned_commit torchbench)
-  git clone https://github.com/pytorch/benchmark torchbench
-  pushd torchbench
-  git checkout "$commit"
-
-  python install.py --continue_on_fail
-
-  echo "Print all dependencies after TorchBench is installed"
-  python -mpip freeze
-  popd
-
-  chown -R jenkins torchbench
-  chown -R jenkins /opt/conda
+  # Clean up
+  conda_run pip uninstall -y torch torchvision triton
 }
 
 # Pango is needed for weasyprint which is needed for doctr
 conda_install pango
-
-# Stable packages are ok here, just to satisfy TorchBench check
-pip_install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
-
-install_torchbench
 install_huggingface
 install_timm
-
-# Clean up
-conda_run pip uninstall -y torch torchvision torchaudio triton torchao

.ci/docker/common/install_nccl.sh

@@ -7,8 +7,6 @@ if [[ ${CUDA_VERSION:0:2} == "11" ]]; then
   NCCL_VERSION=$(cat ci_commit_pins/nccl-cu11.txt)
 elif [[ ${CUDA_VERSION:0:2} == "12" ]]; then
   NCCL_VERSION=$(cat ci_commit_pins/nccl-cu12.txt)
-elif [[ ${CUDA_VERSION:0:2} == "13" ]]; then
-  NCCL_VERSION=$(cat ci_commit_pins/nccl-cu13.txt)
 else
   echo "Unexpected CUDA_VERSION ${CUDA_VERSION}"
   exit 1

.ci/docker/common/install_onnx.sh

@@ -19,8 +19,8 @@ pip_install \
   transformers==4.36.2
 
 pip_install coloredlogs packaging
-pip_install onnxruntime==1.22.1
-pip_install onnxscript==0.4.0
+pip_install onnxruntime==1.18.1
+pip_install onnxscript==0.3.1
 
 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers
 # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/

.ci/docker/common/install_rocm.sh

@@ -2,22 +2,15 @@
 
 set -ex
 
-# for pip_install function
-source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
-
-ROCM_COMPOSABLE_KERNEL_VERSION="$(cat $(dirname $0)/../ci_commit_pins/rocm-composable-kernel.txt)"
-
 ver() {
     printf "%3d%03d%03d%03d" $(echo "$1" | tr '.' ' ');
 }
 
 install_ubuntu() {
     apt-get update
-    # gpg-agent is not available by default
-    apt-get install -y --no-install-recommends gpg-agent
-    if [[ $(ver $UBUNTU_VERSION) -ge $(ver 22.04) ]]; then
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' \
-            | sudo tee /etc/apt/preferences.d/rocm-pin-600
+    if [[ $UBUNTU_VERSION == 20.04 ]]; then
+      # gpg-agent is not available by default on 20.04
+      apt-get install -y --no-install-recommends gpg-agent
     fi
     apt-get install -y kmod
     apt-get install -y wget
@@ -35,19 +28,16 @@ EOF
 
     # we want the patch version of 6.4 instead
     if [[ $(ver $ROCM_VERSION) -eq $(ver 6.4) ]]; then
-        ROCM_VERSION="${ROCM_VERSION}.2"
+        ROCM_VERSION="${ROCM_VERSION}.1"
     fi
 
-    # Default url values
-    rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}"
-    amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu"
-
     # Add amdgpu repository
     UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'`
-    echo "deb [arch=amd64] ${amdgpu_baseurl} ${UBUNTU_VERSION_NAME} main" > /etc/apt/sources.list.d/amdgpu.list
+    echo "deb [arch=amd64] https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu ${UBUNTU_VERSION_NAME} main" > /etc/apt/sources.list.d/amdgpu.list
 
     # Add rocm repository
     wget -qO - http://repo.radeon.com/rocm/rocm.gpg.key | apt-key add -
+    local rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}"
     echo "deb [arch=amd64] ${rocm_baseurl} ${UBUNTU_VERSION_NAME} main" > /etc/apt/sources.list.d/rocm.list
     apt-get update --allow-insecure-repositories
 
@@ -81,39 +71,33 @@ EOF
     done
 
     # ROCm 6.3 had a regression where initializing static code objects had significant overhead
-    # CI no longer builds for ROCm 6.3, but
     # ROCm 6.4 did not yet fix the regression, also HIP branch names are different
-    if [[ $(ver $ROCM_VERSION) -ge $(ver 6.4) ]] && [[ $(ver $ROCM_VERSION) -lt $(ver 7.0) ]]; then
-        if [[ $(ver $ROCM_VERSION) -eq $(ver 6.4.2) ]]; then
-            HIP_TAG=rocm-6.4.2
-            CLR_HASH=74d78ba3ac4bac235d02bcb48511c30b5cfdd457  # branch release/rocm-rel-6.4.2-statco-hotfix
-        elif [[ $(ver $ROCM_VERSION) -eq $(ver 6.4.1) ]]; then
-            HIP_TAG=rocm-6.4.1
-            CLR_HASH=efe6c35790b9206923bfeed1209902feff37f386  # branch release/rocm-rel-6.4.1-statco-hotfix
+    if [[ $(ver $ROCM_VERSION) -ge $(ver 6.3) ]] && [[ $(ver $ROCM_VERSION) -lt $(ver 7.0) ]]; then
+        if [[ $(ver $ROCM_VERSION) -eq $(ver 6.4.1) ]]; then
+            HIP_BRANCH=release/rocm-rel-6.4
+            VER_STR=6.4
+            VER_PATCH=.1
         elif [[ $(ver $ROCM_VERSION) -eq $(ver 6.4) ]]; then
-            HIP_TAG=rocm-6.4.0
-            CLR_HASH=600f5b0d2baed94d5121e2174a9de0851b040b0c  # branch release/rocm-rel-6.4-statco-hotfix
+            HIP_BRANCH=release/rocm-rel-6.4
+            VER_STR=6.4
+        elif [[ $(ver $ROCM_VERSION) -eq $(ver 6.3) ]]; then
+            HIP_BRANCH=rocm-6.3.x
+            VER_STR=6.3
         fi
         # clr build needs CppHeaderParser but can only find it using conda's python
-        python -m pip install CppHeaderParser
-        git clone https://github.com/ROCm/HIP -b $HIP_TAG
+        /opt/conda/bin/python -m pip install CppHeaderParser
+        git clone https://github.com/ROCm/HIP -b $HIP_BRANCH
         HIP_COMMON_DIR=$(readlink -f HIP)
-        git clone https://github.com/jeffdaily/clr
-        pushd clr
-        git checkout $CLR_HASH
-        popd
+        git clone https://github.com/jeffdaily/clr -b release/rocm-rel-${VER_STR}${VER_PATCH}-statco-hotfix
         mkdir -p clr/build
         pushd clr/build
-        # Need to point CMake to the correct python installation to find CppHeaderParser
-        cmake .. -DPython3_EXECUTABLE=/opt/conda/envs/py_${ANACONDA_PYTHON_VERSION}/bin/python3 -DCLR_BUILD_HIP=ON -DHIP_COMMON_DIR=$HIP_COMMON_DIR
+        cmake .. -DCLR_BUILD_HIP=ON -DHIP_COMMON_DIR=$HIP_COMMON_DIR
         make -j
-        cp hipamd/lib/libamdhip64.so.6.4.* /opt/rocm/lib/libamdhip64.so.6.4.*
+        cp hipamd/lib/libamdhip64.so.${VER_STR}.* /opt/rocm/lib/libamdhip64.so.${VER_STR}.*
         popd
         rm -rf HIP clr
     fi
 
-    pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION"
-
     # Cleanup
     apt-get autoclean && apt-get clean
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
@@ -177,8 +161,6 @@ install_centos() {
       sqlite3 $kdb "PRAGMA journal_mode=off; PRAGMA VACUUM;"
   done
 
-  pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION"
-
   # Cleanup
   yum clean all
   rm -rf /var/cache/yum

.ci/docker/common/install_rocm_magma.sh

@@ -5,15 +5,10 @@ set -eou pipefail
 
 function do_install() {
     rocm_version=$1
-    if [[ ${rocm_version} =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
-        # chop off any patch version
-        rocm_version="${rocm_version%.*}"
-    fi
+    rocm_version_nodot=${1//./}
 
-    rocm_version_nodot=${rocm_version//./}
-
-    # https://github.com/icl-utk-edu/magma/pull/65
-    MAGMA_VERSION=d6e4117bc88e73f06d26c6c2e14f064e8fc3d1ec
+    # Version 2.7.2 + ROCm related updates
+    MAGMA_VERSION=a1625ff4d9bc362906bd01f805dbbe12612953f6
     magma_archive="magma-rocm${rocm_version_nodot}-${MAGMA_VERSION}-1.tar.bz2"
 
     rocm_dir="/opt/rocm"

.ci/docker/common/install_triton.sh

@@ -57,7 +57,7 @@ if [ ! -f setup.py ]; then
   cd python
 fi
 
-pip_install pybind11==3.0.1
+pip_install pybind11==2.13.6
 
 # TODO: remove patch setup.py once we have a proper fix for https://github.com/triton-lang/triton/issues/4527
 as_jenkins sed -i -e 's/https:\/\/tritonlang.blob.core.windows.net\/llvm-builds/https:\/\/oaitriton.blob.core.windows.net\/public\/llvm-builds/g' setup.py
@@ -66,15 +66,15 @@ if [ -n "${UBUNTU_VERSION}" ] && [ -n "${GCC_VERSION}" ] && [[ "${GCC_VERSION}"
   # Triton needs at least gcc-9 to build
   apt-get install -y g++-9
 
-  CXX=g++-9 conda_run python -m build --wheel --no-isolation
+  CXX=g++-9 conda_run python setup.py bdist_wheel
 elif [ -n "${UBUNTU_VERSION}" ] && [ -n "${CLANG_VERSION}" ]; then
   # Triton needs <filesystem> which surprisingly is not available with clang-9 toolchain
   add-apt-repository -y ppa:ubuntu-toolchain-r/test
   apt-get install -y g++-9
 
-  CXX=g++-9 conda_run python -m build --wheel --no-isolation
+  CXX=g++-9 conda_run python setup.py bdist_wheel
 else
-  conda_run python -m build --wheel --no-isolation
+  conda_run python setup.py bdist_wheel
 fi
 
 # Copy the wheel to /opt for multi stage docker builds
@@ -98,10 +98,6 @@ fi
 if [ -n "${NUMPY_VERSION}" ]; then
   pip_install "numpy==${NUMPY_VERSION}"
 fi
-
-# IMPORTANT: helion needs to be installed without dependencies.
-# It depends on torch and triton. We don't want to install
-# triton and torch from production on Docker CI images
 if [[ "$ANACONDA_PYTHON_VERSION" != 3.9* ]]; then
-  pip_install helion --no-deps
+  pip_install helion
 fi

.ci/docker/common/install_ucc.sh

@@ -44,12 +44,8 @@ function install_ucc() {
 
   ./autogen.sh
 
-  if [[ -n "$CUDA_VERSION"  && $CUDA_VERSION == 13* ]]; then
-    NVCC_GENCODE="-gencode=arch=compute_86,code=compute_86"
-  else
-    # We only run distributed tests on Tesla M60 and A10G
-    NVCC_GENCODE="-gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=compute_86"
-  fi
+  # We only run distributed tests on Tesla M60 and A10G
+  NVCC_GENCODE="-gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=compute_86"
 
   if [[ -n "$ROCM_VERSION" ]]; then
     if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then

.ci/docker/common/install_xpu.sh

@@ -34,27 +34,18 @@ function install_ubuntu() {
 
     # The xpu-smi packages
     apt-get install -y flex bison xpu-smi
-
-    if [[ "${XPU_DRIVER_TYPE,,}" == "lts" ]]; then
-        # Compute and Media Runtimes
-        apt-get install -y \
-            intel-opencl-icd intel-level-zero-gpu level-zero \
-            intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \
-            libegl-mesa0 libegl1-mesa libegl1-mesa-dev libgbm1 libgl1-mesa-dev libgl1-mesa-dri \
-            libglapi-mesa libgles2-mesa-dev libglx-mesa0 libigdgmm12 libxatracker2 mesa-va-drivers \
-            mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo
-        # Development Packages
-        apt-get install -y libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev
-    else # rolling driver
-        apt-get install -y \
-            intel-opencl-icd libze-intel-gpu1 libze1 \
-            intel-media-va-driver-non-free libmfx-gen1 libvpl2 \
-            libegl-mesa0 libegl1-mesa libegl1-mesa-dev libgbm1 libgl1-mesa-dev libgl1-mesa-dri \
-            libglapi-mesa libglx-mesa0 libigdgmm12 libxatracker2 mesa-va-drivers \
-            mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo intel-ocloc
-        apt-get install -y libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev libze-dev
+    # Compute and Media Runtimes
+    apt-get install -y \
+        intel-opencl-icd intel-level-zero-gpu level-zero \
+        intel-media-va-driver-non-free libmfx1 libmfxgen1 libvpl2 \
+        libegl-mesa0 libegl1-mesa libegl1-mesa-dev libgbm1 libgl1-mesa-dev libgl1-mesa-dri \
+        libglapi-mesa libgles2-mesa-dev libglx-mesa0 libigdgmm12 libxatracker2 mesa-va-drivers \
+        mesa-vdpau-drivers mesa-vulkan-drivers va-driver-all vainfo hwinfo clinfo
+    if [[ "${XPU_DRIVER_TYPE,,}" == "rolling" ]]; then
+        apt-get install -y intel-ocloc
     fi
-
+    # Development Packages
+    apt-get install -y libigc-dev intel-igc-cm libigdfcl-dev libigfxcmrt-dev level-zero-dev
     # Install Intel Support Packages
     apt-get install -y ${XPU_PACKAGES}
 
@@ -143,18 +134,18 @@ function install_sles() {
 
 }
 
-# Default use GPU driver rolling releases
-XPU_DRIVER_VERSION=""
-if [[ "${XPU_DRIVER_TYPE,,}" == "lts" ]]; then
-    # Use GPU driver LTS releases
-    XPU_DRIVER_VERSION="/lts/2350"
+# Default use GPU driver LTS releases
+XPU_DRIVER_VERSION="/lts/2350"
+if [[ "${XPU_DRIVER_TYPE,,}" == "rolling" ]]; then
+    # Use GPU driver rolling releases
+    XPU_DRIVER_VERSION=""
 fi
 
-# Default use Intel® oneAPI Deep Learning Essentials 2025.1
-if [[ "$XPU_VERSION" == "2025.2" ]]; then
-    XPU_PACKAGES="intel-deep-learning-essentials-2025.2"
-else
+# Default use Intel® oneAPI Deep Learning Essentials 2025.0
+if [[ "$XPU_VERSION" == "2025.1" ]]; then
     XPU_PACKAGES="intel-deep-learning-essentials-2025.1"
+else
+    XPU_PACKAGES="intel-deep-learning-essentials-2025.0"
 fi
 
 # The installation depends on the base OS

.ci/docker/libtorch/Dockerfile

@@ -69,19 +69,6 @@ RUN bash ./install_cuda.sh 12.9
 RUN bash ./install_magma.sh 12.9
 RUN ln -sf /usr/local/cuda-12.9 /usr/local/cuda
 
-FROM cuda as cuda13.0
-RUN bash ./install_cuda.sh 13.0
-RUN bash ./install_magma.sh 13.0
-RUN ln -sf /usr/local/cuda-13.0 /usr/local/cuda
-
-# Install libibverbs for libtorch and copy to CUDA directory
-RUN apt-get update -y && \
-    apt-get install -y libibverbs-dev librdmacm-dev && \
-    cp /usr/lib/x86_64-linux-gnu/libmlx5.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/librdmacm.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/libibverbs.so* /usr/local/cuda/lib64/ && \
-    cp /usr/lib/x86_64-linux-gnu/libnl* /usr/local/cuda/lib64/
-
 FROM cpu as rocm
 ARG ROCM_VERSION
 ARG PYTORCH_ROCM_ARCH

.ci/docker/libtorch/build.sh

@@ -39,17 +39,9 @@ case ${DOCKER_TAG_PREFIX} in
         DOCKER_GPU_BUILD_ARG=""
         ;;
     rocm*)
-        # we want the patch version of 6.4 instead
-        if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
-        fi
         BASE_TARGET=rocm
         GPU_IMAGE=rocm/dev-ubuntu-22.04:${GPU_ARCH_VERSION}-complete
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
-        # add gfx950 conditionally starting in ROCm 7.0
-        if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then
-            PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH};gfx950"
-        fi
         DOCKER_GPU_BUILD_ARG="--build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg ROCM_VERSION=${GPU_ARCH_VERSION}"
         ;;
     *)

.ci/docker/linter/Dockerfile

@@ -27,7 +27,5 @@ COPY ./common/install_linter.sh install_linter.sh
 RUN bash ./install_linter.sh
 RUN rm install_linter.sh
 
-RUN chown -R jenkins:jenkins /var/lib/jenkins/ci_env
-
 USER jenkins
 CMD ["bash"]

.ci/docker/manywheel/Dockerfile_2_28

@@ -175,6 +175,6 @@ ENV XPU_DRIVER_TYPE ROLLING
 RUN python3 -m pip install --upgrade pip && \
     python3 -mpip install cmake==3.28.4
 ADD ./common/install_xpu.sh install_xpu.sh
-ENV XPU_VERSION 2025.2
+ENV XPU_VERSION 2025.1
 RUN bash ./install_xpu.sh && rm install_xpu.sh
 RUN pushd /opt/_internal && tar -xJf static-libs-for-embedding-only.tar.xz && popd

.ci/docker/manywheel/Dockerfile_cxx11-abi

@@ -0,0 +1,71 @@
+FROM centos:8 as base
+
+ENV LC_ALL en_US.UTF-8
+ENV LANG en_US.UTF-8
+ENV LANGUAGE en_US.UTF-8
+ENV PATH /opt/rh/gcc-toolset-11/root/bin/:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+
+# change to a valid repo
+RUN sed -i 's|#baseurl=http://mirror.centos.org|baseurl=http://vault.centos.org|g' /etc/yum.repos.d/CentOS-Linux-*.repo
+# enable to install ninja-build
+RUN sed -i 's|enabled=0|enabled=1|g' /etc/yum.repos.d/CentOS-Linux-PowerTools.repo
+
+RUN yum -y update
+RUN yum install -y wget curl perl util-linux xz bzip2 git patch which zlib-devel sudo
+RUN yum install -y autoconf automake make cmake gdb gcc-toolset-11-gcc-c++
+
+
+FROM base as openssl
+ADD ./common/install_openssl.sh install_openssl.sh
+RUN bash ./install_openssl.sh && rm install_openssl.sh
+
+# Install python
+FROM base as python
+RUN yum install -y openssl-devel zlib-devel bzip2-devel ncurses-devel sqlite-devel readline-devel tk-devel gdbm-devel libpcap-devel xz-devel libffi-devel
+ADD common/install_cpython.sh install_cpython.sh
+RUN bash ./install_cpython.sh && rm install_cpython.sh
+
+FROM base as conda
+ADD ./common/install_conda_docker.sh install_conda.sh
+RUN bash ./install_conda.sh && rm install_conda.sh
+RUN /opt/conda/bin/conda install -y cmake
+
+FROM base as intel
+# Install MKL
+COPY --from=python             /opt/python                           /opt/python
+COPY --from=python             /opt/_internal                        /opt/_internal
+COPY --from=conda              /opt/conda                            /opt/conda
+ENV PATH=/opt/conda/bin:$PATH
+ADD ./common/install_mkl.sh install_mkl.sh
+RUN bash ./install_mkl.sh && rm install_mkl.sh
+
+FROM base as patchelf
+ADD ./common/install_patchelf.sh install_patchelf.sh
+RUN bash ./install_patchelf.sh && rm install_patchelf.sh
+RUN cp $(which patchelf) /patchelf
+
+FROM base as jni
+ADD ./common/install_jni.sh install_jni.sh
+ADD ./java/jni.h jni.h
+RUN bash ./install_jni.sh && rm install_jni.sh
+
+FROM base as libpng
+ADD ./common/install_libpng.sh install_libpng.sh
+RUN bash ./install_libpng.sh && rm install_libpng.sh
+
+FROM base as final
+COPY --from=openssl            /opt/openssl                          /opt/openssl
+COPY --from=python             /opt/python                           /opt/python
+COPY --from=python             /opt/_internal                        /opt/_internal
+COPY --from=intel              /opt/intel                            /opt/intel
+COPY --from=conda              /opt/conda                            /opt/conda
+COPY --from=patchelf           /usr/local/bin/patchelf               /usr/local/bin/patchelf
+COPY --from=jni                /usr/local/include/jni.h              /usr/local/include/jni.h
+COPY --from=libpng             /usr/local/bin/png*                   /usr/local/bin/
+COPY --from=libpng             /usr/local/bin/libpng*                /usr/local/bin/
+COPY --from=libpng             /usr/local/include/png*               /usr/local/include/
+COPY --from=libpng             /usr/local/include/libpng*            /usr/local/include/
+COPY --from=libpng             /usr/local/lib/libpng*                /usr/local/lib/
+COPY --from=libpng             /usr/local/lib/pkgconfig              /usr/local/lib/pkgconfig
+
+RUN yum install -y ninja-build

.ci/docker/manywheel/Dockerfile_s390x

@@ -131,8 +131,6 @@ RUN pip3 install flatbuffers && \
   git clone https://github.com/microsoft/onnxruntime && \
   cd onnxruntime && git checkout v1.21.0 && \
   git submodule update --init --recursive && \
-  wget https://github.com/microsoft/onnxruntime/commit/f57db79743c4d1a3553aa05cf95bcd10966030e6.patch && \
-  patch -p1 < f57db79743c4d1a3553aa05cf95bcd10966030e6.patch && \
   ./build.sh --config Release --parallel 0 --enable_pybind \
   --build_wheel --enable_training --enable_training_apis \
   --enable_training_ops --skip_tests --allow_running_as_root \

.ci/docker/manywheel/build.sh

@@ -43,6 +43,12 @@ case ${image} in
         MANY_LINUX_VERSION="2_28_aarch64"
         OPENBLAS_VERSION="v0.3.30"
         ;;
+    manylinuxcxx11-abi-builder:cpu-cxx11-abi)
+        TARGET=final
+        GPU_IMAGE=""
+        DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=9"
+        MANY_LINUX_VERSION="cxx11-abi"
+        ;;
     manylinuxs390x-builder:cpu-s390x)
         TARGET=final
         GPU_IMAGE=s390x/almalinux:8
@@ -61,12 +67,6 @@ case ${image} in
         DOCKER_GPU_BUILD_ARG="--build-arg BASE_CUDA_VERSION=${GPU_ARCH_VERSION} --build-arg DEVTOOLSET_VERSION=13"
         MANY_LINUX_VERSION="2_28"
         ;;
-    manylinux2_28-builder:cuda13*)
-        TARGET=cuda_final
-        GPU_IMAGE=amd64/almalinux:8
-        DOCKER_GPU_BUILD_ARG="--build-arg BASE_CUDA_VERSION=${GPU_ARCH_VERSION} --build-arg DEVTOOLSET_VERSION=13"
-        MANY_LINUX_VERSION="2_28"
-        ;;
     manylinuxaarch64-builder:cuda*)
         TARGET=cuda_final
         GPU_IMAGE=amd64/almalinux:8
@@ -75,19 +75,11 @@ case ${image} in
         DOCKERFILE_SUFFIX="_cuda_aarch64"
         ;;
     manylinux2_28-builder:rocm*)
-        # we want the patch version of 6.4 instead
-        if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
-        fi
         TARGET=rocm_final
         MANY_LINUX_VERSION="2_28"
         DEVTOOLSET_VERSION="11"
         GPU_IMAGE=rocm/dev-almalinux-8:${GPU_ARCH_VERSION}-complete
         PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
-        # add gfx950 conditionally starting in ROCm 7.0
-        if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then
-            PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH};gfx950"
-        fi
         DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}"
         ;;
     manylinux2_28-builder:xpu)

.ci/docker/requirements-ci.txt

@@ -10,18 +10,12 @@ boto3==1.35.42
 #Pinned versions: 1.19.12, 1.16.34
 #test that import:
 
-build==1.3.0
-#Description: A simple, correct Python build frontend.
-#Pinned versions: 1.3.0
-#test that import:
-
 click
 #Description: Command Line Interface Creation Kit
 #Pinned versions:
 #test that import:
 
 coremltools==5.0b5 ; python_version < "3.12"
-coremltools==8.3 ; python_version == "3.12"
 #Description: Apple framework for ML integration
 #Pinned versions: 5.0b5
 #test that import:
@@ -55,7 +49,7 @@ flatbuffers==24.12.23
 hypothesis==5.35.1
 # Pin hypothesis to avoid flakiness: https://github.com/pytorch/pytorch/issues/31136
 #Description: advanced library for generating parametrized tests
-#Pinned versions: 5.35.1
+#Pinned versions: 3.44.6, 4.53.2
 #test that import: test_xnnpack_integration.py, test_pruning_op.py, test_nn.py
 
 junitparser==2.1.1
@@ -68,12 +62,10 @@ lark==0.12.0
 #Pinned versions: 0.12.0
 #test that import:
 
-librosa>=0.6.2 ; python_version < "3.11" and platform_machine != "s390x"
-librosa==0.10.2 ; python_version == "3.12" and platform_machine != "s390x"
+librosa>=0.6.2 ; python_version < "3.11"
 #Description: A python package for music and audio analysis
 #Pinned versions: >=0.6.2
 #test that import: test_spectral_ops.py
-#librosa depends on numba; disable it for s390x while numba is disabled too
 
 #mkl #this breaks linux-bionic-rocm4.5-py3.7
 #Description: Intel oneAPI Math Kernel Library
@@ -98,9 +90,8 @@ librosa==0.10.2 ; python_version == "3.12" and platform_machine != "s390x"
 #Pinned versions:
 #test that import:
 
-mypy==1.16.0 ; platform_system != "Windows"
+mypy==1.16.0
 # Pin MyPy version because new errors are likely to appear with each release
-# Skip on Windows as lots of type annotations are POSIX specific
 #Description: linter
 #Pinned versions: 1.16.0
 #test that import: test_typing.py, test_type_hints.py
@@ -111,19 +102,19 @@ networkx==2.8.8
 #Pinned versions: 2.8.8
 #test that import: functorch
 
-ninja==1.11.1.4
+ninja==1.11.1.3
 #Description: build system. Used in some tests. Used in build to generate build
 #time tracing information
-#Pinned versions: 1.11.1.4
+#Pinned versions: 1.11.1.3
 #test that import: run_test.py, test_cpp_extensions_aot.py,test_determination.py
 
-numba==0.55.2 ; python_version == "3.10" and platform_machine != "s390x"
-numba==0.60.0 ; python_version == "3.12" and platform_machine != "s390x"
+numba==0.49.0 ; python_version < "3.9"
+numba==0.55.2 ; python_version == "3.9"
+numba==0.55.2 ; python_version == "3.10"
 #Description: Just-In-Time Compiler for Numerical Functions
 #Pinned versions: 0.54.1, 0.49.0, <=0.49.1
 #test that import: test_numba_integration.py
 #For numba issue see https://github.com/pytorch/pytorch/issues/51511
-#Need release > 0.61.2 for s390x due to https://github.com/numba/numba/pull/10073
 
 #numpy
 #Description: Provides N-dimensional arrays and linear algebra
@@ -137,7 +128,7 @@ numba==0.60.0 ; python_version == "3.12" and platform_machine != "s390x"
 #test_nn.py, test_namedtensor.py, test_linalg.py, test_jit_cuda_fuser.py,
 #test_jit.py, test_indexing.py, test_datapipe.py, test_dataloader.py,
 #test_binary_ufuncs.py
-numpy==1.22.4; python_version == "3.10"
+numpy==1.22.4; python_version == "3.9" or python_version == "3.10"
 numpy==1.26.2; python_version == "3.11" or python_version == "3.12"
 numpy==2.1.2; python_version >= "3.13"
 
@@ -172,9 +163,9 @@ pillow==11.0.0
 #Pinned versions: 10.3.0
 #test that import:
 
-protobuf==5.29.5
+protobuf==5.29.4
 #Description:  Google's data interchange format
-#Pinned versions: 5.29.5
+#Pinned versions: 5.29.4
 #test that import: test_tensorboard.py, test/onnx/*
 
 psutil
@@ -227,9 +218,9 @@ pygments==2.15.0
 #Pinned versions: 2.12.0
 #test that import: the doctests
 
-#pyyaml
+#PyYAML
 #Description: data serialization format
-#Pinned versions: 6.0.2
+#Pinned versions:
 #test that import:
 
 #requests
@@ -239,7 +230,7 @@ pygments==2.15.0
 
 #rich
 #Description: rich text and beautiful formatting in the terminal
-#Pinned versions: 14.1.0
+#Pinned versions: 10.9.0
 #test that import:
 
 scikit-image==0.19.3 ; python_version < "3.10"
@@ -267,6 +258,11 @@ scipy==1.14.1 ; python_version >= "3.12"
 #Pinned versions:
 #test that import:
 
+tb-nightly==2.13.0a20230426
+#Description: TensorBoard
+#Pinned versions:
+#test that import:
+
 # needed by torchgen utils
 typing-extensions>=4.10.0
 #Description: type hints for python
@@ -308,7 +304,7 @@ pytest-cpp==2.3.0
 #Pinned versions: 2.3.0
 #test that import:
 
-z3-solver==4.15.1.0 ; platform_machine != "s390x"
+z3-solver==4.12.6.0
 #Description: The Z3 Theorem Prover Project
 #Pinned versions:
 #test that import:
@@ -329,6 +325,8 @@ pywavelets==1.7.0 ; python_version >= "3.12"
 lxml==5.3.0
 #Description: This is a requirement of unittest-xml-reporting
 
+# Python-3.9 binaries
+
 PyGithub==2.3.0
 
 sympy==1.13.3
@@ -341,7 +339,7 @@ onnx==1.18.0
 #Pinned versions:
 #test that import:
 
-onnxscript==0.4.0
+onnxscript==0.3.1
 #Description: Required by mypy and test_public_bindings.py when checking torch.onnx._internal
 #Pinned versions:
 #test that import:
@@ -360,11 +358,12 @@ pwlf==2.2.1
 #Pinned versions: 2.2.1
 #test that import: test_sac_estimator.py
 
+
 # To build PyTorch itself
-pyyaml
+astunparse
+PyYAML
 pyzstd
-setuptools>=70.1.0
-six
+setuptools
 
 scons==4.5.2 ; platform_machine == "aarch64"
 
@@ -378,18 +377,12 @@ dataclasses_json==0.6.7
 #Pinned versions: 0.6.7
 #test that import:
 
-cmake==3.31.6
+cmake==4.0.0
 #Description: required for building
 
-tlparse==0.4.0
+tlparse==0.3.30
 #Description: required for log parsing
 
-cuda-bindings>=12.0,<13.0 ; platform_machine != "s390x"
+cuda-bindings>=12.0,<13.0
 #Description: required for testing CUDAGraph::raw_cuda_graph(). See https://nvidia.github.io/cuda-python/cuda-bindings/latest/support.html for how this version was chosen. Note "Any fix in the latest bindings would be backported to the prior major version" means that only the newest version of cuda-bindings will get fixes. Depending on the latest version of 12.x is okay because all 12.y versions will be supported via "CUDA minor version compatibility". Pytorch builds against 13.z versions of cuda toolkit work with 12.x versions of cuda-bindings as well because newer drivers work with old toolkits.
 #test that import: test_cuda.py
-
-setuptools-git-versioning==2.1.0
-scikit-build==0.18.1
-pyre-extensions==0.0.32
-tabulate==0.9.0
-#Description: These package are needed to build FBGEMM and torchrec on PyTorch CI

.ci/docker/requirements-docs.txt

@@ -1,15 +1,8 @@
 sphinx==5.3.0
 #Description: This is used to generate PyTorch docs
 #Pinned versions: 5.3.0
+-e git+https://github.com/pytorch/pytorch_sphinx_theme.git@722b7e6f9ca512fcc526ad07d62b3d28c50bb6cd#egg=pytorch_sphinx_theme2
 
-standard-imghdr==3.13.0; python_version >= "3.13"
-#Description: This is needed by Sphinx, so it needs to be added here.
-# The reasons are as follows:
-# 1) This module has been removed from the Python standard library since Python 3.13(https://peps.python.org/pep-0594/#imghdr);
-# 2) The current version of Sphinx (5.3.0) is not compatible with Python 3.13.
-# Once Sphinx is upgraded to a version compatible with Python 3.13 or later, we can remove this dependency.
-
--e git+https://github.com/pytorch/pytorch_sphinx_theme.git@71e55749be14ceb56e7f8211a9fb649866b87ad4#egg=pytorch_sphinx_theme2
 # TODO: sphinxcontrib.katex 0.9.0 adds a local KaTeX server to speed up pre-rendering
 # but it doesn't seem to work and hangs around idly. The initial thought that it is probably
 # something related to Docker setup. We can investigate this later.
@@ -26,10 +19,9 @@ sphinx_sitemap==2.6.0
 #Description: This is used to generate sitemap for PyTorch docs
 #Pinned versions: 2.6.0
 
-matplotlib==3.5.3 ; python_version < "3.13"
-matplotlib==3.6.3 ; python_version >= "3.13"
+matplotlib==3.5.3
 #Description: This is used to generate PyTorch docs
-#Pinned versions: 3.6.3 if python > 3.12. Otherwise 3.5.3.
+#Pinned versions: 3.5.3
 
 tensorboard==2.13.0 ; python_version < "3.13"
 tensorboard==2.18.0 ; python_version >= "3.13"

.ci/docker/triton_version.txt

@@ -1 +1 @@
-3.5.0
+3.4.0

.ci/docker/triton_xpu_version.txt

@@ -1 +1 @@
-3.5.0
+3.4.0

.ci/docker/ubuntu-cross-riscv/Dockerfile

@@ -1,155 +0,0 @@
-# Cross-compilation Docker container for RISC-V architecture
-ARG UBUNTU_VERSION
-FROM --platform=linux/amd64 ubuntu:${UBUNTU_VERSION} as base
-
-ARG UBUNTU_VERSION
-
-ENV GCC_VERSION=14
-ENV PYTHON_VERSION=3.12.3
-ENV DEBIAN_FRONTEND=noninteractive
-ENV CC=riscv64-linux-gnu-gcc-${GCC_VERSION}
-ENV CXX=riscv64-linux-gnu-g++-${GCC_VERSION}
-ENV QEMU_LD_PREFIX=/usr/riscv64-linux-gnu/
-ENV SYSROOT=/opt/sysroot
-
-# Install basic dependencies
-RUN apt-get update && apt-get install -y \
-    ninja-build \
-    autoconf \
-    automake \
-    libtool \
-    patchelf \
-    ccache \
-    git \
-    wget \
-    python3-pip \
-    python3-venv \
-    python-is-python3 \
-    cmake \
-    sudo \
-    lsb-release \
-    gcc-${GCC_VERSION}-riscv64-linux-gnu \
-    g++-${GCC_VERSION}-riscv64-linux-gnu \
-    pkg-config \
-    && rm -rf /var/lib/apt/lists/*
-
-# Install user
-COPY ./common/install_user.sh install_user.sh
-RUN bash ./install_user.sh && rm install_user.sh
-
-FROM base as python
-ARG ZLIB_VERSION=1.3.1
-ARG FFI_VERSION=3.4.6
-ARG BZ2_VERSION=1.0.8
-ARG XZ_VERSION=5.4.6
-ARG OPENSSL_VERSION=3.2.1
-
-# Set up sysroot directory for dependencies
-ENV PKG_CONFIG_PATH=${SYSROOT}/lib/pkgconfig
-ENV PKG_CONFIG_SYSROOT_DIR=${SYSROOT}
-
-WORKDIR /opt
-
-# Build zlib (for compression)
-RUN echo "--- Building zlib ---" \
-    && wget -c https://www.zlib.net/zlib-${ZLIB_VERSION}.tar.gz \
-    && tar -xf zlib-${ZLIB_VERSION}.tar.gz --no-same-permissions --no-same-owner \
-    && cd zlib-${ZLIB_VERSION}/ \
-    && mkdir build && cd build \
-    && ../configure --prefix=${SYSROOT} \
-    && make -j$(nproc) && make install \
-    && cd ../..
-
-# Build libffi (for ctypes module)
-RUN echo "--- Building libffi ---" \
-    && wget -c https://github.com/libffi/libffi/releases/download/v${FFI_VERSION}/libffi-${FFI_VERSION}.tar.gz \
-    && tar -xf libffi-${FFI_VERSION}.tar.gz --no-same-permissions --no-same-owner \
-    && cd libffi-${FFI_VERSION}/ \
-    && mkdir build && cd build \
-    && ../configure --prefix=${SYSROOT} --host=riscv64-linux-gnu --build=x86_64-linux-gnu \
-    && make -j$(nproc) && make install \
-    && cd ../..
-
-# Build bzip2 (for bz2 module)
-RUN echo "--- Building bzip2 ---" \
-    && wget -c https://sourceware.org/pub/bzip2/bzip2-${BZ2_VERSION}.tar.gz \
-    && tar -xf bzip2-${BZ2_VERSION}.tar.gz --no-same-permissions --no-same-owner \
-    && cd bzip2-${BZ2_VERSION}/ \
-    && make CC=riscv64-linux-gnu-gcc-${GCC_VERSION} bzip2 bzip2recover libbz2.a \
-    && make CC=riscv64-linux-gnu-gcc-${GCC_VERSION} -f Makefile-libbz2_so \
-    && make install PREFIX=${SYSROOT} \
-    && cp libbz2.so.${BZ2_VERSION} ${SYSROOT}/lib/ \
-    && cd ${SYSROOT}/lib/ \
-    && ln -sf libbz2.so.${BZ2_VERSION} libbz2.so.1.0 \
-    && ln -sf libbz2.so.1.0 libbz2.so \
-    && cd /opt/
-
-# Build xz (for lzma module)
-RUN echo "--- Building xz ---" \
-    && wget -c https://github.com/tukaani-project/xz/releases/download/v${XZ_VERSION}/xz-${XZ_VERSION}.tar.gz \
-    && tar -xf xz-${XZ_VERSION}.tar.gz --no-same-permissions --no-same-owner \
-    && cd xz-${XZ_VERSION} \
-    && mkdir build && cd build \
-    && ../configure --prefix=${SYSROOT} --host=riscv64-linux-gnu --build=x86_64-linux-gnu \
-    && make -j$(nproc) && make install \
-    && cd ../..
-
-# Build OpenSSL (for ssl module)
-RUN echo "--- Building OpenSSL ---" \
-    && wget -c https://www.openssl.org/source/openssl-${OPENSSL_VERSION}.tar.gz \
-    && tar -xf openssl-${OPENSSL_VERSION}.tar.gz --no-same-permissions --no-same-owner \
-    && cd openssl-${OPENSSL_VERSION}/ \
-    && mkdir build && cd build \
-    && ../Configure linux64-riscv64 --prefix=${SYSROOT} \
-    && make -j$(nproc) && make install_sw \
-    && cd ../..
-
-# Build SQLite3 (for sqlite3 module)
-RUN echo "--- Building SQLite3 ---" \
-    && wget -c https://www.sqlite.org/2024/sqlite-autoconf-3450200.tar.gz \
-    && tar -xf sqlite-autoconf-3450200.tar.gz --no-same-permissions --no-same-owner \
-    && cd sqlite-autoconf-3450200 \
-    && mkdir build && cd build \
-    && ../configure --prefix=${SYSROOT} --host=riscv64-linux-gnu --build=x86_64-linux-gnu \
-    && make -j$(nproc) && make install \
-    && cd ../..
-
-# Build and install RISC-V Python with all modules
-RUN wget -c https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tgz \
-    && tar -xf Python-${PYTHON_VERSION}.tgz --no-same-permissions --no-same-owner \
-    && cd Python-${PYTHON_VERSION} \
-    && mkdir build && cd build \
-    && ../configure \
-        --host=riscv64-linux-gnu \
-        --build=x86_64-linux-gnu \
-        --prefix=${SYSROOT} \
-        --enable-shared \
-        --disable-ipv6 \
-        --with-build-python=/usr/bin/python3 \
-        --with-ensurepip=no \
-        ac_cv_file__dev_ptmx=yes \
-        ac_cv_file__dev_ptc=no \
-    && make -j$(nproc) \
-    && make install
-
-FROM base as final
-COPY --from=python             /opt/sysroot                       /opt/sysroot
-
-# Install crossenv and cmake
-RUN pip install crossenv cmake==4.0.0 --break-system-packages \
-    && /usr/bin/python3 -m crossenv ${SYSROOT}/bin/python3 /opt/riscv-cross-env
-
-# Add pip-installed cmake binaries to PATH
-ENV PATH="/usr/local/bin:${PATH}"
-
-# Set up cross Python environment
-SHELL ["/bin/bash", "-c"]
-RUN source /opt/riscv-cross-env/bin/activate \
-    && pip install setuptools pyyaml typing_extensions wheel
-
-# Set default environment variables for PyTorch build
-ENV Python_ROOT_DIR=${SYSROOT}
-ENV OPENSSL_ROOT_DIR=${SYSROOT}
-
-USER jenkins
-CMD ["bash"]

.ci/docker/ubuntu-rocm/Dockerfile

@@ -52,13 +52,9 @@ ENV INSTALLED_VISION ${VISION}
 
 # Install rocm
 ARG ROCM_VERSION
-RUN mkdir ci_commit_pins
-COPY ./common/common_utils.sh common_utils.sh
-COPY ./ci_commit_pins/rocm-composable-kernel.txt ci_commit_pins/rocm-composable-kernel.txt
 COPY ./common/install_rocm.sh install_rocm.sh
 RUN bash ./install_rocm.sh
-RUN rm install_rocm.sh common_utils.sh
-RUN rm -r ci_commit_pins
+RUN rm install_rocm.sh
 COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
 RUN bash ./install_rocm_magma.sh ${ROCM_VERSION}
 RUN rm install_rocm_magma.sh
@@ -100,11 +96,10 @@ ARG ANACONDA_PYTHON_VERSION
 ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
 COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
 COPY ./common/common_utils.sh common_utils.sh
-COPY ci_commit_pins/huggingface-requirements.txt huggingface-requirements.txt
+COPY ci_commit_pins/huggingface.txt huggingface.txt
 COPY ci_commit_pins/timm.txt timm.txt
-COPY ci_commit_pins/torchbench.txt torchbench.txt
 RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
-RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt
+RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
 
 # (optional) Install non-default Ninja version
 ARG NINJA_VERSION

.ci/docker/ubuntu-xpu/Dockerfile

@@ -56,10 +56,10 @@ RUN rm install_openssl.sh
 ARG INDUCTOR_BENCHMARKS
 COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
 COPY ./common/common_utils.sh common_utils.sh
-COPY ci_commit_pins/huggingface-requirements.txt huggingface-requirements.txt
+COPY ci_commit_pins/huggingface.txt huggingface.txt
 COPY ci_commit_pins/timm.txt timm.txt
 RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
-RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt
+RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
 
 # Install XPU Dependencies
 ARG XPU_VERSION

.ci/docker/ubuntu/Dockerfile

@@ -66,7 +66,6 @@ ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/"
 # (optional) Install UCC
 ARG UCX_COMMIT
 ARG UCC_COMMIT
-ARG CUDA_VERSION
 ENV UCX_COMMIT $UCX_COMMIT
 ENV UCC_COMMIT $UCC_COMMIT
 ENV UCX_HOME /usr
@@ -97,11 +96,10 @@ RUN rm install_openssl.sh
 ARG INDUCTOR_BENCHMARKS
 COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
 COPY ./common/common_utils.sh common_utils.sh
-COPY ci_commit_pins/huggingface-requirements.txt huggingface-requirements.txt
+COPY ci_commit_pins/huggingface.txt huggingface.txt
 COPY ci_commit_pins/timm.txt timm.txt
-COPY ci_commit_pins/torchbench.txt torchbench.txt
 RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
-RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt
+RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
 
 ARG TRITON
 ARG TRITON_CPU
@@ -182,6 +180,7 @@ COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
 RUN if [ -n "${SKIP_LLVM_SRC_BUILD_INSTALL}" ]; then set -eu; rm -rf /opt/llvm; fi
 
 # AWS specific CUDA build guidance
+ENV TORCH_CUDA_ARCH_LIST Maxwell
 ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all"
 ENV CUDA_PATH /usr/local/cuda
 

.ci/libtorch/build.sh

@@ -7,4 +7,4 @@ set -ex
 
 SCRIPTPATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
-USE_NVSHMEM=0 USE_CUSPARSELT=0 BUILD_PYTHONLESS=1 DESIRED_PYTHON="3.10" ${SCRIPTPATH}/../manywheel/build.sh
+USE_CUSPARSELT=0 BUILD_PYTHONLESS=1 DESIRED_PYTHON="3.9" ${SCRIPTPATH}/../manywheel/build.sh

.ci/lumen_cli/README.md

@@ -1,31 +0,0 @@
-# 🔧 Lumen_cli
-A Python CLI tool for building and testing PyTorch-based components, using a YAML configuration file for structured, repeatable workflows.
-
-
-## Features
-- **Build**
-    - external projects (e.g. vLLM)
-
-## 📦 Installation
-at the root of the pytorch repo
-```bash
-pip install -e .ci/lumen_cli
-```
-
-## Run the cli tool
-The cli tool must be used at root of pytorch repo, as example to run build external vllm:
-```bash
-python -m cli.run build external vllm
-```
-this will run the build steps with default behaviour for vllm project.
-
-to see help messages, run
-```bash
-python3 -m cli.run --help
-```
-
-## Add customized external build logics
-To add a new external build, for instance, add a new external build logics:
-1. create the build function in cli/lib folder
-2. register your target and the main build function at  EXTERNAL_BUILD_TARGET_DISPATCH in `cli/build_cli/register_build.py`
-3. [optional] create your ci config file in .github/ci_configs/${EXTERNAL_PACKAGE_NAME}.yaml

.ci/lumen_cli/cli/build_cli/register_build.py

@@ -1,37 +0,0 @@
-import argparse
-import logging
-
-from cli.lib.common.cli_helper import register_targets, RichHelp, TargetSpec
-from cli.lib.core.vllm.vllm_build import VllmBuildRunner
-
-
-logger = logging.getLogger(__name__)
-
-# Maps targets to their argparse configuration and runner
-# it adds new target to path python -m cli.run build external {target} with buildrunner
-_TARGETS: dict[str, TargetSpec] = {
-    "vllm": {
-        "runner": VllmBuildRunner,
-        "help": "Build vLLM using docker buildx.",
-    }
-    # add yours ...
-}
-
-
-def register_build_commands(subparsers: argparse._SubParsersAction) -> None:
-    build_parser = subparsers.add_parser(
-        "build",
-        help="Build related commands",
-        formatter_class=RichHelp,
-    )
-    build_subparsers = build_parser.add_subparsers(dest="build_command", required=True)
-    overview = "\n".join(
-        f"  {name:12} {spec.get('help', '')}" for name, spec in _TARGETS.items()
-    )
-    external_parser = build_subparsers.add_parser(
-        "external",
-        help="Build external targets",
-        description="Build third-party targets.\n\nAvailable targets:\n" + overview,
-        formatter_class=RichHelp,
-    )
-    register_targets(external_parser, _TARGETS)

.ci/lumen_cli/cli/lib/common/cli_helper.py

@@ -1,71 +0,0 @@
-"""
-Cli Argparser Utility helpers for CLI tasks.
-
-"""
-
-import argparse
-from abc import ABC, abstractmethod
-
-
-try:
-    from typing import Any, Callable, Required, TypedDict  # Python 3.11+
-except ImportError:
-    from typing import Any, Callable, TypedDict
-
-    from typing_extensions import Required  # Fallback for Python <3.11
-
-
-class BaseRunner(ABC):
-    def __init__(self, args: Any) -> None:
-        self.args = args
-
-    @abstractmethod
-    def run(self) -> None:
-        """runs main logics, required"""
-
-
-# Pretty help: keep newlines + show defaults
-class RichHelp(
-    argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter
-):
-    pass
-
-
-class TargetSpec(TypedDict, total=False):
-    """CLI subcommand specification with bA."""
-
-    runner: Required[type[BaseRunner]]
-    help: str
-    description: str
-    add_arguments: Callable[[argparse.ArgumentParser], None]
-
-
-def register_targets(
-    parser: argparse.ArgumentParser,
-    target_specs: dict[str, TargetSpec],
-    common_args: Callable[[argparse.ArgumentParser], None] = lambda _: None,
-) -> None:
-    """Register target subcommands."""
-    targets = parser.add_subparsers(
-        dest="target",
-        required=True,
-        metavar="{" + ",".join(target_specs.keys()) + "}",
-    )
-
-    for name, spec in target_specs.items():
-        desc = spec.get("description") or spec["runner"].__doc__ or ""
-
-        p = targets.add_parser(
-            name,
-            help=spec.get("help", ""),
-            description=desc.strip(),
-            formatter_class=RichHelp,
-        )
-        p.set_defaults(
-            func=lambda args, cls=spec["runner"]: cls(args).run(),
-            _runner_class=spec["runner"],
-        )
-        if "add_arguments" in spec and callable(spec["add_arguments"]):
-            spec["add_arguments"](p)
-        if common_args:
-            common_args(p)

.ci/lumen_cli/cli/lib/common/docker_helper.py

@@ -1,42 +0,0 @@
-"""
-Docker Utility helpers for CLI tasks.
-"""
-
-import logging
-from typing import Optional
-
-import docker
-from docker.errors import APIError, NotFound
-
-
-logger = logging.getLogger(__name__)
-
-# lazy singleton so we don't reconnect every call
-_docker_client: Optional[docker.DockerClient] = None
-
-
-def _get_client() -> docker.DockerClient:
-    global _docker_client
-    if _docker_client is None:
-        _docker_client = docker.from_env()
-    return _docker_client
-
-
-def local_image_exists(
-    image_name: str, client: Optional[docker.DockerClient] = None
-) -> bool:
-    """Return True if a local Docker image exists."""
-    if not image_name:
-        return False
-
-    client = client or _get_client()
-    try:
-        client.images.get(image_name)
-        return True
-    except (NotFound, APIError) as e:
-        logger.error(
-            "Error when checking Docker image '%s': %s",
-            image_name,
-            e.explanation if hasattr(e, "explanation") else str(e),
-        )
-        return False

.ci/lumen_cli/cli/lib/common/envs_helper.py

@@ -1,110 +0,0 @@
-"""
-Environment Variables and Dataclasses Utility helpers for CLI tasks.
-"""
-
-import os
-from dataclasses import field, fields, is_dataclass, MISSING
-from pathlib import Path
-from textwrap import indent
-from typing import Optional, Union
-
-from cli.lib.common.utils import str2bool
-
-
-def get_env(name: str, default: str = "") -> str:
-    """Get environment variable with default fallback."""
-    return os.environ.get(name) or default
-
-
-def env_path_optional(
-    name: str,
-    default: Optional[Union[str, Path]] = None,
-    resolve: bool = True,
-) -> Optional[Path]:
-    """Get environment variable as optional Path."""
-    val = get_env(name) or default
-    if not val:
-        return None
-
-    path = Path(val)
-    return path.resolve() if resolve else path
-
-
-def env_path(
-    name: str,
-    default: Optional[Union[str, Path]] = None,
-    resolve: bool = True,
-) -> Path:
-    """Get environment variable as Path, raise if missing."""
-    path = env_path_optional(name, default, resolve)
-    if not path:
-        raise ValueError(f"Missing path value for {name}")
-    return path
-
-
-def env_bool(
-    name: str,
-    default: bool = False,
-) -> bool:
-    val = get_env(name)
-    if not val:
-        return default
-    return str2bool(val)
-
-
-def env_bool_field(
-    name: str,
-    default: bool = False,
-):
-    return field(default_factory=lambda: env_bool(name, default))
-
-
-def env_path_field(
-    name: str,
-    default: Union[str, Path] = "",
-    *,
-    resolve: bool = True,
-) -> Path:
-    return field(default_factory=lambda: env_path(name, default, resolve=resolve))
-
-
-def env_str_field(
-    name: str,
-    default: str = "",
-) -> str:
-    return field(default_factory=lambda: get_env(name, default))
-
-
-def generate_dataclass_help(cls) -> str:
-    """Auto-generate help text for dataclass fields."""
-    if not is_dataclass(cls):
-        raise TypeError(f"{cls} is not a dataclass")
-
-    def get_value(f):
-        if f.default is not MISSING:
-            return f.default
-        if f.default_factory is not MISSING:
-            try:
-                return f.default_factory()
-            except Exception as e:
-                return f"<error: {e}>"
-        return "<required>"
-
-    lines = [f"{f.name:<22} = {repr(get_value(f))}" for f in fields(cls)]
-    return indent("\n".join(lines), "    ")
-
-
-def with_params_help(params_cls: type, title: str = "Parameter defaults"):
-    """
-    Class decorator that appends a help table generated from another dataclass
-    (e.g., VllmParameters) to the decorated class's docstring.
-    """
-    if not is_dataclass(params_cls):
-        raise TypeError(f"{params_cls} must be a dataclass")
-
-    def _decorator(cls: type) -> type:
-        block = generate_dataclass_help(params_cls)
-        cls.__doc__ = (cls.__doc__ or "") + f"\n\n{title}:\n{block}"
-        return cls
-
-    return _decorator

.ci/lumen_cli/cli/lib/common/gh_summary.py

@@ -1,143 +0,0 @@
-from __future__ import annotations
-
-import logging
-import os
-import textwrap
-from pathlib import Path
-from typing import TYPE_CHECKING
-
-from cli.lib.common.utils import get_wheels
-from jinja2 import Template
-
-
-if TYPE_CHECKING:
-    from collections.abc import Iterable, Mapping
-
-
-logger = logging.getLogger(__name__)
-
-_TPL_CONTENT = Template(
-    textwrap.dedent("""\
-    ## {{ title }}
-
-    ```{{ lang }}
-    {{ content }}
-    ```
-""")
-)
-
-_TPL_LIST_ITEMS = Template(
-    textwrap.dedent("""\
-    ## {{ title }}
-    {% for it in items %}
-    - {{ it.pkg }}: {{ it.relpath }}
-    {% else %}
-    _(no item found)_
-    {% endfor %}
-    """)
-)
-
-_TPL_TABLE = Template(
-    textwrap.dedent("""\
-    {%- if rows %}
-    | {{ cols | join(' | ') }} |
-    |{%- for _ in cols %} --- |{%- endfor %}
-    {%- for r in rows %}
-    | {%- for c in cols %} {{ r.get(c, "") }} |{%- endfor %}
-    {%- endfor %}
-    {%- else %}
-    _(no data)_
-    {%- endif %}
-""")
-)
-
-
-def gh_summary_path() -> Path | None:
-    """Return the Path to the GitHub step summary file, or None if not set."""
-    p = os.environ.get("GITHUB_STEP_SUMMARY")
-    return Path(p) if p else None
-
-
-def write_gh_step_summary(md: str, *, append_content: bool = True) -> bool:
-    """
-    Write Markdown content to the GitHub Step Summary file if GITHUB_STEP_SUMMARY is set.
-    append_content: default true, if True, append to the end of the file, else overwrite the whole file
-
-    Returns:
-        True if written successfully (in GitHub Actions environment),
-        False if skipped (e.g., running locally where the variable is not set).
-    """
-    sp = gh_summary_path()
-    if not sp:
-        logger.info("[gh-summary] GITHUB_STEP_SUMMARY not set, skipping write.")
-        return False
-
-    md_clean = textwrap.dedent(md).strip() + "\n"
-
-    mode = "a" if append_content else "w"
-    with sp.open(mode, encoding="utf-8") as f:
-        f.write(md_clean)
-    return True
-
-
-def md_heading(text: str, level: int = 2) -> str:
-    """Generate a Markdown heading string with the given level (1-6)."""
-    return f"{'#' * max(1, min(level, 6))} {text}\n"
-
-
-def md_details(summary: str, content: str) -> str:
-    """Generate a collapsible <details> block with a summary and inner content."""
-    return f"<details>\n<summary>{summary}</summary>\n\n{content}\n\n</details>\n"
-
-
-def summarize_content_from_file(
-    output_dir: Path,
-    freeze_file: str,
-    title: str = "Content from file",
-    code_lang: str = "",  # e.g. "text" or "ini"
-) -> bool:
-    f = Path(output_dir) / freeze_file
-    if not f.exists():
-        return False
-    content = f.read_text(encoding="utf-8").strip()
-    md = render_content(content, title=title, lang=code_lang)
-    return write_gh_step_summary(md)
-
-
-def summarize_wheels(path: Path, title: str = "Wheels", max_depth: int = 3):
-    items = get_wheels(path, max_depth=max_depth)
-    if not items:
-        return False
-    md = render_list(items, title=title)
-    return write_gh_step_summary(md)
-
-
-def md_kv_table(rows: Iterable[Mapping[str, str | int | float]]) -> str:
-    """
-    Render a list of dicts as a Markdown table using Jinja template.
-    """
-    rows = list(rows)
-    cols = list({k for r in rows for k in r.keys()})
-    md = _TPL_TABLE.render(cols=cols, rows=rows).strip() + "\n"
-    return md
-
-
-def render_list(
-    items: Iterable[str],
-    *,
-    title: str = "List",
-) -> str:
-    tpl = _TPL_LIST_ITEMS
-    md = tpl.render(title=title, items=items)
-    return md
-
-
-def render_content(
-    content: str,
-    *,
-    title: str = "Content",
-    lang: str = "text",
-) -> str:
-    tpl = _TPL_CONTENT
-    md = tpl.render(title=title, content=content, lang=lang)
-    return md

.ci/lumen_cli/cli/lib/common/git_helper.py

@@ -1,69 +0,0 @@
-"""
-Git Utility helpers for CLI tasks.
-"""
-
-import logging
-from pathlib import Path
-
-from cli.lib.common.path_helper import remove_dir
-from git import GitCommandError, RemoteProgress, Repo
-
-
-logger = logging.getLogger(__name__)
-
-
-class PrintProgress(RemoteProgress):
-    """Simple progress logger for git operations."""
-
-    def __init__(self, interval: int = 5):
-        super().__init__()
-        self._last_percent = -1
-        self._interval = interval
-
-    def update(self, op_code, cur, max=None, message=""):
-        msg = self._cur_line or message
-        if max and cur:
-            percent = int(cur / max * 100)
-            if percent != self._last_percent and percent % self._interval == 0:
-                self._last_percent = percent
-                logger.info("Progress: %d%% - %s", percent, msg)
-        elif msg:
-            logger.info(msg)
-
-
-def clone_external_repo(target: str, repo: str, dst: str = "", update_submodules=False):
-    """Clone repository with pinned commit and optional submodules."""
-    dst = dst or target
-
-    try:
-        logger.info("Cloning %s to %s", target, dst)
-
-        # Clone and fetch
-        remove_dir(dst)
-        r = Repo.clone_from(repo, dst, progress=PrintProgress())
-        r.git.fetch("--all", "--tags")
-
-        # Checkout pinned commit
-        commit = get_post_build_pinned_commit(target)
-        logger.info("Checking out pinned %s commit %s", target, commit)
-        r.git.checkout(commit)
-
-        # Update submodules if requested
-        if update_submodules and r.submodules:
-            logger.info("Updating %d submodule(s)", len(r.submodules))
-            for sm in r.submodules:
-                sm.update(init=True, recursive=True, progress=PrintProgress())
-
-        logger.info("Successfully cloned %s", target)
-        return r, commit
-
-    except GitCommandError as e:
-        logger.error("Git operation failed: %s", e)
-        raise
-
-
-def get_post_build_pinned_commit(name: str, prefix=".github/ci_commit_pins") -> str:
-    path = Path(prefix) / f"{name}.txt"
-    if not path.exists():
-        raise FileNotFoundError(f"Pin file not found: {path}")
-    return path.read_text(encoding="utf-8").strip()

.ci/lumen_cli/cli/lib/common/logger.py

@@ -1,14 +0,0 @@
-"""
-Logger Utility helpers for CLI tasks.
-"""
-
-import logging
-import sys
-
-
-def setup_logging(level: int = logging.INFO):
-    logging.basicConfig(
-        level=level,
-        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-        stream=sys.stdout,
-    )

.ci/lumen_cli/cli/lib/common/path_helper.py

@@ -1,62 +0,0 @@
-"""Path utility helpers for CLI tasks."""
-
-import logging
-import shutil
-from pathlib import Path
-from typing import Union
-
-
-logger = logging.getLogger(__name__)
-
-
-def get_path(path: Union[str, Path], resolve: bool = False) -> Path:
-    """Convert to Path object, optionally resolving to absolute path."""
-    if not path:
-        raise ValueError("Path cannot be None or empty")
-    result = Path(path)
-    return result.resolve() if resolve else result
-
-
-def ensure_dir_exists(path: Union[str, Path]) -> Path:
-    """Create directory if it doesn't exist."""
-    path_obj = get_path(path)
-    path_obj.mkdir(parents=True, exist_ok=True)
-    return path_obj
-
-
-def remove_dir(path: Union[str, Path, None]) -> None:
-    """Remove directory if it exists."""
-    if not path:
-        return
-    path_obj = get_path(path)
-    if path_obj.exists():
-        shutil.rmtree(path_obj)
-
-
-def force_create_dir(path: Union[str, Path]) -> Path:
-    """Remove directory if exists, then create fresh empty directory."""
-    remove_dir(path)
-    return ensure_dir_exists(path)
-
-
-def copy(src: Union[str, Path], dst: Union[str, Path]) -> None:
-    """Copy file or directory from src to dst."""
-    src_path = get_path(src, resolve=True)
-    dst_path = get_path(dst, resolve=True)
-
-    if not src_path.exists():
-        raise FileNotFoundError(f"Source does not exist: {src_path}")
-
-    dst_path.parent.mkdir(parents=True, exist_ok=True)
-
-    if src_path.is_file():
-        shutil.copy2(src_path, dst_path)
-    elif src_path.is_dir():
-        shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
-    else:
-        raise ValueError(f"Unsupported path type: {src_path}")
-
-
-def is_path_exist(path: Union[str, Path, None]) -> bool:
-    """Check if path exists."""
-    return bool(path and get_path(path).exists())

.ci/lumen_cli/cli/lib/common/pip_helper.py

@@ -1,71 +0,0 @@
-import glob
-import logging
-import shlex
-import shutil
-import sys
-from collections.abc import Iterable
-from importlib.metadata import PackageNotFoundError, version  # noqa: UP035
-from typing import Optional, Union
-
-from cli.lib.common.utils import run_command
-
-
-logger = logging.getLogger(__name__)
-
-
-def pip_install_packages(
-    packages: Iterable[str] = (),
-    env=None,
-    *,
-    requirements: Optional[str] = None,
-    constraints: Optional[str] = None,
-    prefer_uv: bool = False,
-) -> None:
-    use_uv = prefer_uv and shutil.which("uv") is not None
-    base = (
-        [sys.executable, "-m", "uv", "pip", "install"]
-        if use_uv
-        else [sys.executable, "-m", "pip", "install"]
-    )
-    cmd = base[:]
-    if requirements:
-        cmd += ["-r", requirements]
-    if constraints:
-        cmd += ["-c", constraints]
-    cmd += list(packages)
-    logger.info("pip installing packages: %s", " ".join(map(shlex.quote, cmd)))
-    run_command(" ".join(map(shlex.quote, cmd)), env=env)
-
-
-def pip_install_first_match(pattern: str, extras: Optional[str] = None, pref_uv=False):
-    wheel = first_matching_pkg(pattern)
-    target = f"{wheel}[{extras}]" if extras else wheel
-    logger.info("Installing %s...", target)
-    pip_install_packages([target], prefer_uv=pref_uv)
-
-
-def run_python(args: Union[str, list[str]], env=None):
-    """
-    Run the python in the current environment.
-    """
-    if isinstance(args, str):
-        args = shlex.split(args)
-    cmd = [sys.executable] + args
-    run_command(" ".join(map(shlex.quote, cmd)), env=env)
-
-
-def pkg_exists(name: str) -> bool:
-    try:
-        pkg_version = version(name)
-        logger.info("%s already exist with version: %s", name, pkg_version)
-        return True
-    except PackageNotFoundError:
-        logger.info("%s is not installed", name)
-        return False
-
-
-def first_matching_pkg(pattern: str) -> str:
-    matches = sorted(glob.glob(pattern))
-    if not matches:
-        raise FileNotFoundError(f"No wheel matching: {pattern}")
-    return matches[0]

.ci/lumen_cli/cli/lib/common/utils.py

@@ -1,139 +0,0 @@
-"""
-General Utility helpers for CLI tasks.
-"""
-
-import logging
-import os
-import shlex
-import subprocess
-import sys
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Optional
-
-
-logger = logging.getLogger(__name__)
-
-
-def run_command(
-    cmd: str,
-    use_shell: bool = False,
-    log_cmd: bool = True,
-    cwd: Optional[str] = None,
-    env: Optional[dict] = None,
-    check: bool = True,
-) -> int:
-    """Run a command with optional shell execution."""
-    if use_shell:
-        args = cmd
-        log_prefix = "[shell]"
-        executable = "/bin/bash"
-    else:
-        args = shlex.split(cmd)
-        log_prefix = "[cmd]"
-        executable = None
-
-    if log_cmd:
-        display_cmd = cmd if use_shell else " ".join(args)
-        logger.info("%s %s", log_prefix, display_cmd)
-
-    run_env = {**os.environ, **(env or {})}
-
-    proc = subprocess.run(
-        args,
-        shell=use_shell,
-        executable=executable,
-        stdout=sys.stdout,
-        stderr=sys.stderr,
-        cwd=cwd,
-        env=run_env,
-        check=False,
-    )
-
-    if check and proc.returncode != 0:
-        logger.error(
-            "%s Command failed (exit %s): %s", log_prefix, proc.returncode, cmd
-        )
-        raise subprocess.CalledProcessError(
-            proc.returncode, args if not use_shell else cmd
-        )
-
-    return proc.returncode
-
-
-def str2bool(value: Optional[str]) -> bool:
-    """Convert environment variables to boolean values."""
-    if not value:
-        return False
-    if not isinstance(value, str):
-        raise ValueError(
-            f"Expected a string value for boolean conversion, got {type(value)}"
-        )
-    value = value.strip().lower()
-
-    true_value_set = {"1", "true", "t", "yes", "y", "on", "enable", "enabled", "found"}
-    false_value_set = {"0", "false", "f", "no", "n", "off", "disable"}
-
-    if value in true_value_set:
-        return True
-    if value in false_value_set:
-        return False
-    raise ValueError(f"Invalid string value for boolean conversion: {value}")
-
-
-@contextmanager
-def temp_environ(updates: dict[str, str]):
-    """
-    Temporarily set environment variables and restore them after the block.
-    Args:
-        updates: Dict of environment variables to set.
-    """
-    missing = object()
-    old: dict[str, str | object] = {k: os.environ.get(k, missing) for k in updates}
-    try:
-        os.environ.update(updates)
-        yield
-    finally:
-        for k, v in old.items():
-            if v is missing:
-                os.environ.pop(k, None)
-            else:
-                os.environ[k] = v  # type: ignore[arg-type]
-
-
-@contextmanager
-def working_directory(path: str):
-    """
-    Temporarily change the working directory inside a context.
-    """
-    if not path:
-        # No-op context
-        yield
-        return
-    prev_cwd = os.getcwd()
-    try:
-        os.chdir(path)
-        yield
-    finally:
-        os.chdir(prev_cwd)
-
-
-def get_wheels(
-    output_dir: Path,
-    max_depth: Optional[int] = None,
-) -> list[str]:
-    """Return a list of wheels found in the given output directory."""
-    root = Path(output_dir)
-    if not root.exists():
-        return []
-    items = []
-    for dirpath, _, filenames in os.walk(root):
-        depth = Path(dirpath).relative_to(root).parts
-        if max_depth is not None and len(depth) > max_depth:
-            continue
-        for fname in sorted(filenames):
-            if fname.endswith(".whl"):
-                pkg = fname.split("-")[0]
-                relpath = str((Path(dirpath) / fname).relative_to(root))
-                items.append({"pkg": pkg, "relpath": relpath})
-    return items

.ci/lumen_cli/cli/lib/core/vllm/lib.py

@@ -1,292 +0,0 @@
-import logging
-import os
-import textwrap
-from typing import Any
-
-from cli.lib.common.gh_summary import write_gh_step_summary
-from cli.lib.common.git_helper import clone_external_repo
-from cli.lib.common.pip_helper import pip_install_packages
-from cli.lib.common.utils import run_command, temp_environ, working_directory
-from jinja2 import Template
-
-
-logger = logging.getLogger(__name__)
-
-_TPL_VLLM_INFO = Template(
-    textwrap.dedent("""\
-    ##  Vllm against Pytorch CI Test Summary
-    **Vllm Commit**: [{{ vllm_commit }}](https://github.com/vllm-project/vllm/commit/{{ vllm_commit }})
-    {%- if torch_sha %}
-    **Pytorch Commit**: [{{ torch_sha }}](https://github.com/pytorch/pytorch/commit/{{ torch_sha }})
-    {%- endif %}
-""")
-)
-
-
-def sample_vllm_test_library():
-    """
-    Simple sample to unblock the vllm ci development, which is mimic to
-    https://github.com/vllm-project/vllm/blob/main/.buildkite/test-pipeline.yaml
-    see run_test_plan for more details
-    """
-    # TODO(elainewy): Read from yaml file to handle the env and tests for vllm
-    return {
-        "vllm_basic_correctness_test": {
-            "title": "Basic Correctness Test",
-            "id": "vllm_basic_correctness_test",
-            "env_vars": {
-                "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
-            },
-            "steps": [
-                "pytest -v -s basic_correctness/test_cumem.py",
-                "pytest -v -s basic_correctness/test_basic_correctness.py",
-                "pytest -v -s basic_correctness/test_cpu_offload.py",
-            ],
-        },
-        "vllm_basic_models_test": {
-            "title": "Basic models test",
-            "id": "vllm_basic_models_test",
-            "steps": [
-                "pytest -v -s models/test_transformers.py",
-                "pytest -v -s models/test_registry.py",
-                "pytest -v -s models/test_utils.py",
-                "pytest -v -s models/test_vision.py",
-                "pytest -v -s models/test_initialization.py",
-            ],
-        },
-        "vllm_entrypoints_test": {
-            "title": "Entrypoints Test ",
-            "id": "vllm_entrypoints_test",
-            "env_vars": {
-                "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
-            },
-            "steps": [
-                " ".join(
-                    [
-                        "pytest",
-                        "-v",
-                        "-s",
-                        "entrypoints/llm",
-                        "--ignore=entrypoints/llm/test_generate.py",
-                        "--ignore=entrypoints/llm/test_collective_rpc.py",
-                    ]
-                ),
-                "pytest -v -s entrypoints/llm/test_generate.py",
-                "pytest -v -s entrypoints/offline_mode",
-            ],
-        },
-        "vllm_regression_test": {
-            "title": "Regression Test",
-            "id": "vllm_regression_test",
-            "package_install": ["modelscope"],
-            "steps": [
-                "pytest -v -s test_regression.py",
-            ],
-        },
-        "vllm_lora_tp_test_distributed": {
-            "title": "LoRA TP Test (Distributed)",
-            "id": "vllm_lora_tp_test_distributed",
-            "env_vars": {
-                "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
-            },
-            "num_gpus": 4,
-            "steps": [
-                "pytest -v -s -x lora/test_chatglm3_tp.py",
-                "pytest -v -s -x lora/test_llama_tp.py",
-                "pytest -v -s -x lora/test_llm_with_multi_loras.py",
-            ],
-        },
-        "vllm_distributed_test_28_failure_test": {
-            "title": "Distributed Tests (2 GPUs) pytorch 2.8 release failure",
-            "id": "vllm_distributed_test_28_failure_test",
-            "env_vars": {
-                "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
-            },
-            "num_gpus": 4,
-            "steps": [
-                "pytest -v -s distributed/test_sequence_parallel.py",
-            ],
-        },
-        "vllm_lora_28_failure_test": {
-            "title": "LoRA pytorch 2.8 failure test",
-            "id": "vllm_lora_28_failure_test",
-            "steps": ["pytest -v lora/test_quant_model.py"],
-        },
-        "vllm_multi_model_processor_test": {
-            "title": "Multi-Modal Processor Test",
-            "id": "vllm_multi_model_processor_test",
-            "package_install": ["git+https://github.com/TIGER-AI-Lab/Mantis.git"],
-            "steps": [
-                "pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py",
-            ],
-        },
-        "vllm_multi_model_test_28_failure_test": {
-            "title": "Multi-Model Test (Failed 2.8 release)",
-            "id": "vllm_multi_model_test_28_failure_test",
-            "package_install": ["git+https://github.com/TIGER-AI-Lab/Mantis.git"],
-            "steps": [
-                "pytest -v -s models/multimodal/generation/test_voxtral.py",
-                "pytest -v -s models/multimodal/pooling",
-            ],
-        },
-        "vllm_pytorch_compilation_unit_tests": {
-            "title": "PyTorch Compilation Unit Tests",
-            "id": "vllm_pytorch_compilation_unit_tests",
-            "steps": [
-                "pytest -v -s compile/test_pass_manager.py",
-                "pytest -v -s compile/test_fusion.py",
-                "pytest -v -s compile/test_fusion_attn.py",
-                "pytest -v -s compile/test_silu_mul_quant_fusion.py",
-                "pytest -v -s compile/test_sequence_parallelism.py",
-                "pytest -v -s compile/test_async_tp.py",
-                "pytest -v -s compile/test_fusion_all_reduce.py",
-                "pytest -v -s compile/test_decorator.py",
-            ],
-        },
-        "vllm_languagde_model_test_extended_generation_28_failure_test": {
-            "title": "Language Models Test (Extended Generation) 2.8 release failure",
-            "id": "vllm_languagde_model_test_extended_generation_28_failure_test",
-            "package_install": [
-                "--no-build-isolation",
-                "git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8",
-            ],
-            "steps": [
-                "pytest -v -s models/language/generation/test_mistral.py",
-            ],
-        },
-        "vllm_distributed_test_2_gpu_28_failure_test": {
-            "title": "Distributed Tests (2 GPUs) pytorch 2.8 release failure",
-            "id": "vllm_distributed_test_2_gpu_28_failure_test",
-            "env_vars": {
-                "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
-            },
-            "num_gpus": 4,
-            "steps": [
-                "pytest -v -s distributed/test_sequence_parallel.py",
-            ],
-        },
-        # TODO(elainewy):need to add g6 with 4 gpus to run this test
-        "vllm_lora_test": {
-            "title": "LoRA Test %N",
-            "id": "lora_test",
-            "parallelism": 4,
-            "steps": [
-                "echo '[checking] list sharded lora tests:'",
-                " ".join(
-                    [
-                        "pytest -q --collect-only lora",
-                        "--shard-id=$$BUILDKITE_PARALLEL_JOB",
-                        "--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT",
-                        "--ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py",
-                    ]
-                ),
-                "echo '[checking] Done. list lora tests'",
-                " ".join(
-                    [
-                        "pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB",
-                        "--num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT",
-                        "--ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py",
-                    ]
-                ),
-            ],
-        },
-    }
-
-
-def check_parallelism(tests: Any, title: str, shard_id: int = 0, num_shards: int = 0):
-    """
-    a method to check if the test plan is parallelism or not.
-    """
-    parallelism = int(tests.get("parallelism", "0"))
-    is_parallel = parallelism and parallelism > 1
-
-    if not is_parallel:
-        return False
-
-    if shard_id > num_shards:
-        raise RuntimeError(
-            f"Test {title} expects {num_shards} shards, but invalid {shard_id} is provided"
-        )
-
-    if num_shards != parallelism:
-        raise RuntimeError(
-            f"Test {title} expects {parallelism} shards, but invalid {num_shards} is provided"
-        )
-
-    return True
-
-
-def run_test_plan(
-    test_plan: str,
-    test_target: str,
-    tests_map: dict[str, Any],
-    shard_id: int = 0,
-    num_shards: int = 0,
-):
-    """
-    a method to run list of tests based on the test plan.
-    """
-    logger.info("run %s tests.....", test_target)
-    if test_plan not in tests_map:
-        raise RuntimeError(
-            f"test {test_plan} not found, please add it to test plan pool"
-        )
-    tests = tests_map[test_plan]
-    pkgs = tests.get("package_install", [])
-    title = tests.get("title", "unknown test")
-
-    is_parallel = check_parallelism(tests, title, shard_id, num_shards)
-    if is_parallel:
-        title = title.replace("%N", f"{shard_id}/{num_shards}")
-
-    logger.info("Running tests: %s", title)
-    if pkgs:
-        logger.info("Installing packages: %s", pkgs)
-        pip_install_packages(packages=pkgs, prefer_uv=True)
-    with (
-        working_directory(tests.get("working_directory", "tests")),
-        temp_environ(tests.get("env_vars", {})),
-    ):
-        failures = []
-        for step in tests["steps"]:
-            logger.info("Running step: %s", step)
-            if is_parallel:
-                step = replace_buildkite_placeholders(step, shard_id, num_shards)
-                logger.info("Running parallel step: %s", step)
-            code = run_command(cmd=step, check=False, use_shell=True)
-            if code != 0:
-                failures.append(step)
-            logger.info("Finish running step: %s", step)
-        if failures:
-            logger.error("Failed tests: %s", failures)
-            raise RuntimeError(f"{len(failures)} pytest runs failed: {failures}")
-        logger.info("Done. All tests passed")
-
-
-def clone_vllm(dst: str = "vllm"):
-    _, commit = clone_external_repo(
-        target="vllm",
-        repo="https://github.com/vllm-project/vllm.git",
-        dst=dst,
-        update_submodules=True,
-    )
-    return commit
-
-
-def replace_buildkite_placeholders(step: str, shard_id: int, num_shards: int) -> str:
-    mapping = {
-        "$$BUILDKITE_PARALLEL_JOB_COUNT": str(num_shards),
-        "$$BUILDKITE_PARALLEL_JOB": str(shard_id),
-    }
-    for k in sorted(mapping, key=len, reverse=True):
-        step = step.replace(k, mapping[k])
-    return step
-
-
-def summarize_build_info(vllm_commit: str) -> bool:
-    torch_sha = os.getenv("GITHUB_SHA")
-    md = (
-        _TPL_VLLM_INFO.render(vllm_commit=vllm_commit, torch_sha=torch_sha).strip()
-        + "\n"
-    )
-    return write_gh_step_summary(md)

.ci/lumen_cli/cli/lib/core/vllm/vllm_build.py

@@ -1,296 +0,0 @@
-import logging
-import os
-import textwrap
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Optional
-
-from cli.lib.common.cli_helper import BaseRunner
-from cli.lib.common.docker_helper import local_image_exists
-from cli.lib.common.envs_helper import (
-    env_bool_field,
-    env_path_field,
-    env_str_field,
-    with_params_help,
-)
-from cli.lib.common.gh_summary import (
-    gh_summary_path,
-    summarize_content_from_file,
-    summarize_wheels,
-)
-from cli.lib.common.path_helper import (
-    copy,
-    ensure_dir_exists,
-    force_create_dir,
-    get_path,
-    is_path_exist,
-)
-from cli.lib.common.utils import run_command
-from cli.lib.core.vllm.lib import clone_vllm, summarize_build_info
-
-
-logger = logging.getLogger(__name__)
-
-
-# Default path for docker build artifacts
-_DEFAULT_RESULT_PATH = "./shared"
-
-# Temp folder in vllm work place to cp torch whls in vllm work directory for docker build
-_VLLM_TEMP_FOLDER = "tmp"
-
-
-@dataclass
-class VllmBuildParameters:
-    """
-    Parameters defining the vllm external input configurations.
-    Combine with VllmDockerBuildArgs to define the vllm build environment
-    """
-
-    # USE_TORCH_WHEEL: when true, use local Torch wheels; requires TORCH_WHEELS_PATH.
-    # Otherwise docker build pull torch nightly during build
-    # TORCH_WHEELS_PATH: directory containing local torch wheels when use_torch_whl is True
-    use_torch_whl: bool = env_bool_field("USE_TORCH_WHEEL", True)
-    torch_whls_path: Path = env_path_field("TORCH_WHEELS_PATH", "./dist")
-
-    # USE_LOCAL_BASE_IMAGE: when true, use an existing local Docker base image; requires BASE_IMAGE
-    # Otherwise, pull dockerfile's default image remotely
-    # BASE_IMAGE: name:tag (only needed when use_local_base_image is True)
-    use_local_base_image: bool = env_bool_field("USE_LOCAL_BASE_IMAGE", True)
-    base_image: str = env_str_field("BASE_IMAGE")
-
-    # USE_LOCAL_DOCKERFILE: when true("1"), use a local Dockerfile; requires DOCKERFILE_PATH.
-    # otherwise, use vllm's default dockerfile.torch_nightly for build
-    # DOCKERFILE_PATH: path to Dockerfile used when use_local_dockerfile is True"
-    use_local_dockerfile: bool = env_bool_field("USE_LOCAL_DOCKERFILE", True)
-    dockerfile_path: Path = env_path_field(
-        "DOCKERFILE_PATH", ".github/ci_configs/vllm/Dockerfile.tmp_vllm"
-    )
-
-    # the cleaning script to remove torch dependencies from pip
-    cleaning_script: Path = env_path_field(
-        "cleaning_script", ".github/ci_configs/vllm/use_existing_torch.py"
-    )
-
-    # OUTPUT_DIR: where docker buildx (local exporter) will write artifacts
-    output_dir: Path = env_path_field("OUTPUT_DIR", "external/vllm")
-
-    # --- Build args ----------------------------------------------------------
-    target_stage: str = env_str_field("TARGET_STAGE", "export-wheels")
-
-    tag_name: str = env_str_field("TAG", "vllm-wheels")
-
-    cuda_version: str = env_str_field("CUDA_VERSION", "12.8.1")
-
-    python_version: str = env_str_field("PYTHON_VERSION", "3.12")
-
-    max_jobs: str = env_str_field("MAX_JOBS", "64")
-
-    sccache_bucket: str = env_str_field("SCCACHE_BUCKET")
-
-    sccache_region: str = env_str_field("SCCACHE_REGION")
-
-    torch_cuda_arch_list: str = env_str_field("TORCH_CUDA_ARCH_LIST", "8.9")
-
-    def __post_init__(self):
-        checks = [
-            (
-                self.use_torch_whl,  # flag
-                True,  # trigger_value
-                "torch_whls_path",  # resource
-                is_path_exist,  # check_func
-                "TORCH_WHEELS_PATH is not provided, but USE_TORCH_WHEEL is set to 1",
-            ),
-            (
-                self.use_local_base_image,
-                True,
-                "base_image",
-                local_image_exists,
-                f"BASE_IMAGE {self.base_image} does not found, but USE_LOCAL_BASE_IMAGE is set to 1",
-            ),
-            (
-                self.use_local_dockerfile,
-                True,
-                "dockerfile_path",
-                is_path_exist,
-                " DOCKERFILE_PATH path does not found, but USE_LOCAL_DOCKERFILE is set to 1",
-            ),
-        ]
-        for flag, trigger_value, attr_name, check_func, error_msg in checks:
-            value = getattr(self, attr_name)
-            if flag == trigger_value:
-                if not value or not check_func(value):
-                    raise ValueError(error_msg)
-            else:
-                logger.info("flag  %s is not set", flag)
-        if not self.output_dir:
-            raise ValueError("missing required output_dir")
-
-
-@with_params_help(VllmBuildParameters)
-class VllmBuildRunner(BaseRunner):
-    """
-    Build vLLM using docker buildx.
-
-    Environment variable options:
-        "USE_TORCH_WHEEL":      "1: use local wheels; 0: pull nightly from pypi",
-        "TORCH_WHEELS_PATH":    "Path to local wheels (when USE_TORCH_WHEEL=1)",
-
-        "USE_LOCAL_BASE_IMAGE": "1: use local base image; 0: default image",
-         "BASE_IMAGE":           "name:tag to indicate base image the dockerfile depends on (when USE_LOCAL_BASE_IMAGE=1)",
-
-        "USE_LOCAL_DOCKERFILE": "1: use local Dockerfile; 0: vllm repo default dockerfile.torch_nightly",
-        "DOCKERFILE_PATH":      "Path to Dockerfile (when USE_LOCAL_DOCKERFILE=1)",
-
-        "OUTPUT_DIR":           "e.g. './shared'",
-
-        "TORCH_CUDA_ARCH_LIST": "e.g. '8.0' or '8.0;9.0'",
-        "CUDA_VERSION":         "e.g. '12.8.1'",
-        "PYTHON_VERSION":       "e.g. '3.12'",
-        "MAX_JOBS":             "e.g. '64'",
-        "SCCACHE_BUCKET":       "e.g. 'my-bucket'",
-        "SCCACHE_REGION":       "e.g. 'us-west-2'",
-    """
-
-    def __init__(self, args=None):
-        self.work_directory = "vllm"
-
-    def run(self):
-        """
-        main function to run vllm build
-        1. prepare vllm build environment
-        2. prepare the docker build command args
-        3. run docker build
-        """
-        inputs = VllmBuildParameters()
-        logger.info("Running vllm build with inputs: %s", inputs)
-        vllm_commit = clone_vllm()
-
-        self.cp_torch_cleaning_script(inputs)
-        self.cp_dockerfile_if_exist(inputs)
-        # cp torch wheels from root direct to vllm workspace if exist
-        self.cp_torch_whls_if_exist(inputs)
-
-        # make sure the output dir to store the build artifacts exist
-        ensure_dir_exists(Path(inputs.output_dir))
-
-        cmd = self._generate_docker_build_cmd(inputs)
-        logger.info("Running docker build: \n %s", cmd)
-
-        try:
-            run_command(cmd, cwd="vllm", env=os.environ.copy())
-        finally:
-            self.genearte_vllm_build_summary(vllm_commit, inputs)
-
-    def genearte_vllm_build_summary(
-        self, vllm_commit: str, inputs: VllmBuildParameters
-    ):
-        if not gh_summary_path():
-            return logger.info("Skipping, not detect GH Summary env var....")
-        logger.info("Generate GH Summary ...")
-        # summarize vllm build info
-        summarize_build_info(vllm_commit)
-
-        # summarize vllm build artifacts
-        vllm_artifact_dir = inputs.output_dir / "wheels"
-        summarize_content_from_file(
-            vllm_artifact_dir,
-            "build_summary.txt",
-            title="Vllm build env pip package summary",
-        )
-        summarize_wheels(
-            inputs.torch_whls_path, max_depth=3, title="Torch Wheels Artifacts"
-        )
-        summarize_wheels(vllm_artifact_dir, max_depth=3, title="Vllm Wheels Artifacts")
-
-    def cp_torch_whls_if_exist(self, inputs: VllmBuildParameters) -> str:
-        if not inputs.use_torch_whl:
-            return ""
-        tmp_dir = f"./{self.work_directory}/{_VLLM_TEMP_FOLDER}"
-        tmp_path = Path(tmp_dir)
-        force_create_dir(tmp_path)
-        copy(inputs.torch_whls_path, tmp_dir)
-        return tmp_dir
-
-    def cp_torch_cleaning_script(self, inputs: VllmBuildParameters):
-        script = get_path(inputs.cleaning_script, resolve=True)
-        vllm_script = Path(f"./{self.work_directory}/use_existing_torch.py")
-        copy(script, vllm_script)
-
-    def cp_dockerfile_if_exist(self, inputs: VllmBuildParameters):
-        if not inputs.use_local_dockerfile:
-            logger.info("using vllm default dockerfile.torch_nightly for build")
-            return
-        dockerfile_path = get_path(inputs.dockerfile_path, resolve=True)
-        vllm_torch_dockerfile = Path(
-            f"./{self.work_directory}/docker/Dockerfile.nightly_torch"
-        )
-        copy(dockerfile_path, vllm_torch_dockerfile)
-
-    def get_result_path(self, path):
-        """
-        Get the absolute path of the result path
-        """
-        if not path:
-            path = _DEFAULT_RESULT_PATH
-        abs_path = get_path(path, resolve=True)
-        return abs_path
-
-    def _get_torch_wheel_path_arg(self, torch_whl_dir: Optional[Path]) -> str:
-        if not torch_whl_dir:
-            return ""
-        return f"--build-arg TORCH_WHEELS_PATH={_VLLM_TEMP_FOLDER}"
-
-    def _get_base_image_args(self, inputs: VllmBuildParameters) -> tuple[str, str, str]:
-        """
-        Returns:
-            - base_image_arg: docker buildx arg string for base image
-            - final_base_image_arg:  docker buildx arg string for vllm-base stage
-            - pull_flag: --pull=true or --pull=false depending on whether the image exists locally
-        """
-        if not inputs.use_local_base_image:
-            return "", "", ""
-
-        base_image = inputs.base_image
-
-        # set both base image and final base image to the same local image
-        base_image_arg = f"--build-arg BUILD_BASE_IMAGE={base_image}"
-        final_base_image_arg = f"--build-arg FINAL_BASE_IMAGE={base_image}"
-
-        if local_image_exists(base_image):
-            pull_flag = "--pull=false"
-            return base_image_arg, final_base_image_arg, pull_flag
-        logger.info(
-            "[INFO] Local image not found:%s will try to pull from remote", {base_image}
-        )
-        return base_image_arg, final_base_image_arg, ""
-
-    def _generate_docker_build_cmd(
-        self,
-        inputs: VllmBuildParameters,
-    ) -> str:
-        base_image_arg, final_base_image_arg, pull_flag = self._get_base_image_args(
-            inputs
-        )
-        torch_arg = self._get_torch_wheel_path_arg(inputs.torch_whls_path)
-
-        return textwrap.dedent(
-            f"""
-            docker buildx build \
-                --output type=local,dest={inputs.output_dir} \
-                -f docker/Dockerfile.nightly_torch \
-                {pull_flag} \
-                {torch_arg} \
-                {base_image_arg} \
-                {final_base_image_arg} \
-                --build-arg max_jobs={inputs.max_jobs} \
-                --build-arg CUDA_VERSION={inputs.cuda_version} \
-                --build-arg PYTHON_VERSION={inputs.python_version} \
-                --build-arg USE_SCCACHE={int(bool(inputs.sccache_bucket and inputs.sccache_region))} \
-                --build-arg SCCACHE_BUCKET_NAME={inputs.sccache_bucket} \
-                --build-arg SCCACHE_REGION_NAME={inputs.sccache_region} \
-                --build-arg torch_cuda_arch_list='{inputs.torch_cuda_arch_list}' \
-                --target {inputs.target_stage} \
-                -t {inputs.tag_name} \
-                --progress=plain .
-        """
-        ).strip()

.ci/lumen_cli/cli/lib/core/vllm/vllm_test.py

@@ -1,280 +0,0 @@
-import logging
-import os
-import re
-import subprocess
-import sys
-from collections.abc import Iterable
-from dataclasses import dataclass
-from enum import Enum
-from pathlib import Path
-from typing import Any
-
-from cli.lib.common.cli_helper import BaseRunner
-from cli.lib.common.envs_helper import env_path_field, env_str_field, get_env
-from cli.lib.common.path_helper import copy, get_path, remove_dir
-from cli.lib.common.pip_helper import (
-    pip_install_first_match,
-    pip_install_packages,
-    pkg_exists,
-    run_python,
-)
-from cli.lib.common.utils import run_command, working_directory
-from cli.lib.core.vllm.lib import clone_vllm, run_test_plan, sample_vllm_test_library
-
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class VllmTestParameters:
-    """
-    Parameters defining the vllm external test input
-
-    !!!DO NOT ADD SECRETS IN THIS CLASS!!!
-    you can put environment variable name in VllmTestParameters if it's not the same as the secret one
-    fetch secrests directly from env variables during runtime
-    """
-
-    torch_whls_path: Path = env_path_field("WHEELS_PATH", "./dist")
-
-    vllm_whls_path: Path = env_path_field(
-        "VLLM_WHEELS_PATH", "./dist/external/vllm/wheels"
-    )
-
-    torch_cuda_arch_list: str = env_str_field("TORCH_CUDA_ARCH_LIST", "8.9")
-
-    cleaning_script: Path = env_path_field(
-        "cleaning_script", ".github/ci_configs/vllm/use_existing_torch.py"
-    )
-
-    def __post_init__(self):
-        if not self.torch_whls_path.exists():
-            raise ValueError("missing torch_whls_path")
-        if not self.vllm_whls_path.exists():
-            raise ValueError("missing vllm_whls_path")
-
-
-class TestInpuType(Enum):
-    TEST_PLAN = "test_plan"
-    UNKNOWN = "unknown"
-
-
-class VllmTestRunner(BaseRunner):
-    def __init__(self, args: Any):
-        self.work_directory = "vllm"
-        self.test_plan = ""
-        self.test_type = TestInpuType.UNKNOWN
-
-        self.shard_id = args.shard_id
-        self.num_shards = args.num_shards
-
-        if args.test_plan:
-            self.test_plan = args.test_plan
-            self.test_type = TestInpuType.TEST_PLAN
-
-        # Matches the structeur in the artifacts.zip from torcb build
-        self.TORCH_WHL_PATH_REGEX = "torch*.whl"
-        self.TORCH_WHL_EXTRA = "opt-einsum"
-        self.TORCH_ADDITIONAL_WHLS_REGEX = [
-            "vision/torchvision*.whl",
-            "audio/torchaudio*.whl",
-        ]
-
-        # Match the structure of the artifacts.zip from vllm external build
-        self.VLLM_TEST_WHLS_REGEX = [
-            "xformers/*.whl",
-            "vllm/vllm*.whl",
-            "flashinfer-python/flashinfer*.whl",
-        ]
-
-    def prepare(self):
-        """
-        prepare test environment for vllm. This includes clone vllm repo, install all wheels, test dependencies and set env
-        """
-        params = VllmTestParameters()
-        logger.info("Display VllmTestParameters %s", params)
-        self._set_envs(params)
-
-        clone_vllm(dst=self.work_directory)
-        self.cp_torch_cleaning_script(params)
-        with working_directory(self.work_directory):
-            remove_dir(Path("vllm"))
-            self._install_wheels(params)
-            self._install_dependencies()
-        # verify the torches are not overridden by test dependencies
-
-        check_versions()
-
-    def run(self):
-        """
-        main function to run vllm test
-        """
-        self.prepare()
-        try:
-            with working_directory(self.work_directory):
-                if self.test_type == TestInpuType.TEST_PLAN:
-                    if self.num_shards > 1:
-                        run_test_plan(
-                            self.test_plan,
-                            "vllm",
-                            sample_vllm_test_library(),
-                            self.shard_id,
-                            self.num_shards,
-                        )
-                    else:
-                        run_test_plan(
-                            self.test_plan, "vllm", sample_vllm_test_library()
-                        )
-                else:
-                    raise ValueError(f"Unknown test type {self.test_type}")
-        finally:
-            # double check the torches are not overridden by other packages
-            check_versions()
-
-    def cp_torch_cleaning_script(self, params: VllmTestParameters):
-        script = get_path(params.cleaning_script, resolve=True)
-        vllm_script = Path(f"./{self.work_directory}/use_existing_torch.py")
-        copy(script, vllm_script)
-
-    def _install_wheels(self, params: VllmTestParameters):
-        logger.info("Running vllm test with inputs: %s", params)
-        if not pkg_exists("torch"):
-            # install torch from local whls if it's not installed yet.
-            torch_p = f"{str(params.torch_whls_path)}/{self.TORCH_WHL_PATH_REGEX}"
-            pip_install_first_match(torch_p, self.TORCH_WHL_EXTRA)
-
-        torch_whls_path = [
-            f"{str(params.torch_whls_path)}/{whl_path}"
-            for whl_path in self.TORCH_ADDITIONAL_WHLS_REGEX
-        ]
-        for torch_whl in torch_whls_path:
-            pip_install_first_match(torch_whl)
-        logger.info("Done. Installed torch and other torch-related wheels ")
-
-        logger.info("Installing vllm wheels")
-        vllm_whls_path = [
-            f"{str(params.vllm_whls_path)}/{whl_path}"
-            for whl_path in self.VLLM_TEST_WHLS_REGEX
-        ]
-        for vllm_whl in vllm_whls_path:
-            pip_install_first_match(vllm_whl)
-        logger.info("Done. Installed vllm wheels")
-
-    def _install_test_dependencies(self):
-        """
-        This method replaces torch dependencies with local torch wheel info in
-        requirements/test.in file from vllm repo. then generates the test.txt
-        in runtime
-        """
-        logger.info("generate test.txt from requirements/test.in with local torch whls")
-        preprocess_test_in()
-        copy("requirements/test.txt", "snapshot_constraint.txt")
-
-        run_command(
-            f"{sys.executable} -m uv pip compile requirements/test.in "
-            "-o test.txt "
-            "--index-strategy unsafe-best-match "
-            "--constraint snapshot_constraint.txt "
-            "--torch-backend cu128"
-        )
-        pip_install_packages(requirements="test.txt", prefer_uv=True)
-        logger.info("Done. installed requirements for test dependencies")
-
-    def _install_dependencies(self):
-        pip_install_packages(packages=["-e", "tests/vllm_test_utils"], prefer_uv=True)
-        pip_install_packages(packages=["hf_transfer"], prefer_uv=True)
-        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-
-        # using script from vllm repo to remove all torch packages from requirements txt
-        run_python("use_existing_torch.py")
-
-        # install common packages
-        for requirements in ["requirements/common.txt", "requirements/build.txt"]:
-            pip_install_packages(
-                requirements=requirements,
-                prefer_uv=True,
-            )
-        # install test packages
-        self._install_test_dependencies()
-
-    def _set_envs(self, inputs: VllmTestParameters):
-        os.environ["TORCH_CUDA_ARCH_LIST"] = inputs.torch_cuda_arch_list
-        if not validate_cuda(get_env("TORCH_CUDA_ARCH_LIST")):
-            logger.warning(
-                "Missing supported TORCH_CUDA_ARCH_LIST. "
-                "Currently support TORCH_CUDA_ARCH_LIST env var "
-                "with supported arch [8.0, 8.9, 9.0]"
-            )
-
-        os.environ["HF_TOKEN"] = os.getenv("VLLM_TEST_HUGGING_FACE_TOKEN", "")
-        if not get_env("HF_TOKEN"):
-            raise ValueError(
-                "missing required HF_TOKEN, please set VLLM_TEST_HUGGING_FACE_TOKEN env var"
-            )
-        if not get_env("TORCH_CUDA_ARCH_LIST"):
-            raise ValueError(
-                "missing required TORCH_CUDA_ARCH_LIST, please set TORCH_CUDA_ARCH_LIST env var"
-            )
-
-
-def preprocess_test_in(
-    target_file: str = "requirements/test.in", additional_packages: Iterable[str] = ()
-):
-    """
-    This modifies the target_file file in place in vllm work directory.
-    It removes torch and unwanted packages in target_file and replace with local torch whls
-    package  with format "$WHEEL_PACKAGE_NAME @ file://<LOCAL_PATH>"
-    """
-    additional_package_to_move = list(additional_packages or ())
-    pkgs_to_remove = [
-        "torch",
-        "torchvision",
-        "torchaudio",
-        "xformers",
-        "mamba_ssm",
-    ] + additional_package_to_move
-    # Read current requirements
-    target_path = Path(target_file)
-    lines = target_path.read_text().splitlines()
-
-    pkgs_to_add = []
-
-    # Remove lines starting with the package names (==, @, >=) — case-insensitive
-    pattern = re.compile(rf"^({'|'.join(pkgs_to_remove)})\s*(==|@|>=)", re.IGNORECASE)
-    kept_lines = [line for line in lines if not pattern.match(line)]
-
-    # Get local installed torch/vision/audio from pip freeze
-    # This is hacky, but it works
-    pip_freeze = subprocess.check_output(["pip", "freeze"], text=True)
-    header_lines = [
-        line
-        for line in pip_freeze.splitlines()
-        if re.match(
-            r"^(torch|torchvision|torchaudio)\s*@\s*file://", line, re.IGNORECASE
-        )
-    ]
-
-    # Write back: header_lines + blank + kept_lines
-    out_lines = header_lines + [""] + kept_lines
-    if pkgs_to_add:
-        out_lines += [""] + pkgs_to_add
-
-    out = "\n".join(out_lines) + "\n"
-    target_path.write_text(out)
-    logger.info("[INFO] Updated %s", target_file)
-
-
-def validate_cuda(value: str) -> bool:
-    VALID_VALUES = {"8.0", "8.9", "9.0"}
-    return all(v in VALID_VALUES for v in value.split())
-
-
-def check_versions():
-    """
-    check installed packages version
-    """
-    logger.info("Double check installed packages")
-    patterns = ["torch", "xformers", "torchvision", "torchaudio", "vllm"]
-    for pkg in patterns:
-        pkg_exists(pkg)
-    logger.info("Done. checked installed packages")

.ci/lumen_cli/cli/run.py

@@ -1,40 +0,0 @@
-# main.py
-
-import argparse
-import logging
-
-from cli.build_cli.register_build import register_build_commands
-from cli.lib.common.logger import setup_logging
-from cli.test_cli.register_test import register_test_commands
-
-
-logger = logging.getLogger(__name__)
-
-
-def main():
-    # Define top-level parser
-    parser = argparse.ArgumentParser(description="Lumos CLI")
-    subparsers = parser.add_subparsers(dest="command", required=True)
-    parser.add_argument(
-        "--log-level", default="INFO", help="Log level (DEBUG, INFO, WARNING, ERROR)"
-    )
-
-    # registers second-level subcommands
-    register_build_commands(subparsers)
-    register_test_commands(subparsers)
-
-    # parse args after all options are registered
-    args = parser.parse_args()
-
-    # setup global logging
-    setup_logging(getattr(logging, args.log_level.upper(), logging.INFO))
-    logger.debug("Parsed args: %s", args)
-
-    if hasattr(args, "func"):
-        args.func(args)
-    else:
-        parser.print_help()
-
-
-if __name__ == "__main__":
-    main()

.ci/lumen_cli/cli/test_cli/register_test.py

@@ -1,62 +0,0 @@
-import argparse
-import logging
-
-from cli.lib.common.cli_helper import register_targets, RichHelp, TargetSpec
-from cli.lib.core.vllm.vllm_test import VllmTestRunner
-
-
-logger = logging.getLogger(__name__)
-
-# Maps targets to their argparse configuration and runner
-# it adds new target to path python -m cli.run build external {target} with buildrunner
-_TARGETS: dict[str, TargetSpec] = {
-    "vllm": {
-        "runner": VllmTestRunner,
-        "help": "test vLLM with pytorch main",
-    }
-    # add yours ...
-}
-
-
-def common_args(parser: argparse.ArgumentParser) -> None:
-    """
-    Add common CLI arguments to the given parser.
-    """
-    parser.add_argument(
-        "--shard-id",
-        type=int,
-        default=1,
-        help="a shard id to run, e.g. '0,1,2,3'",
-    )
-    parser.add_argument(
-        "--num-shards",
-        type=int,
-        default=1,
-        help="a number of shards to run, e.g. '4'",
-    )
-    group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument(
-        "-tp",
-        "--test-plan",
-        type=str,
-        help="a pre-defined test plan to run, e.g. 'basic_correctness_test'",
-    )
-
-
-def register_test_commands(subparsers: argparse._SubParsersAction) -> None:
-    build_parser = subparsers.add_parser(
-        "test",
-        help="test related commands",
-        formatter_class=RichHelp,
-    )
-    build_subparsers = build_parser.add_subparsers(dest="test_command", required=True)
-    overview = "\n".join(
-        f"  {name:12} {spec.get('help', '')}" for name, spec in _TARGETS.items()
-    )
-    external_parser = build_subparsers.add_parser(
-        "external",
-        help="Test external targets",
-        description="Test third-party targets.\n\nAvailable targets:\n" + overview,
-        formatter_class=RichHelp,
-    )
-    register_targets(external_parser, _TARGETS, common_args=common_args)

.ci/lumen_cli/pyproject.toml

@@ -1,23 +0,0 @@
-[project]
-name = "lumen-ci"
-version = "0.1.0"
-dependencies = [
-    "pyyaml==6.0.2",
-    "GitPython==3.1.45",
-    "docker==7.1.0",
-    "pytest==7.3.2",
-    "uv==0.8.6"
-]
-
-[tool.setuptools]
-packages = ["cli"]
-
-[tool.setuptools.package-dir]
-cli = "cli"
-
-[tool.ruff.lint]
-# Enable preview mode for linting
-preview = true
-
-# Now you can select your preview rules, like RUF048
-extend-select = ["RUF048"]

.ci/lumen_cli/tests/test_app.py

@@ -1,47 +0,0 @@
-# tests/test_cli.py
-import io
-import sys
-import unittest
-from contextlib import redirect_stderr, redirect_stdout
-from unittest.mock import patch
-
-from cli.run import main
-
-
-class TestArgparseCLI(unittest.TestCase):
-    @patch("cli.build_cli.register_build.VllmBuildRunner.run", return_value=None)
-    @patch("cli.build_cli.register_build.VllmBuildRunner.__init__", return_value=None)
-    def test_cli_run_build_external(self, mock_init, mock_run):
-        from cli.run import main  # import after patches if needed
-
-        test_args = ["cli.run", "build", "external", "vllm"]
-        with patch.object(sys, "argv", test_args):
-            # argparse may call sys.exit on error; capture to avoid test aborts
-            try:
-                main()
-            except SystemExit:
-                pass
-        mock_init.assert_called_once()  # got constructed
-        mock_run.assert_called_once_with()  # run() called
-
-    def test_build_help(self):
-        test_args = ["cli.run", "build", "--help"]
-
-        with patch.object(sys, "argv", test_args):
-            stdout = io.StringIO()
-            stderr = io.StringIO()
-
-            # --help always raises SystemExit(0)
-            with self.assertRaises(SystemExit) as cm:
-                with redirect_stdout(stdout), redirect_stderr(stderr):
-                    main()
-
-            self.assertEqual(cm.exception.code, 0)
-
-            output = stdout.getvalue()
-            self.assertIn("usage", output)
-            self.assertIn("external", output)
-
-
-if __name__ == "__main__":
-    unittest.main()

.ci/lumen_cli/tests/test_cli_helper.py

@@ -1,115 +0,0 @@
-import argparse
-import io
-import unittest
-from contextlib import redirect_stderr
-from unittest.mock import patch
-
-from cli.lib.common.cli_helper import BaseRunner, register_targets, RichHelp, TargetSpec
-
-
-# ---- Dummy runners for unittests----
-class FooRunner(BaseRunner):
-    """Foo description from docstring."""
-
-    def run(self) -> None:  # replaced by mock
-        pass
-
-
-class BarRunner(BaseRunner):
-    def run(self) -> None:  # replaced by mock
-        pass
-
-
-def add_foo_args(p: argparse.ArgumentParser) -> None:
-    p.add_argument("--x", type=int, required=True, help="x value")
-
-
-def common_args(p: argparse.ArgumentParser) -> None:
-    p.add_argument("--verbose", action="store_true", help="verbose flag")
-
-
-def build_parser(specs: dict[str, TargetSpec]) -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(prog="app", formatter_class=RichHelp)
-    register_targets(
-        parser=parser,
-        target_specs=specs,
-        common_args=common_args,
-    )
-    return parser
-
-
-def get_subparser(
-    parser: argparse.ArgumentParser, name: str
-) -> argparse.ArgumentParser:
-    subparsers_action = next(
-        a
-        for a in parser._subparsers._group_actions  # type: ignore[attr-defined]
-        if isinstance(a, argparse._SubParsersAction)
-    )
-    return subparsers_action.choices[name]
-
-
-class TestRegisterTargets(unittest.TestCase):
-    def test_metavar_lists_targets(self):
-        specs: dict[str, TargetSpec] = {
-            "foo": {"runner": FooRunner, "add_arguments": add_foo_args},
-            "bar": {"runner": BarRunner},
-        }
-        parser = build_parser(specs)
-        subparsers_action = next(
-            a
-            for a in parser._subparsers._group_actions  # type: ignore[attr-defined]
-            if isinstance(a, argparse._SubParsersAction)
-        )
-        self.assertEqual(subparsers_action.metavar, "{foo,bar}")
-
-    def test_add_arguments_and_common_args_present(self):
-        specs: dict[str, TargetSpec] = {
-            "foo": {"runner": FooRunner, "add_arguments": add_foo_args},
-        }
-        parser = build_parser(specs)
-        foo = get_subparser(parser, "foo")
-        help_text = foo.format_help()
-        self.assertIn("--x", help_text)
-        self.assertIn("--verbose", help_text)
-
-    def test_runner_constructed_with_ns_and_run_called(self):
-        specs: dict[str, TargetSpec] = {
-            "foo": {"runner": FooRunner, "add_arguments": add_foo_args},
-        }
-        parser = build_parser(specs)
-
-        with (
-            patch.object(FooRunner, "__init__", return_value=None) as mock_init,
-            patch.object(FooRunner, "run", return_value=None) as mock_run,
-        ):
-            ns = parser.parse_args(["foo", "--x", "3", "--verbose"])
-            ns.func(ns)  # set by register_targets
-            # __init__ received the Namespace
-            self.assertEqual(mock_init.call_count, 1)
-            (called_ns,), _ = mock_init.call_args
-            self.assertIsInstance(called_ns, argparse.Namespace)
-            # run() called with no args
-            mock_run.assert_called_once_with()
-
-    def test_runner_docstring_used_as_description_when_missing(self):
-        specs: dict[str, TargetSpec] = {
-            "foo": {"runner": FooRunner, "add_arguments": add_foo_args},
-        }
-        parser = build_parser(specs)
-        foo = get_subparser(parser, "foo")
-        help_text = foo.format_help()
-        self.assertIn("Foo description from docstring.", help_text)
-
-    def test_missing_target_raises_systemexit_with_usage(self):
-        specs: dict[str, TargetSpec] = {"foo": {"runner": FooRunner}}
-        parser = build_parser(specs)
-        buf = io.StringIO()
-        with self.assertRaises(SystemExit), redirect_stderr(buf):
-            parser.parse_args([])
-        err = buf.getvalue()
-        self.assertIn("usage:", err)
-
-
-if __name__ == "__main__":
-    unittest.main()

.ci/lumen_cli/tests/test_docker_helper.py

@@ -1,75 +0,0 @@
-import unittest
-from unittest import mock
-from unittest.mock import MagicMock
-
-import docker.errors as derr
-from cli.lib.common.docker_helper import _get_client, local_image_exists
-
-
-class TestDockerImageHelpers(unittest.TestCase):
-    def setUp(self):
-        # Reset the singleton in the target module
-        patcher = mock.patch("cli.lib.common.docker_helper._docker_client", None)
-        self.addCleanup(patcher.stop)
-        patcher.start()
-
-    def test_local_image_exists_true(self):
-        # Mock a docker client whose images.get returns an object (no exception)
-        mock_client = MagicMock()
-        mock_client.images.get.return_value = object()
-        ok = local_image_exists("repo:tag", client=mock_client)
-        self.assertTrue(ok)
-
-    def test_local_image_exists_not_found_false(self):
-        mock_client = MagicMock()
-        # Raise docker.errors.NotFound
-        mock_client.images.get.side_effect = derr.NotFound("nope")
-        ok = local_image_exists("missing:latest", client=mock_client)
-        self.assertFalse(ok)
-
-    def test_local_image_exists_api_error_false(self):
-        mock_client = MagicMock()
-        mock_client.images.get.side_effect = derr.APIError("boom", None)
-
-        ok = local_image_exists("broken:tag", client=mock_client)
-        self.assertFalse(ok)
-
-    def test_local_image_exists_uses_lazy_singleton(self):
-        # Patch docker.from_env used by _get_client()
-        with mock.patch(
-            "cli.lib.common.docker_helper.docker.from_env"
-        ) as mock_from_env:
-            mock_docker_client = MagicMock()
-            mock_from_env.return_value = mock_docker_client
-
-            # First call should create and cache the client
-            c1 = _get_client()
-            self.assertIs(c1, mock_docker_client)
-            mock_from_env.assert_called_once()
-
-            # Second call should reuse cached client (no extra from_env calls)
-            c2 = _get_client()
-            self.assertIs(c2, mock_docker_client)
-            mock_from_env.assert_called_once()  # still once
-
-    def test_local_image_exists_without_client_param_calls_get_client_once(self):
-        # Ensure _get_client is called and cached; local_image_exists should reuse it
-        with mock.patch("cli.lib.common.docker_helper._get_client") as mock_get_client:
-            mock_client = MagicMock()
-            mock_get_client.return_value = mock_client
-
-            # 1st call
-            local_image_exists("repo:tag")
-            # 2nd call
-            local_image_exists("repo:tag2")
-
-            # local_image_exists should call _get_client each time,
-            # but your _get_client itself caches docker.from_env.
-            self.assertEqual(mock_get_client.call_count, 2)
-            self.assertEqual(mock_client.images.get.call_count, 2)
-            mock_client.images.get.assert_any_call("repo:tag")
-            mock_client.images.get.assert_any_call("repo:tag2")
-
-
-if __name__ == "__main__":
-    unittest.main()

.ci/lumen_cli/tests/test_envs_helper.py

@@ -1,149 +0,0 @@
-import os
-import unittest
-from dataclasses import dataclass
-from pathlib import Path
-from unittest.mock import patch
-
-import cli.lib.common.envs_helper as m
-
-
-class TestEnvHelpers(unittest.TestCase):
-    def setUp(self):
-        # Keep a copy of the original environment to restore later
-        self._env_backup = dict(os.environ)
-
-    def tearDown(self):
-        # Restore environment to original state
-        os.environ.clear()
-        os.environ.update(self._env_backup)
-
-    # -------- get_env --------
-    def test_get_env_unset_returns_default(self):
-        with patch.dict(os.environ, {}, clear=True):
-            self.assertEqual(m.get_env("FOO", "default"), "default")
-
-    def test_get_env_empty_returns_default(self):
-        with patch.dict(os.environ, {"FOO": ""}, clear=True):
-            self.assertEqual(m.get_env("FOO", "default"), "default")
-
-    def test_get_env_set_returns_value(self):
-        with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
-            self.assertEqual(m.get_env("FOO", "default"), "bar")
-
-    def test_get_env_not_exist_returns_default(self):
-        with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
-            self.assertEqual(m.get_env("TEST_NOT_EXIST", "default"), "default")
-
-    def test_get_env_not_exist_without_default(self):
-        with patch.dict(os.environ, {"FOO": "bar"}, clear=True):
-            self.assertEqual(m.get_env("TEST_NOT_EXIST"), "")
-
-    # -------- env_bool --------
-    def test_env_bool_uses_default_when_unset(self):
-        with patch.dict(os.environ, {}, clear=True):
-            self.assertTrue(m.env_bool("FLAG", default=True))
-            self.assertFalse(m.env_bool("FLAG", default=False))
-
-    def test_env_bool_uses_str2bool_when_set(self):
-        # Patch str2bool used by env_bool so we don't depend on its exact behavior
-        def fake_str2bool(s: str) -> bool:
-            return s.lower() in {"1", "true", "yes", "on", "y"}
-
-        with (
-            patch.dict(os.environ, {"FLAG": "yEs"}, clear=True),
-            patch.object(m, "str2bool", fake_str2bool),
-        ):
-            self.assertTrue(m.env_bool("FLAG", default=False))
-
-    # -------- env_path_optional / env_path --------
-    def test_env_path_optional_unset_returns_none_by_default(self):
-        with patch.dict(os.environ, {}, clear=True):
-            self.assertIsNone(m.env_path_optional("P"))
-
-    def test_env_path_optional_unset_returns_none_when_env_var_is_empty(self):
-        with patch.dict(os.environ, {"P": ""}, clear=True):
-            self.assertIsNone(m.env_path_optional("P"))
-
-    def test_env_path_optional_unset_returns_default_str(self):
-        # default as string; resolve=True by default -> absolute path
-        default_str = "x/y"
-        with patch.dict(os.environ, {}, clear=True):
-            p = m.env_path_optional("P", default=default_str)
-            self.assertIsInstance(p, Path)
-            self.assertIsNotNone(p)
-            if p:
-                self.assertTrue(p.is_absolute())
-                self.assertEqual(p.parts[-2:], ("x", "y"))
-
-    def test_env_path_optional_unset_returns_default_path_no_resolve(self):
-        d = Path("z")
-        with patch.dict(os.environ, {}, clear=True):
-            p = m.env_path_optional("P", default=d, resolve=False)
-            self.assertEqual(p, d)
-
-    def test_env_path_optional_respects_resolve_true(self):
-        with patch.dict(os.environ, {"P": "a/b"}, clear=True):
-            p = m.env_path_optional("P", resolve=True)
-            self.assertIsInstance(p, Path)
-            if p:
-                self.assertTrue(p.is_absolute())
-
-    def test_env_path_optional_respects_resolve_false(self):
-        with patch.dict(os.environ, {"P": "rel/dir"}, clear=True):
-            p = m.env_path_optional("P", resolve=False)
-            self.assertEqual(p, Path("rel/dir"))
-            if p:
-                self.assertFalse(p.is_absolute())
-
-    def test_env_path_raises_when_missing_and_default_none(self):
-        with patch.dict(os.environ, {}, clear=True):
-            with self.assertRaises(ValueError):
-                m.env_path("P", None, resolve=True)
-
-    def test_env_path_returns_path_when_present(self):
-        tmp = Path("./b").resolve()
-        with patch.dict(os.environ, {"P": str(tmp)}, clear=True):
-            p = m.env_path("P", None, resolve=True)
-            self.assertEqual(p, tmp)
-
-    # -------- dataclass field helpers --------
-    def test_dataclass_fields_read_env_at_instantiation(self):
-        @dataclass
-        class Cfg:
-            flag: bool = m.env_bool_field("FLAG", default=False)
-            out: Path = m.env_path_field("OUT", default="ab", resolve=True)
-            name: str = m.env_str_field("NAME", default="anon")
-
-        # First instantiation
-        with patch.dict(
-            os.environ, {"FLAG": "true", "OUT": "outdir", "NAME": "alice"}, clear=True
-        ):
-            cfg1 = Cfg()
-            self.assertTrue(cfg1.flag)
-            self.assertIsInstance(cfg1.out, Path)
-            self.assertTrue(cfg1.out.is_absolute())
-            self.assertEqual(cfg1.name, "alice")
-            cfg1.name = "bob"  # change instance value
-            self.assertEqual(cfg1.name, "bob")  # change is reflected
-
-        # Change env; new instance should reflect new values
-        with patch.dict(os.environ, {"FLAG": "false", "NAME": ""}, clear=True):
-            cfg2 = Cfg()
-            self.assertFalse(cfg2.flag)  # str2bool("false") -> False
-            self.assertTrue("ab" in str(cfg2.out))
-            self.assertIsInstance(cfg2.out, Path)
-            self.assertTrue(cfg2.out.is_absolute())
-            self.assertEqual(cfg2.name, "anon")  # empty -> fallback to default
-
-    def test_dataclass_path_field_with_default_value(self):
-        @dataclass
-        class C2:
-            out: Path = m.env_path_field("OUT", default="some/dir", resolve=False)
-
-        with patch.dict(os.environ, {}, clear=True):
-            c = C2()
-            self.assertEqual(c.out, Path("some/dir"))
-
-
-if __name__ == "__main__":
-    unittest.main()

.ci/lumen_cli/tests/test_path_helper.py

@@ -1,122 +0,0 @@
-# test_path_utils.py
-# Run: pytest -q
-
-import os
-import unittest
-from pathlib import Path
-from tempfile import TemporaryDirectory
-
-from cli.lib.common.path_helper import (
-    copy,
-    ensure_dir_exists,
-    force_create_dir,
-    get_path,
-    is_path_exist,
-    remove_dir,
-)
-
-
-class TestPathHelper(unittest.TestCase):
-    def setUp(self):
-        self.tmpdir = TemporaryDirectory()
-        self.tmp_path = Path(self.tmpdir.name)
-
-    def tearDown(self):
-        self.tmpdir.cleanup()
-
-    # -------- get_path --------
-    def test_get_path_returns_path_for_str(self):
-        # Use relative path to avoid absolute-ness
-        rel_str = "sub/f.txt"
-        os.chdir(self.tmp_path)
-        p = get_path(rel_str, resolve=False)
-        self.assertIsInstance(p, Path)
-        self.assertFalse(p.is_absolute())
-        self.assertEqual(str(p), rel_str)
-
-    def test_get_path_resolves(self):
-        rel_str = "sub/f.txt"
-        p = get_path(str(self.tmp_path / rel_str), resolve=True)
-        self.assertTrue(p.is_absolute())
-        self.assertTrue(str(p).endswith(rel_str))
-
-    def test_get_path_with_path_input(self):
-        p_in = self.tmp_path / "sub/f.txt"
-        p_out = get_path(p_in, resolve=False)
-        self.assertTrue(str(p_out) == str(p_in))
-
-    def test_get_path_with_none_raises(self):
-        with self.assertRaises(ValueError):
-            get_path(None)  # type: ignore[arg-type]
-
-    def test_get_path_invalid_type_raises(self):
-        with self.assertRaises(TypeError):
-            get_path(123)  # type: ignore[arg-type]
-
-    # -------- ensure_dir_exists / force_create_dir / remove_dir --------
-    def test_ensure_dir_exists_creates_and_is_idempotent(self):
-        d = self.tmp_path / "made"
-        ensure_dir_exists(d)
-        self.assertTrue(d.exists() and d.is_dir())
-        ensure_dir_exists(d)
-
-    def test_force_create_dir_clears_existing(self):
-        d = self.tmp_path / "fresh"
-        (d / "inner").mkdir(parents=True)
-        (d / "inner" / "f.txt").write_text("x")
-        force_create_dir(d)
-        self.assertTrue(d.exists())
-        self.assertEqual(list(d.iterdir()), [])
-
-    def test_remove_dir_none_is_noop(self):
-        remove_dir(None)  # type: ignore[arg-type]
-
-    def test_remove_dir_nonexistent_is_noop(self):
-        ghost = self.tmp_path / "ghost"
-        remove_dir(ghost)
-
-    def test_remove_dir_accepts_str(self):
-        d = self.tmp_path / "to_rm"
-        d.mkdir()
-        remove_dir(str(d))
-        self.assertFalse(d.exists())
-
-    # -------- copy --------
-    def test_copy_file_to_file(self):
-        src = self.tmp_path / "src.txt"
-        dst = self.tmp_path / "out" / "dst.txt"
-        src.write_text("hello")
-        copy(src, dst)
-        self.assertEqual(dst.read_text(), "hello")
-
-    def test_copy_dir_to_new_dir(self):
-        src = self.tmp_path / "srcdir"
-        (src / "a").mkdir(parents=True)
-        (src / "a" / "f.txt").write_text("content")
-        dst = self.tmp_path / "destdir"
-        copy(src, dst)
-        self.assertEqual((dst / "a" / "f.txt").read_text(), "content")
-
-    def test_copy_dir_into_existing_dir_overwrite_true_merges(self):
-        src = self.tmp_path / "srcdir"
-        dst = self.tmp_path / "destdir"
-        (src / "x").mkdir(parents=True)
-        (src / "x" / "new.txt").write_text("new")
-        dst.mkdir()
-        (dst / "existing.txt").write_text("old")
-        copy(src, dst)
-        self.assertEqual((dst / "existing.txt").read_text(), "old")
-        self.assertEqual((dst / "x" / "new.txt").read_text(), "new")
-
-    def test_is_str_path_exist(self):
-        p = self.tmp_path / "x.txt"
-        p.write_text("1")
-        self.assertTrue(is_path_exist(str(p)))
-        self.assertTrue(is_path_exist(p))
-        self.assertFalse(is_path_exist(str(self.tmp_path / "missing")))
-        self.assertFalse(is_path_exist(self.tmp_path / "missing"))
-        self.assertFalse(is_path_exist(""))
-
-
-if __name__ == "__main__":
-    unittest.main()

.ci/lumen_cli/tests/test_run_plan.py

@@ -1,185 +0,0 @@
-# tests/test_run_test_plan.py
-import importlib
-from contextlib import nullcontext
-from types import SimpleNamespace
-from unittest.mock import MagicMock
-
-import pytest
-
-
-MOD = "cli.lib.core.vllm.lib"
-
-# We import inside tests so the MOD override above applies everywhere
-run_test_plan_import_path = f"{MOD}.run_test_plan"
-
-
-def _get_cmd(c):
-    # Support both kwargs and positional args
-    return c.kwargs.get("cmd", c.args[0] if c.args else None)
-
-
-def _get_check(c):
-    if "check" in c.kwargs:
-        return c.kwargs["check"]
-    # If positional, assume second arg is 'check' when present; default False
-    return c.args[1] if len(c.args) > 1 else False
-
-
-@pytest.fixture
-def patch_module(monkeypatch):
-    """
-    Patch helpers ('pip_install_packages', 'temp_environ', 'working_directory',
-    'run_command', 'logger') inside the target module and expose them.
-    """
-    module = importlib.import_module(MOD)
-
-    # Create fakes/mocks
-    pip_install_packages = MagicMock(name="pip_install_packages")
-    run_command = MagicMock(name="run_command", return_value=0)
-
-    # temp_environ / working_directory: record calls but act as context managers
-    temp_calls: list[dict] = []
-    workdir_calls: list[str] = []
-
-    def fake_working_directory(path: str):
-        workdir_calls.append(path)
-        return nullcontext()
-
-    def fake_temp_env(map: dict[str, str]):
-        temp_calls.append(map)
-        return nullcontext()
-
-    logger = SimpleNamespace(
-        info=MagicMock(name="logger.info"),
-        error=MagicMock(name="logger.error"),
-    )
-
-    # Apply patches (raise if attribute doesn't exist)
-    monkeypatch.setattr(
-        module, "pip_install_packages", pip_install_packages, raising=True
-    )
-    monkeypatch.setattr(module, "run_command", run_command, raising=True)
-    monkeypatch.setattr(
-        module, "working_directory", fake_working_directory, raising=True
-    )
-    monkeypatch.setattr(module, "temp_environ", fake_temp_env, raising=True)
-    monkeypatch.setattr(module, "logger", logger, raising=True)
-
-    return SimpleNamespace(
-        module=module,
-        run_test_plan=module.run_test_plan,  # expose to avoid getattr("constant") (Ruff B009)
-        pip_install_packages=pip_install_packages,
-        run_command=run_command,
-        temp_calls=temp_calls,
-        workdir_calls=workdir_calls,
-        logger=logger,
-    )
-
-
-def test_success_runs_all_steps_and_uses_env_and_workdir(monkeypatch, patch_module):
-    run_test_plan = patch_module.run_test_plan
-
-    tests_map = {
-        "basic": {
-            "title": "Basic suite",
-            "package_install": [],
-            "working_directory": "tests",
-            "env_vars": {"GLOBAL_FLAG": "1"},
-            "steps": [
-                "export A=x && pytest -q",
-                "export B=y && pytest -q tests/unit",
-            ],
-        }
-    }
-
-    # One exit code per step (export + two pytest)
-    patch_module.run_command.side_effect = [0, 0, 0]
-
-    run_test_plan("basic", "cpu", tests_map)
-
-    calls = patch_module.run_command.call_args_list
-    cmds = [_get_cmd(c) for c in calls]
-    checks = [_get_check(c) for c in calls]
-
-    assert cmds == [
-        "export A=x && pytest -q",
-        "export B=y && pytest -q tests/unit",
-    ]
-    assert all(chk is False for chk in checks)
-
-    assert patch_module.workdir_calls == ["tests"]
-    assert patch_module.temp_calls == [{"GLOBAL_FLAG": "1"}]
-
-
-def test_installs_packages_when_present(monkeypatch, patch_module):
-    run_test_plan = patch_module.module.run_test_plan
-
-    tests_map = {
-        "with_pkgs": {
-            "title": "Needs deps",
-            "package_install": ["timm==1.0.0", "flash-attn"],
-            "steps": ["pytest -q"],
-        }
-    }
-
-    patch_module.run_command.return_value = 0
-
-    run_test_plan("with_pkgs", "gpu", tests_map)
-
-    patch_module.pip_install_packages.assert_called_once_with(
-        packages=["timm==1.0.0", "flash-attn"],
-        prefer_uv=True,
-    )
-
-
-def test_raises_on_missing_plan(patch_module):
-    run_test_plan = patch_module.module.run_test_plan
-    with pytest.raises(RuntimeError) as ei:
-        run_test_plan("nope", "cpu", tests_map={})
-
-    assert "test nope not found" in str(ei.value)
-
-
-def test_aggregates_failures_and_raises(monkeypatch, patch_module):
-    run_test_plan = patch_module.module.run_test_plan
-
-    tests_map = {
-        "mix": {
-            "title": "Some pass some fail",
-            "steps": [
-                "pytest test_a.py",  # 0 → pass
-                "pytest test_b.py",  # 1 → fail
-                "pytest test_c.py",  # 2 → fail
-            ],
-        }
-    }
-
-    # Simulate pass, fail, fail
-    patch_module.run_command.side_effect = [0, 1, 2]
-
-    with pytest.raises(RuntimeError) as ei:
-        run_test_plan("mix", "cpu", tests_map)
-
-    msg = str(ei.value)
-    assert "2 pytest runs failed" in msg
-    # Ensure logger captured failed tests list
-    patch_module.logger.error.assert_called_once()
-    # And we attempted all three commands
-    assert patch_module.run_command.call_count == 3
-
-
-def test_custom_working_directory_used(patch_module):
-    run_test_plan = patch_module.module.run_test_plan
-
-    tests_map = {
-        "customwd": {
-            "title": "Custom wd",
-            "working_directory": "examples/ci",
-            "steps": ["pytest -q"],
-        }
-    }
-
-    patch_module.run_command.return_value = 0
-    run_test_plan("customwd", "cpu", tests_map)
-
-    assert patch_module.workdir_calls == ["examples/ci"]

.ci/lumen_cli/tests/test_utils.py

@@ -1,143 +0,0 @@
-import os
-import tempfile
-import unittest
-from pathlib import Path
-
-from cli.lib.common.utils import temp_environ, working_directory  # <-- replace import
-
-
-class EnvIsolatedTestCase(unittest.TestCase):
-    """Base class that snapshots os.environ and CWD for isolation."""
-
-    def setUp(self):
-        import os
-        import tempfile
-
-        self._env_backup = dict(os.environ)
-
-        # Snapshot/repair CWD if it's gone
-        try:
-            self._cwd_backup = os.getcwd()
-        except FileNotFoundError:
-            # If CWD no longer exists, switch to a safe place and record that
-            self._cwd_backup = tempfile.gettempdir()
-            os.chdir(self._cwd_backup)
-
-        # Create a temporary directory for the test to run in
-        self._temp_dir = tempfile.mkdtemp()
-        os.chdir(self._temp_dir)
-
-    def tearDown(self):
-        import os
-        import shutil
-        import tempfile
-
-        # Restore cwd first (before cleaning up temp dir)
-        try:
-            os.chdir(self._cwd_backup)
-        except OSError:
-            os.chdir(tempfile.gettempdir())
-
-        # Clean up temporary directory
-        try:
-            shutil.rmtree(self._temp_dir, ignore_errors=True)
-        except Exception:
-            pass  # Ignore cleanup errors
-
-        # Restore env
-        to_del = set(os.environ.keys()) - set(self._env_backup.keys())
-        for k in to_del:
-            os.environ.pop(k, None)
-        for k, v in self._env_backup.items():
-            os.environ[k] = v
-
-
-class TestTempEnviron(EnvIsolatedTestCase):
-    def test_sets_and_restores_new_var(self):
-        var = "TEST_TMP_ENV_NEW"
-        self.assertNotIn(var, os.environ)
-
-        with temp_environ({var: "123"}):
-            self.assertEqual(os.environ[var], "123")
-
-        self.assertNotIn(var, os.environ)  # removed after exit
-
-    def test_overwrites_and_restores_existing_var(self):
-        var = "TEST_TMP_ENV_OVERWRITE"
-        os.environ[var] = "orig"
-
-        with temp_environ({var: "override"}):
-            self.assertEqual(os.environ[var], "override")
-
-        self.assertEqual(os.environ[var], "orig")  # restored
-
-    def test_multiple_vars_and_missing_cleanup(self):
-        v1, v2 = "TEST_ENV_V1", "TEST_ENV_V2"
-        os.environ.pop(v1, None)
-        os.environ[v2] = "keep"
-
-        with temp_environ({v1: "a", v2: "b"}):
-            self.assertEqual(os.environ[v1], "a")
-            self.assertEqual(os.environ[v2], "b")
-
-        self.assertNotIn(v1, os.environ)  # newly-added -> removed
-        self.assertEqual(os.environ[v2], "keep")  # pre-existing -> restored
-
-    def test_restores_even_on_exception(self):
-        var = "TEST_TMP_ENV_EXCEPTION"
-        self.assertNotIn(var, os.environ)
-
-        with self.assertRaises(RuntimeError):
-            with temp_environ({var: "x"}):
-                self.assertEqual(os.environ[var], "x")
-                raise RuntimeError("boom")
-
-        self.assertNotIn(var, os.environ)  # removed after exception
-
-
-class TestWorkingDirectory(EnvIsolatedTestCase):
-    def test_changes_and_restores(self):
-        start = Path.cwd()
-        with tempfile.TemporaryDirectory() as td:
-            target = Path(td) / "wd"
-            target.mkdir()
-
-            with working_directory(str(target)):
-                self.assertEqual(Path.cwd().resolve(), target.resolve())
-
-        self.assertEqual(Path.cwd(), start)
-
-    def test_noop_when_empty_path(self):
-        start = Path.cwd()
-        with working_directory(""):
-            self.assertEqual(Path.cwd(), start)
-        self.assertEqual(Path.cwd(), start)
-
-    def test_restores_on_exception(self):
-        start = Path.cwd()
-
-        with tempfile.TemporaryDirectory() as td:
-            target = Path(td) / "wd_exc"
-            target.mkdir()
-
-            with self.assertRaises(ValueError):
-                with working_directory(str(target)):
-                    # Normalize both sides to handle /var -> /private/var
-                    self.assertEqual(Path.cwd().resolve(), target.resolve())
-                    raise ValueError("boom")
-
-        self.assertEqual(Path.cwd().resolve(), start.resolve())
-
-    def test_raises_for_missing_dir(self):
-        start = Path.cwd()
-        with tempfile.TemporaryDirectory() as td:
-            missing = Path(td) / "does_not_exist"
-            with self.assertRaises(FileNotFoundError):
-                # os.chdir should raise before yielding
-                with working_directory(str(missing)):
-                    pass
-        self.assertEqual(Path.cwd(), start)
-
-
-if __name__ == "__main__":
-    unittest.main(verbosity=2)

.ci/lumen_cli/tests/test_vllm.py

@@ -1,176 +0,0 @@
-import os
-import tempfile
-import unittest
-from pathlib import Path
-from unittest.mock import MagicMock, patch
-
-import cli.lib.core.vllm.vllm_build as vllm_build
-
-
-_VLLM_BUILD_MODULE = "cli.lib.core.vllm.vllm_build"
-
-
-class TestVllmBuildParameters(unittest.TestCase):
-    @patch(f"{_VLLM_BUILD_MODULE}.local_image_exists", return_value=True)
-    @patch(f"{_VLLM_BUILD_MODULE}.is_path_exist", return_value=True)
-    @patch(
-        "cli.lib.common.envs_helper.env_path_optional",
-        side_effect=lambda name, default=None, resolve=True: {
-            "DOCKERFILE_PATH": Path("/abs/vllm/Dockerfile"),
-            "TORCH_WHEELS_PATH": Path("/abs/dist"),
-            "OUTPUT_DIR": Path("/abs/shared"),
-        }.get(name, Path(default) if default is not None else None),
-    )
-    @patch.dict(
-        os.environ,
-        {
-            "USE_TORCH_WHEEL": "1",
-            "USE_LOCAL_BASE_IMAGE": "1",
-            "USE_LOCAL_DOCKERFILE": "1",
-            "BASE_IMAGE": "my/image:tag",
-            "DOCKERFILE_PATH": "vllm/Dockerfile",
-            "TORCH_WHEELS_PATH": "dist",
-            "OUTPUT_DIR": "shared",
-        },
-        clear=True,
-    )
-    def test_params_success_normalizes_and_validates(
-        self, mock_env_path, mock_is_path, mock_local_img
-    ):
-        params = vllm_build.VllmBuildParameters()
-        self.assertEqual(params.torch_whls_path, Path("/abs/dist"))
-        self.assertEqual(params.dockerfile_path, Path("/abs/vllm/Dockerfile"))
-        self.assertEqual(params.output_dir, Path("/abs/shared"))
-        self.assertEqual(params.base_image, "my/image:tag")
-
-    @patch(f"{_VLLM_BUILD_MODULE}.is_path_exist", return_value=False)
-    @patch.dict(
-        os.environ, {"USE_TORCH_WHEEL": "1", "TORCH_WHEELS_PATH": "dist"}, clear=True
-    )
-    def test_params_missing_torch_whls_raises(self, _is_path):
-        with tempfile.TemporaryDirectory() as td:
-            os.chdir(td)
-            with self.assertRaises(ValueError) as cm:
-                vllm_build.VllmBuildParameters(
-                    use_local_base_image=False,
-                    use_local_dockerfile=False,
-                )
-        err = cm.exception
-        self.assertIn("TORCH_WHEELS_PATH", str(err))
-
-    @patch(f"{_VLLM_BUILD_MODULE}.local_image_exists", return_value=False)
-    @patch.dict(
-        os.environ, {"USE_LOCAL_BASE_IMAGE": "1", "BASE_IMAGE": "img:tag"}, clear=True
-    )
-    def test_params_missing_local_base_image_raises(self, _local_img):
-        with tempfile.TemporaryDirectory() as td:
-            os.chdir(td)
-            with self.assertRaises(ValueError) as cm:
-                vllm_build.VllmBuildParameters(
-                    use_torch_whl=False,
-                    use_local_dockerfile=False,
-                )
-        err = cm.exception
-        self.assertIn("BASE_IMAGE", str(err))
-
-    @patch(f"{_VLLM_BUILD_MODULE}.is_path_exist", return_value=False)
-    @patch.dict(
-        os.environ,
-        {"USE_LOCAL_DOCKERFILE": "1", "DOCKERFILE_PATH": "Dockerfile"},
-        clear=True,
-    )
-    def test_params_missing_dockerfile_raises(self, _is_path):
-        with tempfile.TemporaryDirectory() as td:
-            os.chdir(td)
-            with self.assertRaises(ValueError) as cm:
-                vllm_build.VllmBuildParameters(
-                    use_torch_whl=False,
-                    use_local_base_image=False,
-                )
-        err = cm.exception
-        self.assertIn("DOCKERFILE_PATH", str(err))
-
-    @patch(f"{_VLLM_BUILD_MODULE}.is_path_exist", return_value=False)
-    @patch.dict(
-        os.environ,
-        {"OUTPUT_DIR": ""},
-        clear=True,
-    )
-    def test_params_missing_output_dir(self, _is_path):
-        with self.assertRaises(FileNotFoundError):
-            vllm_build.VllmBuildParameters()
-
-
-class TestBuildCmdAndRun(unittest.TestCase):
-    @patch(f"{_VLLM_BUILD_MODULE}.local_image_exists", return_value=True)
-    def test_generate_docker_build_cmd_includes_bits(self, _exists):
-        runner = vllm_build.VllmBuildRunner()
-        inputs = MagicMock()
-        inputs.output_dir = Path("/abs/out")
-        inputs.use_local_base_image = True
-        inputs.base_image = "img:tag"
-        inputs.torch_whls_path = Path("./vllm/tmp")
-        inputs.max_jobs = 64
-        inputs.cuda_version = "12.8.1"
-        inputs.python_version = "3.12"
-        inputs.sccache_bucket = "my-bucket"
-        inputs.sccache_region = "us-west-2"
-        inputs.torch_cuda_arch_list = "8.0;9.0"
-        inputs.target_stage = "export-wheels"
-        inputs.tag_name = "vllm-wheels"
-
-        cmd = runner._generate_docker_build_cmd(inputs)
-        squashed = " ".join(cmd.split())
-
-        self.assertIn("--output type=local,dest=/abs/out", squashed)
-        self.assertIn("-f docker/Dockerfile.nightly_torch", squashed)
-        self.assertIn("--pull=false", squashed)
-        self.assertIn("--build-arg TORCH_WHEELS_PATH=tmp", squashed)
-        self.assertIn("--build-arg BUILD_BASE_IMAGE=img:tag", squashed)
-        self.assertIn("--build-arg FINAL_BASE_IMAGE=img:tag", squashed)
-        self.assertIn("--build-arg max_jobs=64", squashed)
-        self.assertIn("--build-arg CUDA_VERSION=12.8.1", squashed)
-        self.assertIn("--build-arg PYTHON_VERSION=3.12", squashed)
-        self.assertIn("--build-arg USE_SCCACHE=1", squashed)
-        self.assertIn("--build-arg SCCACHE_BUCKET_NAME=my-bucket", squashed)
-        self.assertIn("--build-arg SCCACHE_REGION_NAME=us-west-2", squashed)
-        self.assertIn("--build-arg torch_cuda_arch_list='8.0;9.0'", squashed)
-        self.assertIn("--target export-wheels", squashed)
-        self.assertIn("-t vllm-wheels", squashed)
-
-    @patch(f"{_VLLM_BUILD_MODULE}.run_command")
-    @patch(f"{_VLLM_BUILD_MODULE}.ensure_dir_exists")
-    @patch(f"{_VLLM_BUILD_MODULE}.clone_vllm")
-    @patch.object(
-        vllm_build.VllmBuildRunner,
-        "_generate_docker_build_cmd",
-        return_value="docker buildx ...",
-    )
-    @patch.dict(
-        os.environ,
-        {
-            "USE_TORCH_WHEEL": "0",
-            "USE_LOCAL_BASE_IMAGE": "0",
-            "USE_LOCAL_DOCKERFILE": "0",
-            "OUTPUT_DIR": "shared",
-        },
-        clear=True,
-    )
-    def test_run_calls_clone_prepare_and_build(
-        self, mock_gen, mock_clone, mock_ensure, mock_run
-    ):
-        params = MagicMock()
-        params.output_dir = Path("shared")
-        params.use_local_dockerfile = False
-        params.use_torch_whl = False
-
-        with patch(f"{_VLLM_BUILD_MODULE}.VllmBuildParameters", return_value=params):
-            runner = vllm_build.VllmBuildRunner()
-            runner.run()
-
-        mock_clone.assert_called_once()
-        mock_ensure.assert_called_once_with(Path("shared"))
-        mock_gen.assert_called_once_with(params)
-        mock_run.assert_called_once()
-        _, kwargs = mock_run.call_args
-        assert kwargs.get("cwd") == "vllm"

.ci/magma-rocm/Makefile

@@ -1,11 +1,11 @@
 SHELL=/usr/bin/env bash
 
 DOCKER_CMD ?= docker
-DESIRED_ROCM ?= 7.0
+DESIRED_ROCM ?= 6.4
 DESIRED_ROCM_SHORT = $(subst .,,$(DESIRED_ROCM))
 PACKAGE_NAME = magma-rocm
 # inherit this from underlying docker image, do not pass this env var to docker
-#PYTORCH_ROCM_ARCH ?= gfx900;gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201
+#PYTORCH_ROCM_ARCH ?= gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201
 
 DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
 	-v $(shell git rev-parse --show-toplevel)/.ci:/builder \
@@ -16,7 +16,6 @@ DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
 	magma-rocm/build_magma.sh
 
 .PHONY: all
-all: magma-rocm70
 all: magma-rocm64
 all: magma-rocm63
 
@@ -25,11 +24,6 @@ clean:
 	$(RM) -r magma-*
 	$(RM) -r output
 
-.PHONY: magma-rocm70
-magma-rocm70: DESIRED_ROCM := 7.0
-magma-rocm70:
-	$(DOCKER_RUN)
-
 .PHONY: magma-rocm64
 magma-rocm64: DESIRED_ROCM := 6.4
 magma-rocm64:

.ci/magma-rocm/build_magma.sh

@@ -6,8 +6,8 @@ set -eou pipefail
 # The script expects DESIRED_CUDA and PACKAGE_NAME to be set
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
 
-# https://github.com/icl-utk-edu/magma/pull/65
-MAGMA_VERSION=d6e4117bc88e73f06d26c6c2e14f064e8fc3d1ec
+# Version 2.7.2 + ROCm related updates
+MAGMA_VERSION=a1625ff4d9bc362906bd01f805dbbe12612953f6
 
 # Folders for the build
 PACKAGE_FILES=${ROOT_DIR}/magma-rocm/package_files # metadata
@@ -20,7 +20,7 @@ mkdir -p ${PACKAGE_DIR} ${PACKAGE_OUTPUT}/linux-64 ${PACKAGE_BUILD} ${PACKAGE_RE
 
 # Fetch magma sources and verify checksum
 pushd ${PACKAGE_DIR}
-git clone https://github.com/jeffdaily/magma
+git clone https://bitbucket.org/icl/magma.git
 pushd magma
 git checkout ${MAGMA_VERSION}
 popd

.ci/magma/Makefile

@@ -16,7 +16,6 @@ DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
 	magma/build_magma.sh
 
 .PHONY: all
-all: magma-cuda130
 all: magma-cuda129
 all: magma-cuda128
 all: magma-cuda126
@@ -26,12 +25,6 @@ clean:
 	$(RM) -r magma-*
 	$(RM) -r output
 
-.PHONY: magma-cuda130
-magma-cuda130: DESIRED_CUDA := 13.0
-magma-cuda130: CUDA_ARCH_LIST := -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
-magma-cuda130:
-	$(DOCKER_RUN)
-
 .PHONY: magma-cuda129
 magma-cuda129: DESIRED_CUDA := 12.9
 magma-cuda129: CUDA_ARCH_LIST += -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120

.ci/magma/build_magma.sh

@@ -28,7 +28,6 @@ pushd ${PACKAGE_DIR}/magma-${MAGMA_VERSION}
 patch < ${PACKAGE_FILES}/CMake.patch
 patch < ${PACKAGE_FILES}/cmakelists.patch
 patch -p0 < ${PACKAGE_FILES}/thread_queue.patch
-patch -p1 < ${PACKAGE_FILES}/cuda13.patch
 patch -p1 < ${PACKAGE_FILES}/getrf_shfl.patch
 patch -p1 < ${PACKAGE_FILES}/getrf_nbparam.patch
 # The build.sh script expects to be executed from the sources root folder
@@ -38,7 +37,6 @@ popd
 # Package recipe, license and tarball
 # Folder and package name are backward compatible for the build workflow
 cp ${PACKAGE_FILES}/build.sh ${PACKAGE_RECIPE}/build.sh
-cp ${PACKAGE_FILES}/cuda13.patch ${PACKAGE_RECIPE}/cuda13.patch
 cp ${PACKAGE_FILES}/thread_queue.patch ${PACKAGE_RECIPE}/thread_queue.patch
 cp ${PACKAGE_FILES}/cmakelists.patch ${PACKAGE_RECIPE}/cmakelists.patch
 cp ${PACKAGE_FILES}/getrf_shfl.patch ${PACKAGE_RECIPE}/getrf_shfl.patch

.ci/magma/package_files/cuda13.patch

@@ -1,26 +0,0 @@
-diff --git a/interface_cuda/interface.cpp b/interface_cuda/interface.cpp
-index 73fed1b20..e77519bfe 100644
---- a/interface_cuda/interface.cpp
-+++ b/interface_cuda/interface.cpp
-@@ -438,14 +438,20 @@ magma_print_environment()
-         cudaDeviceProp prop;
-         err = cudaGetDeviceProperties( &prop, dev );
-         check_error( err );
-+        #ifdef MAGMA_HAVE_CUDA
-+#if CUDA_VERSION < 13000
-         printf( "%% device %d: %s, %.1f MHz clock, %.1f MiB memory, capability %d.%d\n",
-                 dev,
-                 prop.name,
-                 prop.clockRate / 1000.,
-+#else
-+        printf( "%% device %d: %s, ??? MHz clock, %.1f MiB memory, capability %d.%d\n",
-+                dev,
-+                prop.name,
-+#endif
-                 prop.totalGlobalMem / (1024.*1024.),
-                 prop.major,
-                 prop.minor );
--        #ifdef MAGMA_HAVE_CUDA
-         int arch = prop.major*100 + prop.minor*10;
-         if ( arch < MAGMA_CUDA_ARCH_MIN ) {
-             printf("\n"

.ci/manywheel/build.sh

@@ -5,6 +5,10 @@ set -ex
 SCRIPTPATH="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
 case "${GPU_ARCH_TYPE:-BLANK}" in
+    BLANK)
+        # Legacy behavior for CircleCI
+        bash "${SCRIPTPATH}/build_cuda.sh"
+        ;;
     cuda)
         bash "${SCRIPTPATH}/build_cuda.sh"
         ;;

.ci/manywheel/build_common.sh

@@ -97,7 +97,7 @@ if [[ -z "$PYTORCH_ROOT" ]]; then
     exit 1
 fi
 pushd "$PYTORCH_ROOT"
-retry pip install -qUr requirements-build.txt
+retry pip install -q cmake
 python setup.py clean
 retry pip install -qr requirements.txt
 case ${DESIRED_PYTHON} in
@@ -138,11 +138,28 @@ fi
 
 echo "Calling setup.py bdist at $(date)"
 
-time CMAKE_ARGS=${CMAKE_ARGS[@]} \
-    EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
+if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+    echo "Calling setup.py bdist_wheel for split build (BUILD_LIBTORCH_WHL)"
+    time EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
+    BUILD_LIBTORCH_WHL=1 BUILD_PYTHON_ONLY=0 \
     BUILD_LIBTORCH_CPU_WITH_DEBUG=$BUILD_DEBUG_INFO \
     USE_NCCL=${USE_NCCL} USE_RCCL=${USE_RCCL} USE_KINETO=${USE_KINETO} \
-    python -m build --wheel --no-isolation --outdir /tmp/$WHEELHOUSE_DIR
+    python setup.py bdist_wheel -d /tmp/$WHEELHOUSE_DIR
+    echo "Finished setup.py bdist_wheel for split build (BUILD_LIBTORCH_WHL)"
+    echo "Calling setup.py bdist_wheel for split build (BUILD_PYTHON_ONLY)"
+    time EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
+    BUILD_LIBTORCH_WHL=0 BUILD_PYTHON_ONLY=1 \
+    BUILD_LIBTORCH_CPU_WITH_DEBUG=$BUILD_DEBUG_INFO \
+    USE_NCCL=${USE_NCCL} USE_RCCL=${USE_RCCL} USE_KINETO=${USE_KINETO} \
+    CMAKE_FRESH=1 python setup.py bdist_wheel -d /tmp/$WHEELHOUSE_DIR
+    echo "Finished setup.py bdist_wheel for split build (BUILD_PYTHON_ONLY)"
+else
+    time CMAKE_ARGS=${CMAKE_ARGS[@]} \
+        EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
+        BUILD_LIBTORCH_CPU_WITH_DEBUG=$BUILD_DEBUG_INFO \
+        USE_NCCL=${USE_NCCL} USE_RCCL=${USE_RCCL} USE_KINETO=${USE_KINETO} \
+        python setup.py bdist_wheel -d /tmp/$WHEELHOUSE_DIR
+fi
 echo "Finished setup.py bdist at $(date)"
 
 # Build libtorch packages
@@ -255,6 +272,10 @@ ls /tmp/$WHEELHOUSE_DIR
 mkdir -p "/$WHEELHOUSE_DIR"
 mv /tmp/$WHEELHOUSE_DIR/torch*linux*.whl /$WHEELHOUSE_DIR/
 
+if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+    mv /tmp/$WHEELHOUSE_DIR/torch_no_python*.whl /$WHEELHOUSE_DIR/ || true
+fi
+
 if [[ -n "$BUILD_PYTHONLESS" ]]; then
     mkdir -p /$LIBTORCH_HOUSE_DIR
     mv /tmp/$LIBTORCH_HOUSE_DIR/*.zip /$LIBTORCH_HOUSE_DIR
@@ -431,8 +452,16 @@ if [[ -z "$BUILD_PYTHONLESS" ]]; then
   pushd $PYTORCH_ROOT/test
 
   # Install the wheel for this Python version
+  if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+    pip uninstall -y "$TORCH_NO_PYTHON_PACKAGE_NAME" || true
+  fi
+
   pip uninstall -y "$TORCH_PACKAGE_NAME"
 
+  if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+    pip install "$TORCH_NO_PYTHON_PACKAGE_NAME" --no-index -f /$WHEELHOUSE_DIR --no-dependencies -v
+  fi
+
   pip install "$TORCH_PACKAGE_NAME" --no-index -f /$WHEELHOUSE_DIR --no-dependencies -v
 
   # Print info on the libraries installed in this wheel

.ci/manywheel/build_cuda.sh

@@ -66,9 +66,6 @@ case ${CUDA_VERSION} in
             TORCH_CUDA_ARCH_LIST="7.5;8.0;9.0;10.0;12.0+PTX"
         fi
         ;;
-    13.0)
-        TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0;10.0;12.0+PTX"
-        ;;
     12.6)
         TORCH_CUDA_ARCH_LIST="5.0;6.0;7.0;7.5;8.0;8.6;9.0"
         ;;
@@ -113,18 +110,13 @@ DEPS_SONAME=(
 )
 
 
-# CUDA_VERSION 12.*, 13.*
-if [[ $CUDA_VERSION == 12* || $CUDA_VERSION == 13* ]]; then
+# CUDA_VERSION 12.6, 12.8, 12.9
+if [[ $CUDA_VERSION == 12* ]]; then
     export USE_STATIC_CUDNN=0
     # Try parallelizing nvcc as well
-    TORCH_NVCC_FLAGS="-Xfatbin -compress-all --threads 2"
-    # Compress the fatbin with -compress-mode=size for CUDA 13
-    if [[ $CUDA_VERSION == 13* ]]; then
-        export TORCH_NVCC_FLAGS="$TORCH_NVCC_FLAGS -compress-mode=size"
-    fi
+    export TORCH_NVCC_FLAGS="-Xfatbin -compress-all --threads 2"
     if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
         echo "Bundling with cudnn and cublas."
-
         DEPS_LIST+=(
             "/usr/local/cuda/lib64/libcudnn_adv.so.9"
             "/usr/local/cuda/lib64/libcudnn_cnn.so.9"
@@ -134,12 +126,14 @@ if [[ $CUDA_VERSION == 12* || $CUDA_VERSION == 13* ]]; then
             "/usr/local/cuda/lib64/libcudnn_engines_precompiled.so.9"
             "/usr/local/cuda/lib64/libcudnn_heuristic.so.9"
             "/usr/local/cuda/lib64/libcudnn.so.9"
+            "/usr/local/cuda/lib64/libcublas.so.12"
+            "/usr/local/cuda/lib64/libcublasLt.so.12"
             "/usr/local/cuda/lib64/libcusparseLt.so.0"
+            "/usr/local/cuda/lib64/libcudart.so.12"
+            "/usr/local/cuda/lib64/libnvrtc.so.12"
             "/usr/local/cuda/lib64/libnvrtc-builtins.so"
             "/usr/local/cuda/lib64/libcufile.so.0"
             "/usr/local/cuda/lib64/libcufile_rdma.so.1"
-            "/usr/local/cuda/lib64/libnvshmem_host.so.3"
-            "/usr/local/cuda/extras/CUPTI/lib64/libnvperf_host.so"
         )
         DEPS_SONAME+=(
             "libcudnn_adv.so.9"
@@ -150,83 +144,33 @@ if [[ $CUDA_VERSION == 12* || $CUDA_VERSION == 13* ]]; then
             "libcudnn_engines_precompiled.so.9"
             "libcudnn_heuristic.so.9"
             "libcudnn.so.9"
+            "libcublas.so.12"
+            "libcublasLt.so.12"
             "libcusparseLt.so.0"
+            "libcudart.so.12"
+            "libnvrtc.so.12"
             "libnvrtc-builtins.so"
-            "libnvshmem_host.so.3"
             "libcufile.so.0"
             "libcufile_rdma.so.1"
-            "libnvperf_host.so"
         )
-        # Add libnvToolsExt only if CUDA version is not 12.9
-        if [[ $CUDA_VERSION == 13* ]]; then
-            DEPS_LIST+=(
-                "/usr/local/cuda/lib64/libcublas.so.13"
-                "/usr/local/cuda/lib64/libcublasLt.so.13"
-                "/usr/local/cuda/lib64/libcudart.so.13"
-                "/usr/local/cuda/lib64/libnvrtc.so.13"
-                "/usr/local/cuda/extras/CUPTI/lib64/libcupti.so.13"
-                "/usr/local/cuda/lib64/libibverbs.so.1"
-                "/usr/local/cuda/lib64/librdmacm.so.1"
-                "/usr/local/cuda/lib64/libmlx5.so.1"
-                "/usr/local/cuda/lib64/libnl-3.so.200"
-                "/usr/local/cuda/lib64/libnl-route-3.so.200")
-            DEPS_SONAME+=(
-                "libcublas.so.13"
-                "libcublasLt.so.13"
-                "libcudart.so.13"
-                "libnvrtc.so.13"
-                "libcupti.so.13"
-                "libibverbs.so.1"
-                "librdmacm.so.1"
-                "libmlx5.so.1"
-                "libnl-3.so.200"
-                "libnl-route-3.so.200")
-            export USE_CUPTI_SO=1
-            export ATEN_STATIC_CUDA=0
-            export USE_CUDA_STATIC_LINK=0
-            export USE_CUFILE=0
-        else
-            DEPS_LIST+=(
-                "/usr/local/cuda/lib64/libnvToolsExt.so.1"
-                "/usr/local/cuda/lib64/libcublas.so.12"
-                "/usr/local/cuda/lib64/libcublasLt.so.12"
-                "/usr/local/cuda/lib64/libcudart.so.12"
-                "/usr/local/cuda/lib64/libnvrtc.so.12"
-                "/usr/local/cuda/extras/CUPTI/lib64/libcupti.so.12")
-            DEPS_SONAME+=(
-                "libnvToolsExt.so.1"
-                "libcublas.so.12"
-                "libcublasLt.so.12"
-                "libcudart.so.12"
-                "libnvrtc.so.12"
-                "libcupti.so.12")
-        fi
     else
         echo "Using nvidia libs from pypi."
         CUDA_RPATHS=(
+            '$ORIGIN/../../nvidia/cublas/lib'
+            '$ORIGIN/../../nvidia/cuda_cupti/lib'
+            '$ORIGIN/../../nvidia/cuda_nvrtc/lib'
+            '$ORIGIN/../../nvidia/cuda_runtime/lib'
             '$ORIGIN/../../nvidia/cudnn/lib'
-            '$ORIGIN/../../nvidia/nvshmem/lib'
-            '$ORIGIN/../../nvidia/nccl/lib'
+            '$ORIGIN/../../nvidia/cufft/lib'
+            '$ORIGIN/../../nvidia/curand/lib'
+            '$ORIGIN/../../nvidia/cusolver/lib'
+            '$ORIGIN/../../nvidia/cusparse/lib'
             '$ORIGIN/../../nvidia/cusparselt/lib'
+            '$ORIGIN/../../cusparselt/lib'
+            '$ORIGIN/../../nvidia/nccl/lib'
+            '$ORIGIN/../../nvidia/nvtx/lib'
+            '$ORIGIN/../../nvidia/cufile/lib'
         )
-        if [[ $CUDA_VERSION == 13* ]]; then
-            CUDA_RPATHS+=('$ORIGIN/../../nvidia/cu13/lib')
-        else
-            CUDA_RPATHS+=(
-                '$ORIGIN/../../nvidia/cublas/lib'
-                '$ORIGIN/../../nvidia/cuda_cupti/lib'
-                '$ORIGIN/../../nvidia/cuda_nvrtc/lib'
-                '$ORIGIN/../../nvidia/cuda_runtime/lib'
-                '$ORIGIN/../../nvidia/cufft/lib'
-                '$ORIGIN/../../nvidia/curand/lib'
-                '$ORIGIN/../../nvidia/cusolver/lib'
-                '$ORIGIN/../../nvidia/cusparse/lib'
-                '$ORIGIN/../../cusparselt/lib'
-                '$ORIGIN/../../nvidia/nvtx/lib'
-                '$ORIGIN/../../nvidia/cufile/lib'
-            )
-        fi
-
         CUDA_RPATHS=$(IFS=: ; echo "${CUDA_RPATHS[*]}")
         export C_SO_RPATH=$CUDA_RPATHS':$ORIGIN:$ORIGIN/lib'
         export LIB_SO_RPATH=$CUDA_RPATHS':$ORIGIN'

.ci/manywheel/build_libtorch.sh

@@ -92,7 +92,7 @@ if [[ -z "$PYTORCH_ROOT" ]]; then
     exit 1
 fi
 pushd "$PYTORCH_ROOT"
-retry pip install -qUr requirements-build.txt
+retry pip install -q cmake
 python setup.py clean
 retry pip install -qr requirements.txt
 retry pip install -q numpy==2.0.1
@@ -104,7 +104,7 @@ if [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
     export ROCclr_DIR=/opt/rocm/rocclr/lib/cmake/rocclr
 fi
 
-echo "Calling -m pip install . -v --no-build-isolation at $(date)"
+echo "Calling setup.py install at $(date)"
 
 if [[ $LIBTORCH_VARIANT = *"static"* ]]; then
     STATIC_CMAKE_FLAG="-DTORCH_STATIC=1"
@@ -120,7 +120,7 @@ fi
         # TODO: Remove this flag once https://github.com/pytorch/pytorch/issues/55952 is closed
         CFLAGS='-Wno-deprecated-declarations' \
         BUILD_LIBTORCH_CPU_WITH_DEBUG=1 \
-        python -m pip install --no-build-isolation -v .
+        python setup.py install
 
     mkdir -p libtorch/{lib,bin,include,share}
 

.ci/manywheel/build_rocm.sh

@@ -194,7 +194,7 @@ ROCBLAS_LIB_SRC=$ROCM_HOME/lib/rocblas/library
 ROCBLAS_LIB_DST=lib/rocblas/library
 ROCBLAS_ARCH_SPECIFIC_FILES=$(ls $ROCBLAS_LIB_SRC | grep -E $ARCH)
 ROCBLAS_OTHER_FILES=$(ls $ROCBLAS_LIB_SRC | grep -v gfx)
-ROCBLAS_LIB_FILES=($ROCBLAS_ARCH_SPECIFIC_FILES $ROCBLAS_OTHER_FILES)
+ROCBLAS_LIB_FILES=($ROCBLAS_ARCH_SPECIFIC_FILES $OTHER_FILES)
 
 # hipblaslt library files
 HIPBLASLT_LIB_SRC=$ROCM_HOME/lib/hipblaslt/library

.ci/manywheel/build_xpu.sh

@@ -25,7 +25,6 @@ source /opt/intel/oneapi/mpi/latest/env/vars.sh
 export USE_STATIC_MKL=1
 export USE_ONEMKL=1
 export USE_XCCL=1
-export USE_MPI=0
 
 WHEELHOUSE_DIR="wheelhousexpu"
 LIBTORCH_HOUSE_DIR="libtorch_housexpu"

.ci/onnx/test.sh

@@ -19,7 +19,7 @@ git config --global --add safe.directory /var/lib/jenkins/workspace
 
 if [[ "$BUILD_ENVIRONMENT" == *onnx* ]]; then
   # TODO: This can be removed later once vision is also part of the Docker image
-  pip install -q --no-use-pep517 "git+https://github.com/pytorch/vision.git@$(cat .github/ci_commit_pins/vision.txt)"
+  pip install -q --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@$(cat .github/ci_commit_pins/vision.txt)"
   # JIT C++ extensions require ninja, so put it into PATH.
   export PATH="/var/lib/jenkins/.local/bin:$PATH"
   # NB: ONNX test is fast (~15m) so it's ok to retry it few more times to avoid any flaky issue, we

.ci/pytorch/build-mobile.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# DO NOT ADD 'set -x' not to reveal CircleCI secret context environment variables
+set -eu -o pipefail
+
+# This script uses linux host toolchain + mobile build options in order to
+# build & test mobile libtorch without having to setup Android/iOS
+# toolchain/simulator.
+
+# shellcheck source=./common.sh
+source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
+# shellcheck source=./common-build.sh
+source "$(dirname "${BASH_SOURCE[0]}")/common-build.sh"
+
+# Install torch & torchvision - used to download & trace test model.
+# Ideally we should use the libtorch built on the PR so that backward
+# incompatible changes won't break this script - but it will significantly slow
+# down mobile CI jobs.
+# Here we install nightly instead of stable so that we have an option to
+# temporarily skip mobile CI jobs on BC-breaking PRs until they are in nightly.
+retry pip install --pre torch torchvision \
+  -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html \
+  --progress-bar off
+
+# Run end-to-end process of building mobile library, linking into the predictor
+# binary, and running forward pass with a real model.
+if [[ "$BUILD_ENVIRONMENT" == *-mobile-custom-build-static* ]]; then
+  TEST_CUSTOM_BUILD_STATIC=1 test/mobile/custom_build/build.sh
+elif [[ "$BUILD_ENVIRONMENT" == *-mobile-lightweight-dispatch* ]]; then
+  test/mobile/lightweight_dispatch/build.sh
+else
+  TEST_DEFAULT_BUILD=1 test/mobile/custom_build/build.sh
+fi
+
+print_sccache_stats

.ci/pytorch/build.sh

@@ -11,6 +11,10 @@ source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 # shellcheck source=./common-build.sh
 source "$(dirname "${BASH_SOURCE[0]}")/common-build.sh"
 
+if [[ "$BUILD_ENVIRONMENT" == *-mobile-*build* ]]; then
+  exec "$(dirname "${BASH_SOURCE[0]}")/build-mobile.sh" "$@"
+fi
+
 echo "Python version:"
 python --version
 
@@ -50,6 +54,9 @@ if [[ ${BUILD_ENVIRONMENT} == *"parallelnative"* ]]; then
   export ATEN_THREADING=NATIVE
 fi
 
+# Enable LLVM dependency for TensorExpr testing
+export USE_LLVM=/opt/llvm
+export LLVM_DIR=/opt/llvm/lib/cmake/llvm
 
 if ! which conda; then
   # In ROCm CIs, we are doing cross compilation on build machines with
@@ -92,27 +99,6 @@ if [[ "$BUILD_ENVIRONMENT" == *aarch64* ]]; then
   export ACL_ROOT_DIR=/ComputeLibrary
 fi
 
-if [[ "$BUILD_ENVIRONMENT" == *riscv64* ]]; then
-  if [[ -f /opt/riscv-cross-env/bin/activate ]]; then
-    # shellcheck disable=SC1091
-    source /opt/riscv-cross-env/bin/activate
-  else
-    echo "Activation file not found"
-    exit 1
-  fi
-
-  export CMAKE_CROSSCOMPILING=TRUE
-  export CMAKE_SYSTEM_NAME=Linux
-  export CMAKE_SYSTEM_PROCESSOR=riscv64
-
-  export USE_CUDA=0
-  export USE_MKLDNN=0
-
-  export SLEEF_TARGET_EXEC_USE_QEMU=ON
-  sudo chown -R jenkins /var/lib/jenkins/workspace /opt
-
-fi
-
 if [[ "$BUILD_ENVIRONMENT" == *libtorch* ]]; then
   POSSIBLE_JAVA_HOMES=()
   POSSIBLE_JAVA_HOMES+=(/usr/local)
@@ -138,8 +124,26 @@ if [[ "$BUILD_ENVIRONMENT" == *libtorch* ]]; then
 fi
 
 # Use special scripts for Android builds
+if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
+  export ANDROID_NDK=/opt/ndk
+  build_args=()
+  if [[ "${BUILD_ENVIRONMENT}" == *-arm-v7a* ]]; then
+    build_args+=("-DANDROID_ABI=armeabi-v7a")
+  elif [[ "${BUILD_ENVIRONMENT}" == *-arm-v8a* ]]; then
+    build_args+=("-DANDROID_ABI=arm64-v8a")
+  elif [[ "${BUILD_ENVIRONMENT}" == *-x86_32* ]]; then
+    build_args+=("-DANDROID_ABI=x86")
+  elif [[ "${BUILD_ENVIRONMENT}" == *-x86_64* ]]; then
+    build_args+=("-DANDROID_ABI=x86_64")
+  fi
+  if [[ "${BUILD_ENVIRONMENT}" == *vulkan* ]]; then
+    build_args+=("-DUSE_VULKAN=ON")
+  fi
+  build_args+=("-DUSE_LITE_INTERPRETER_PROFILER=OFF")
+  exec ./scripts/build_android.sh "${build_args[@]}" "$@"
+fi
 
-if [[ "$BUILD_ENVIRONMENT" == *vulkan* ]]; then
+if [[ "$BUILD_ENVIRONMENT" != *android* && "$BUILD_ENVIRONMENT" == *vulkan* ]]; then
   export USE_VULKAN=1
   # shellcheck disable=SC1091
   source /var/lib/jenkins/vulkansdk/setup-env.sh
@@ -173,7 +177,6 @@ if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
   source /opt/intel/oneapi/mpi/latest/env/vars.sh
   # Enable XCCL build
   export USE_XCCL=1
-  export USE_MPI=0
   # XPU kineto feature dependencies are not fully ready, disable kineto build as temp WA
   export USE_KINETO=0
   export TORCH_XPU_ARCH_LIST=pvc
@@ -195,16 +198,10 @@ fi
 
 # We only build FlashAttention files for CUDA 8.0+, and they require large amounts of
 # memory to build and will OOM
-
-if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && echo "${TORCH_CUDA_ARCH_LIST}" | tr ' ' '\n' | sed 's/$/>= 8.0/' | bc | grep -q 1; then
-  J=2  # default to 2 jobs
-  case "$RUNNER" in
-    linux.12xlarge.memory|linux.24xlarge.memory)
-      J=24
-      ;;
-  esac
-  echo "Building FlashAttention with job limit $J"
-  export BUILD_CUSTOM_STEP="ninja -C build flash_attention -j ${J}"
+if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && [[ 1 -eq $(echo "${TORCH_CUDA_ARCH_LIST} >= 8.0" | bc) ]] && [ -z "$MAX_JOBS_OVERRIDE" ]; then
+  echo "WARNING: FlashAttention files require large amounts of memory to build and will OOM"
+  echo "Setting MAX_JOBS=(nproc-2)/3 to reduce memory usage"
+  export MAX_JOBS="$(( $(nproc --ignore=2) / 3 ))"
 fi
 
 if [[ "${BUILD_ENVIRONMENT}" == *clang* ]]; then
@@ -219,6 +216,7 @@ if [[ "$BUILD_ENVIRONMENT" == *-clang*-asan* ]]; then
   export USE_ASAN=1
   export REL_WITH_DEB_INFO=1
   export UBSAN_FLAGS="-fno-sanitize-recover=all"
+  unset USE_LLVM
 fi
 
 if [[ "${BUILD_ENVIRONMENT}" == *no-ops* ]]; then
@@ -229,7 +227,7 @@ if [[ "${BUILD_ENVIRONMENT}" == *-pch* ]]; then
     export USE_PRECOMPILED_HEADERS=1
 fi
 
-if [[ "${BUILD_ENVIRONMENT}" != *cuda* ]]; then
+if [[ "${BUILD_ENVIRONMENT}" != *android* && "${BUILD_ENVIRONMENT}" != *cuda* ]]; then
   export BUILD_STATIC_RUNTIME_BENCHMARK=ON
 fi
 
@@ -239,7 +237,7 @@ fi
 
 # Do not change workspace permissions for ROCm and s390x CI jobs
 # as it can leave workspace with bad permissions for cancelled jobs
-if [[ "$BUILD_ENVIRONMENT" != *rocm* && "$BUILD_ENVIRONMENT" != *s390x* && "$BUILD_ENVIRONMENT" != *riscv64* && -d /var/lib/jenkins/workspace ]]; then
+if [[ "$BUILD_ENVIRONMENT" != *rocm* && "$BUILD_ENVIRONMENT" != *s390x* && -d /var/lib/jenkins/workspace ]]; then
   # Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
   WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
   cleanup_workspace() {
@@ -284,38 +282,32 @@ else
     # XLA test build fails when WERROR=1
     # set only when building other architectures
     # or building non-XLA tests.
-    if [[ "$BUILD_ENVIRONMENT" != *rocm*  && "$BUILD_ENVIRONMENT" != *xla* && "$BUILD_ENVIRONMENT" != *riscv64* ]]; then
+    if [[ "$BUILD_ENVIRONMENT" != *rocm*  &&
+          "$BUILD_ENVIRONMENT" != *xla* ]]; then
       # Install numpy-2.0.2 for builds which are backward compatible with 1.X
       python -mpip install numpy==2.0.2
 
       WERROR=1 python setup.py clean
 
-      WERROR=1 python -m build --wheel --no-isolation
+      if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+        python3 tools/packaging/split_wheel.py bdist_wheel
+      else
+        WERROR=1 python setup.py bdist_wheel
+      fi
     else
       python setup.py clean
       if [[ "$BUILD_ENVIRONMENT" == *xla* ]]; then
         source .ci/pytorch/install_cache_xla.sh
       fi
-      python -m build --wheel --no-isolation
+      if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
+        echo "USE_SPLIT_BUILD cannot be used with xla or rocm"
+        exit 1
+      else
+        python setup.py bdist_wheel
+      fi
     fi
     pip_install_whl "$(echo dist/*.whl)"
 
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *vision* ]]; then
-      install_torchvision
-    fi
-
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *audio* ]]; then
-      install_torchaudio
-    fi
-
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchrec* || "${BUILD_ADDITIONAL_PACKAGES:-}" == *fbgemm* ]]; then
-      install_torchrec_and_fbgemm
-    fi
-
-    if [[ "${BUILD_ADDITIONAL_PACKAGES:-}" == *torchao* ]]; then
-      install_torchao
-    fi
-
     if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
       echo "Checking that xpu is compiled"
       pushd dist/
@@ -403,8 +395,10 @@ else
     # This is an attempt to mitigate flaky libtorch build OOM error. By default, the build parallelization
     # is set to be the number of CPU minus 2. So, let's try a more conservative value here. A 4xlarge has
     # 16 CPUs
-    MAX_JOBS=$(nproc --ignore=4)
-    export MAX_JOBS
+    if [ -z "$MAX_JOBS_OVERRIDE" ]; then
+      MAX_JOBS=$(nproc --ignore=4)
+      export MAX_JOBS
+    fi
 
     # NB: Install outside of source directory (at the same level as the root
     # pytorch folder) so that it doesn't get cleaned away prior to docker push.
@@ -421,7 +415,7 @@ if [[ "$BUILD_ENVIRONMENT" != *libtorch* && "$BUILD_ENVIRONMENT" != *bazel* ]];
   # don't do this for libtorch as libtorch is C++ only and thus won't have python tests run on its build
   python tools/stats/export_test_times.py
 fi
-# don't do this for bazel or s390x or riscv64 as they don't use sccache
-if [[ "$BUILD_ENVIRONMENT" != *s390x* && "$BUILD_ENVIRONMENT" != *riscv64* && "$BUILD_ENVIRONMENT" != *-bazel-* ]]; then
+# don't do this for bazel or s390x as they don't use sccache
+if [[ "$BUILD_ENVIRONMENT" != *s390x* && "$BUILD_ENVIRONMENT" != *-bazel-* ]]; then
   print_sccache_stats
 fi

.ci/pytorch/check_binary.sh

@@ -300,3 +300,24 @@ except RuntimeError as e:
     exit 1
   fi
 fi
+
+###############################################################################
+# Check for C++ ABI compatibility to GCC-11 - GCC 13
+###############################################################################
+if [[ "$(uname)" == 'Linux' &&  "$PACKAGE_TYPE" == 'manywheel' ]]; then
+  pushd /tmp
+  # Per https://gcc.gnu.org/onlinedocs/gcc/C_002b_002b-Dialect-Options.html
+  # gcc-11 is ABI16, gcc-13 is ABI18, gcc-14 is ABI19
+  # gcc 11 - CUDA 11.8, xpu, rocm
+  # gcc 13 - CUDA 12.6, 12.8 and cpu
+  # Please see issue for reference: https://github.com/pytorch/pytorch/issues/152426
+  if [[ "$(uname -m)" == "s390x" ]]; then
+    cxx_abi="19"
+  elif [[ "$DESIRED_CUDA" != 'xpu' && "$DESIRED_CUDA" != 'rocm'* ]]; then
+    cxx_abi="18"
+  else
+    cxx_abi="16"
+  fi
+  python -c "import torch; exit(0 if torch._C._PYBIND11_BUILD_ABI == '_cxxabi10${cxx_abi}' else 1)"
+  popd
+fi

.ci/pytorch/common-build.sh

@@ -13,13 +13,6 @@ if [[ "$BUILD_ENVIRONMENT" != *win-* ]]; then
     fi
 
     if which sccache > /dev/null; then
-        # Clear SCCACHE_BUCKET and SCCACHE_REGION if they are empty, otherwise
-        # sccache will complain about invalid bucket configuration
-        if [[ -z "${SCCACHE_BUCKET:-}" ]]; then
-          unset SCCACHE_BUCKET
-          unset SCCACHE_REGION
-        fi
-
         # Save sccache logs to file
         sccache --stop-server > /dev/null  2>&1 || true
         rm -f ~/sccache_error.log || true

.ci/pytorch/common_utils.sh

@@ -78,34 +78,6 @@ function pip_install_whl() {
   fi
 }
 
-function pip_build_and_install() {
-  local build_target=$1
-  local wheel_dir=$2
-
-  local found_whl=0
-  for file in "${wheel_dir}"/*.whl
-  do
-    if [[ -f "${file}" ]]; then
-      found_whl=1
-      break
-    fi
-  done
-
-  # Build the wheel if it doesn't exist
-  if [ "${found_whl}" == "0" ]; then
-    python3 -m pip wheel \
-      --no-build-isolation \
-      --no-deps \
-      --no-use-pep517 \
-      -w "${wheel_dir}" \
-      "${build_target}"
-  fi
-
-  for file in "${wheel_dir}"/*.whl
-  do
-    pip_install_whl "${file}"
-  done
-}
 
 function pip_install() {
   # retry 3 times
@@ -149,23 +121,17 @@ function get_pinned_commit() {
   cat .github/ci_commit_pins/"${1}".txt
 }
 
-function detect_cuda_arch() {
-  if [[ "${BUILD_ENVIRONMENT}" == *cuda* ]]; then
-    if command -v nvidia-smi; then
-      TORCH_CUDA_ARCH_LIST=$(nvidia-smi --query-gpu=compute_cap --format=csv | tail -n 1)
-    elif [[ "${TEST_CONFIG}" == *nogpu* ]]; then
-      # There won't be nvidia-smi in nogpu tests, so just set TORCH_CUDA_ARCH_LIST to the default
-      # minimum supported value here
-      TORCH_CUDA_ARCH_LIST=8.0
-    fi
-    export TORCH_CUDA_ARCH_LIST
-  fi
-}
-
 function install_torchaudio() {
   local commit
   commit=$(get_pinned_commit audio)
-  pip_build_and_install "git+https://github.com/pytorch/audio.git@${commit}" dist/audio
+  if [[ "$1" == "cuda" ]]; then
+    # TODO: This is better to be passed as a parameter from _linux-test workflow
+    # so that it can be consistent with what is set in build
+    TORCH_CUDA_ARCH_LIST="8.0;8.6" pip_install --no-use-pep517 --user "git+https://github.com/pytorch/audio.git@${commit}"
+  else
+    pip_install --no-use-pep517 --user "git+https://github.com/pytorch/audio.git@${commit}"
+  fi
+
 }
 
 function install_torchtext() {
@@ -173,8 +139,8 @@ function install_torchtext() {
   local text_commit
   data_commit=$(get_pinned_commit data)
   text_commit=$(get_pinned_commit text)
-  pip_build_and_install "git+https://github.com/pytorch/data.git@${data_commit}" dist/data
-  pip_build_and_install "git+https://github.com/pytorch/text.git@${text_commit}" dist/text
+  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/data.git@${data_commit}"
+  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/text.git@${text_commit}"
 }
 
 function install_torchvision() {
@@ -187,14 +153,7 @@ function install_torchvision() {
     echo 'char* dlerror(void) { return "";}'|gcc -fpic -shared -o "${HOME}/dlerror.so" -x c -
     LD_PRELOAD=${orig_preload}:${HOME}/dlerror.so
   fi
-
-  if [[ "${BUILD_ENVIRONMENT}" == *cuda* ]]; then
-    # Not sure if both are needed, but why not
-    export FORCE_CUDA=1
-    export WITH_CUDA=1
-  fi
-  pip_build_and_install "git+https://github.com/pytorch/vision.git@${commit}" dist/vision
-
+  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/vision.git@${commit}"
   if [ -n "${LD_PRELOAD}" ]; then
     LD_PRELOAD=${orig_preload}
   fi
@@ -214,85 +173,31 @@ function install_torchrec_and_fbgemm() {
 
   if [[ "$BUILD_ENVIRONMENT" == *rocm* ]] ; then
     # install torchrec first because it installs fbgemm nightly on top of rocm fbgemm
-    pip_build_and_install "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}" dist/torchrec
+    pip_install --no-use-pep517 --user "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}"
     pip_uninstall fbgemm-gpu-nightly
 
-    # Set ROCM_HOME isn't available, use ROCM_PATH if set or /opt/rocm
-    ROCM_HOME="${ROCM_HOME:-${ROCM_PATH:-/opt/rocm}}"
-
-    # Find rocm_version.h header file for ROCm version extract
-    rocm_version_h="${ROCM_HOME}/include/rocm-core/rocm_version.h"
-    if [ ! -f "$rocm_version_h" ]; then
-        rocm_version_h="${ROCM_HOME}/include/rocm_version.h"
-    fi
-
-    # Error out if rocm_version.h not found
-    if [ ! -f "$rocm_version_h" ]; then
-        echo "Error: rocm_version.h not found in expected locations." >&2
-        exit 1
-    fi
-
-    # Extract major, minor and patch ROCm version numbers
-    MAJOR_VERSION=$(grep 'ROCM_VERSION_MAJOR' "$rocm_version_h" | awk '{print $3}')
-    MINOR_VERSION=$(grep 'ROCM_VERSION_MINOR' "$rocm_version_h" | awk '{print $3}')
-    PATCH_VERSION=$(grep 'ROCM_VERSION_PATCH' "$rocm_version_h" | awk '{print $3}')
-    ROCM_INT=$((MAJOR_VERSION * 10000 + MINOR_VERSION * 100 + PATCH_VERSION))
-    echo "ROCm version: $ROCM_INT"
-    export BUILD_ROCM_VERSION="$MAJOR_VERSION.$MINOR_VERSION"
-
     pip_install tabulate  # needed for newer fbgemm
     pip_install patchelf  # needed for rocm fbgemm
-
-    local wheel_dir=dist/fbgemm_gpu
-    local found_whl=0
-    for file in "${wheel_dir}"/*.whl
-    do
-      if [[ -f "${file}" ]]; then
-        found_whl=1
-        break
-      fi
-    done
-
-    # Build the wheel if it doesn't exist
-    if [ "${found_whl}" == "0" ]; then
-      git clone --recursive https://github.com/pytorch/fbgemm
-      pushd fbgemm/fbgemm_gpu
-      git checkout "${fbgemm_commit}" --recurse-submodules
-      # until the fbgemm_commit includes the tbb patch
-      patch <<'EOF'
---- a/FbgemmGpu.cmake
-+++ b/FbgemmGpu.cmake
-@@ -184,5 +184,6 @@ gpu_cpp_library(
-     fbgemm_gpu_tbe_cache
-     fbgemm_gpu_tbe_optimizers
-     fbgemm_gpu_tbe_utils
-+    tbb
-   DESTINATION
-     fbgemm_gpu)
-EOF
-      python setup.py bdist_wheel --build-variant=rocm
-      popd
-
-      # Save the wheel before cleaning up
-      mkdir -p dist/fbgemm_gpu
-      cp fbgemm/fbgemm_gpu/dist/*.whl dist/fbgemm_gpu
-    fi
-
-    for file in "${wheel_dir}"/*.whl
-    do
-      pip_install_whl "${file}"
-    done
-
+    git clone --recursive https://github.com/pytorch/fbgemm
+    pushd fbgemm/fbgemm_gpu
+    git checkout "${fbgemm_commit}"
+    python setup.py install \
+      --package_variant=rocm \
+      -DHIP_ROOT_DIR="${ROCM_PATH}" \
+      -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
+      -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"
+    popd
     rm -rf fbgemm
   else
-    pip_build_and_install "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}" dist/torchrec
-    pip_build_and_install "git+https://github.com/pytorch/FBGEMM.git@${fbgemm_commit}#subdirectory=fbgemm_gpu" dist/fbgemm_gpu
+    # See https://github.com/pytorch/pytorch/issues/106971
+    CUDA_PATH=/usr/local/cuda-12.1 pip_install --no-use-pep517 --user "git+https://github.com/pytorch/FBGEMM.git@${fbgemm_commit}#egg=fbgemm-gpu&subdirectory=fbgemm_gpu"
+    pip_install --no-use-pep517 --user "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}"
   fi
 }
 
 function clone_pytorch_xla() {
   if [[ ! -d ./xla ]]; then
-    git clone --recursive --quiet https://github.com/pytorch/xla.git
+    git clone --recursive -b r2.8 https://github.com/pytorch/xla.git
     pushd xla
     # pin the xla hash so that we don't get broken by changes to xla
     git checkout "$(cat ../.github/ci_commit_pins/xla.txt)"
@@ -302,10 +207,34 @@ function clone_pytorch_xla() {
   fi
 }
 
+function checkout_install_torchbench() {
+  local commit
+  commit=$(get_pinned_commit torchbench)
+  git clone https://github.com/pytorch/benchmark torchbench
+  pushd torchbench
+  git checkout "$commit"
+
+  if [ "$1" ]; then
+    python install.py --continue_on_fail models "$@"
+  else
+    # Occasionally the installation may fail on one model but it is ok to continue
+    # to install and test other models
+    python install.py --continue_on_fail
+  fi
+
+  # TODO (huydhn): transformers-4.44.2 added by https://github.com/pytorch/benchmark/pull/2488
+  # is regressing speedup metric. This needs to be investigated further
+  pip install transformers==4.38.1
+
+  echo "Print all dependencies after TorchBench is installed"
+  python -mpip freeze
+  popd
+}
+
 function install_torchao() {
   local commit
   commit=$(get_pinned_commit torchao)
-  pip_build_and_install "git+https://github.com/pytorch/ao.git@${commit}" dist/ao
+  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/ao.git@${commit}"
 }
 
 function print_sccache_stats() {

.ci/pytorch/cpp_doc_push_script.sh

@@ -58,7 +58,7 @@ time python tools/setup_helpers/generate_code.py \
 
 # Build the docs
 pushd docs/cpp
-time make VERBOSE=1 html
+time make VERBOSE=1 html -j
 
 popd
 popd

.ci/pytorch/create_test_cert.py

@@ -0,0 +1,123 @@
+from datetime import datetime, timedelta, timezone
+from tempfile import mkdtemp
+
+from cryptography import x509
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.x509.oid import NameOID
+
+
+temp_dir = mkdtemp()
+print(temp_dir)
+
+
+def genrsa(path):
+    key = rsa.generate_private_key(
+        public_exponent=65537,
+        key_size=2048,
+    )
+    with open(path, "wb") as f:
+        f.write(
+            key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.TraditionalOpenSSL,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+        )
+    return key
+
+
+def create_cert(path, C, ST, L, O, key):
+    subject = issuer = x509.Name(
+        [
+            x509.NameAttribute(NameOID.COUNTRY_NAME, C),
+            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, ST),
+            x509.NameAttribute(NameOID.LOCALITY_NAME, L),
+            x509.NameAttribute(NameOID.ORGANIZATION_NAME, O),
+        ]
+    )
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(issuer)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.now(timezone.utc))
+        .not_valid_after(
+            # Our certificate will be valid for 10 days
+            datetime.now(timezone.utc) + timedelta(days=10)
+        )
+        .add_extension(
+            x509.BasicConstraints(ca=True, path_length=None),
+            critical=True,
+        )
+        .sign(key, hashes.SHA256())
+    )
+    # Write our certificate out to disk.
+    with open(path, "wb") as f:
+        f.write(cert.public_bytes(serialization.Encoding.PEM))
+    return cert
+
+
+def create_req(path, C, ST, L, O, key):
+    csr = (
+        x509.CertificateSigningRequestBuilder()
+        .subject_name(
+            x509.Name(
+                [
+                    # Provide various details about who we are.
+                    x509.NameAttribute(NameOID.COUNTRY_NAME, C),
+                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, ST),
+                    x509.NameAttribute(NameOID.LOCALITY_NAME, L),
+                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, O),
+                ]
+            )
+        )
+        .sign(key, hashes.SHA256())
+    )
+    with open(path, "wb") as f:
+        f.write(csr.public_bytes(serialization.Encoding.PEM))
+    return csr
+
+
+def sign_certificate_request(path, csr_cert, ca_cert, private_ca_key):
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(csr_cert.subject)
+        .issuer_name(ca_cert.subject)
+        .public_key(csr_cert.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(datetime.now(timezone.utc))
+        .not_valid_after(
+            # Our certificate will be valid for 10 days
+            datetime.now(timezone.utc) + timedelta(days=10)
+            # Sign our certificate with our private key
+        )
+        .sign(private_ca_key, hashes.SHA256())
+    )
+    with open(path, "wb") as f:
+        f.write(cert.public_bytes(serialization.Encoding.PEM))
+    return cert
+
+
+ca_key = genrsa(temp_dir + "/ca.key")
+ca_cert = create_cert(
+    temp_dir + "/ca.pem",
+    "US",
+    "New York",
+    "New York",
+    "Gloo Certificate Authority",
+    ca_key,
+)
+
+pkey = genrsa(temp_dir + "/pkey.key")
+csr = create_req(
+    temp_dir + "/csr.csr",
+    "US",
+    "California",
+    "San Francisco",
+    "Gloo Testing Company",
+    pkey,
+)
+
+cert = sign_certificate_request(temp_dir + "/cert.pem", csr, ca_cert, ca_key)

.ci/pytorch/functorch_doc_push_script.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+# This is where the local pytorch install in the docker image is located
+pt_checkout="/var/lib/jenkins/workspace"
+source "$pt_checkout/.ci/pytorch/common_utils.sh"
+echo "functorch_doc_push_script.sh: Invoked with $*"
+
+set -ex -o pipefail
+
+version=${DOCS_VERSION:-nightly}
+echo "version: $version"
+
+# Build functorch docs
+pushd $pt_checkout/functorch/docs
+make html
+popd
+
+git clone https://github.com/pytorch/functorch -b gh-pages --depth 1 functorch_ghpages
+pushd functorch_ghpages
+
+if [ "$version" == "main" ]; then
+  version=nightly
+fi
+
+git rm -rf "$version" || true
+mv "$pt_checkout/functorch/docs/build/html" "$version"
+
+git add "$version" || true
+git status
+git config user.email "soumith+bot@pytorch.org"
+git config user.name "pytorchbot"
+# If there aren't changes, don't make a commit; push is no-op
+git commit -m "Generate Python docs from pytorch/pytorch@${GITHUB_SHA}" || true
+git status
+
+if [[ "${WITH_PUSH:-}" == true ]]; then
+  git push -u origin gh-pages
+fi
+
+popd

.ci/pytorch/macos-build.sh

@@ -36,11 +36,11 @@ fi
 print_cmake_info
 if [[ ${BUILD_ENVIRONMENT} == *"distributed"* ]]; then
   # Needed for inductor benchmarks, as lots of HF networks make `torch.distribtued` calls
-  USE_DISTRIBUTED=1 USE_OPENMP=1 WERROR=1 python -m build --wheel --no-isolation
+  USE_DISTRIBUTED=1 USE_OPENMP=1 WERROR=1 python setup.py bdist_wheel
 else
   # Explicitly set USE_DISTRIBUTED=0 to align with the default build config on mac. This also serves as the sole CI config that tests
   # that building with USE_DISTRIBUTED=0 works at all. See https://github.com/pytorch/pytorch/issues/86448
-  USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python -m build --wheel --no-isolation -C--build-option=--plat-name=macosx_11_0_arm64
+  USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python setup.py bdist_wheel --plat-name macosx_11_0_arm64
 fi
 if which sccache > /dev/null; then
   print_sccache_stats

.ci/pytorch/macos-test.sh

@@ -55,7 +55,7 @@ test_python_shard() {
 
   setup_test_python
 
-  time python test/run_test.py --verbose --exclude-jit-executor --exclude-distributed-tests --exclude-quantization-tests --shard "$1" "$NUM_TEST_SHARDS"
+  time python test/run_test.py --verbose --exclude-jit-executor --exclude-distributed-tests --shard "$1" "$NUM_TEST_SHARDS"
 
   assert_git_not_dirty
 }
@@ -157,34 +157,6 @@ test_jit_hooks() {
   assert_git_not_dirty
 }
 
-# Shellcheck doesn't like it when you pass no arguments to a function
-# that can take args. See https://www.shellcheck.net/wiki/SC2120
-# shellcheck disable=SC2120
-checkout_install_torchbench() {
-  local commit
-  commit=$(cat .ci/docker/ci_commit_pins/torchbench.txt)
-  git clone https://github.com/pytorch/benchmark torchbench
-  pushd torchbench
-  git checkout "$commit"
-
-  if [ "$1" ]; then
-    python install.py --continue_on_fail models "$@"
-  else
-    # Occasionally the installation may fail on one model but it is ok to continue
-    # to install and test other models
-    python install.py --continue_on_fail
-  fi
-  popd
-
-  pip install -r .ci/docker/ci_commit_pins/huggingface-requirements.txt
-  # https://github.com/pytorch/pytorch/issues/160689 to remove torchao because
-  # its current version 0.12.0 doesn't work with transformers 4.54.0
-  pip uninstall -y torchao
-
-  echo "Print all dependencies after TorchBench is installed"
-  python -mpip freeze
-}
-
 torchbench_setup_macos() {
   git clone --recursive https://github.com/pytorch/vision torchvision
   git clone --recursive https://github.com/pytorch/audio torchaudio
@@ -195,7 +167,7 @@ torchbench_setup_macos() {
   git checkout "$(cat ../.github/ci_commit_pins/vision.txt)"
   git submodule update --init --recursive
   python setup.py clean
-  python -m pip install -e . -v --no-build-isolation
+  python setup.py develop
   popd
 
   pushd torchaudio
@@ -204,14 +176,16 @@ torchbench_setup_macos() {
   git submodule update --init --recursive
   python setup.py clean
   #TODO: Remove me, when figure out how to make TorchAudio find brew installed openmp
-  USE_OPENMP=0 python -m pip install -e . -v --no-build-isolation
+  USE_OPENMP=0 python setup.py develop
   popd
 
+  # Shellcheck doesn't like it when you pass no arguments to a function that can take args. See https://www.shellcheck.net/wiki/SC2120
+  # shellcheck disable=SC2119,SC2120
   checkout_install_torchbench
 }
 
 pip_benchmark_deps() {
-  python -mpip install --no-input requests cython scikit-learn six
+  python -mpip install --no-input astunparse requests cython scikit-learn
 }
 
 
@@ -302,47 +276,6 @@ test_torchbench_smoketest() {
     fi
 
   done
-  echo "Pytorch benchmark on mps device completed"
-}
-
-test_aoti_torchbench_smoketest() {
-  print_cmake_info
-
-  echo "Launching AOTInductor torchbench setup"
-  pip_benchmark_deps
-  # shellcheck disable=SC2119,SC2120
-  torchbench_setup_macos
-
-  TEST_REPORTS_DIR=$(pwd)/test/test-reports
-  mkdir -p "$TEST_REPORTS_DIR"
-
-  local device=mps
-  local dtypes=(undefined float16 bfloat16 notset)
-  local dtype=${dtypes[$1]}
-  local models=(hf_T5 llama BERT_pytorch dcgan hf_GPT2 yolov3 resnet152 sam sam_fast pytorch_unet stable_diffusion_text_encoder speech_transformer Super_SloMo doctr_det_predictor doctr_reco_predictor timm_resnet timm_vovnet vgg16)
-
-  echo "Launching torchbench inference performance run for AOT Inductor and dtype ${dtype}"
-  local dtype_arg="--${dtype}"
-  if [ "$dtype" == notset ]; then
-      dtype_arg="--float32"
-  fi
-  touch "$TEST_REPORTS_DIR/aot_inductor_torchbench_${dtype}_inference_${device}_performance.csv"
-  for model in "${models[@]}"; do
-    PYTHONPATH="$(pwd)"/torchbench python benchmarks/dynamo/torchbench.py \
-      --performance --only "$model" --export-aot-inductor --inference --devices "$device" "$dtype_arg" \
-      --output "$TEST_REPORTS_DIR/aot_inductor_torchbench_${dtype}_inference_${device}_performance.csv" || true
-    PYTHONPATH="$(pwd)"/torchbench python benchmarks/dynamo/torchbench.py \
-      --accuracy --only "$model" --export-aot-inductor --inference --devices "$device" "$dtype_arg" \
-      --output "$TEST_REPORTS_DIR/aot_inductor_torchbench_${dtype}_inference_${device}_accuracy.csv" || true
-  done
-
-  echo "Launching HuggingFace inference performance run for AOT Inductor and dtype ${dtype}"
-  PYTHONPATH="$(pwd)"/torchbench python benchmarks/dynamo/huggingface.py \
-    --performance --export-aot-inductor --inference --devices "$device" "$dtype_arg" \
-    --output "$TEST_REPORTS_DIR/aot_inductor_huggingface_${dtype}_inference_${device}_performance.csv" || true
-  PYTHONPATH="$(pwd)"/torchbench python benchmarks/dynamo/huggingface.py \
-    --accuracy --export-aot-inductor --inference --devices "$device" "$dtype_arg" \
-    --output "$TEST_REPORTS_DIR/aot_inductor_huggingface_${dtype}_inference_${device}_accuracy.csv" || true
 
   echo "Pytorch benchmark on mps device completed"
 }
@@ -391,8 +324,6 @@ elif [[ $TEST_CONFIG == *"perf_timm"* ]]; then
   test_timm_perf
 elif [[ $TEST_CONFIG == *"perf_smoketest"* ]]; then
   test_torchbench_smoketest "${SHARD_NUMBER}"
-elif [[ $TEST_CONFIG == *"aot_inductor_perf_smoketest"* ]]; then
-  test_aoti_torchbench_smoketest "${SHARD_NUMBER}"
 elif [[ $TEST_CONFIG == *"mps"* ]]; then
   test_python_mps
 elif [[ $NUM_TEST_SHARDS -gt 1 ]]; then

.ci/pytorch/multigpu-test.sh

@@ -26,7 +26,6 @@ if [[ "${SHARD_NUMBER:-2}" == "2" ]]; then
     time python test/run_test.py --verbose -i distributed/test_c10d_spawn_gloo
     time python test/run_test.py --verbose -i distributed/test_c10d_spawn_nccl
     time python test/run_test.py --verbose -i distributed/test_compute_comm_reordering
-    time python test/run_test.py --verbose -i distributed/test_aten_comm_compute_reordering
     time python test/run_test.py --verbose -i distributed/test_store
     time python test/run_test.py --verbose -i distributed/test_symmetric_memory
     time python test/run_test.py --verbose -i distributed/test_pg_wrapper
@@ -46,7 +45,6 @@ if [[ "${SHARD_NUMBER:-2}" == "2" ]]; then
     # DTensor tests
     time python test/run_test.py --verbose -i distributed/tensor/test_random_ops
     time python test/run_test.py --verbose -i distributed/tensor/test_dtensor_compile
-    time python test/run_test.py --verbose -i distributed/tensor/test_utils.py
 
     # DeviceMesh test
     time python test/run_test.py --verbose -i distributed/test_device_mesh