Use CPU torch for Linux release packaging

Author: danias

.github/workflows/release.yml

@@ -64,6 +64,36 @@ jobs:
         shell: bash
         run: python scripts/real_backend_smoke.py --binary bitloops-embeddings
 
+      - name: Switch Linux packaging environment to CPU-only torch
+        if: runner.os == 'Linux'
+        shell: bash
+        run: |
+          TORCH_VERSION="$(python - <<'PY'
+          import importlib.metadata
+          print(importlib.metadata.version("torch"))
+          PY
+          )"
+
+          mapfile -t GPU_PACKAGES < <(python - <<'PY'
+          import subprocess
+          import sys
+
+          output = subprocess.check_output([sys.executable, "-m", "pip", "freeze"], text=True)
+          for line in output.splitlines():
+              name = line.split("==", 1)[0]
+              lower = name.lower()
+              if lower.startswith("nvidia-") or lower == "triton":
+                  print(name)
+          PY
+          )
+
+          if ((${#GPU_PACKAGES[@]} > 0)); then
+            python -m pip uninstall -y "${GPU_PACKAGES[@]}"
+          fi
+
+          python -m pip uninstall -y torch
+          python -m pip install "torch==${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu
+
       - name: Import Apple signing certificate
         if: runner.os == 'macOS'
         env:

scripts/package_release.py

@@ -29,34 +29,13 @@
     "IPython",
     "jedi",
     "matplotlib",
-    "nvidia",
     "notebook",
     "pandas",
     "tensorboard",
     "tensorflow",
-    "triton",
     "torchvision",
     "torchaudio",
 ]
-LINUX_GPU_LIBRARY_PREFIXES = (
-    "libc10_cuda",
-    "libcudart",
-    "libcudnn",
-    "libcufft",
-    "libcupti",
-    "libcublas",
-    "libcublasLt",
-    "libcurand",
-    "libcusolver",
-    "libcusparse",
-    "libcusparseLt",
-    "libnccl",
-    "libnvJitLink",
-    "libnvrtc",
-    "libnvToolsExt",
-    "libnpp",
-    "libtorch_cuda",
-)
 
 
 def main() -> None:
@@ -122,8 +101,6 @@ def build_release(*, version: str, target: str, archive_dir: Path) -> tuple[Path
     )
 
     shutil.copytree(pyinstaller_dist / PACKAGE_NAME, bundle_dir)
-    if target.endswith("unknown-linux-gnu"):
-        prune_linux_gpu_payload(bundle_dir)
     shutil.copy2(ROOT / "README.md", staging_dir / "README.md")
     shutil.copy2(ROOT / "LICENSE", staging_dir / "LICENSE")
 
@@ -232,27 +209,5 @@ def write_github_outputs(output_path: Path, outputs: dict[str, str]) -> None:
             file_handle.write(f"{key}={value}\n")
 
 
-def prune_linux_gpu_payload(bundle_dir: Path) -> None:
-    removable_directories = []
-    for pattern in ("**/nvidia", "**/triton"):
-        removable_directories.extend(
-            path for path in bundle_dir.glob(pattern) if path.is_dir()
-        )
-
-    for path in sorted(removable_directories, reverse=True):
-        shutil.rmtree(path, ignore_errors=True)
-
-    for path in bundle_dir.rglob("*"):
-        if not path.is_file():
-            continue
-
-        name = path.name
-        if name.endswith(".pyc"):
-            continue
-
-        if name.startswith(LINUX_GPU_LIBRARY_PREFIXES):
-            path.unlink(missing_ok=True)
-
-
 if __name__ == "__main__":
     main()