update packaging to include triton

packaging/runtime_builder.py

@@ -78,7 +78,7 @@ def main() -> int:
     filtered.write_text(
         "\n".join(
             line for line in requirements.read_text(encoding="utf-8").splitlines()
-            if line.strip().lower() not in {"torch", "torch==", "torchvision"}
+            if line.strip().lower() not in {"torch", "torch==", "torchvision", "triton", "triton-windows"}
         ) + "\n",
         encoding="utf-8",
     )

packaging/runtime_setup.py

@@ -86,47 +86,141 @@ def _log_path(root: Path) -> Path:
     raise OSError("Could not create a writable runtime setup log.")
 
 
+def get_triton_specifier(torch_version: str) -> str:
+    """Return the compatible triton-windows requirement for a given PyTorch version.
+
+    The official triton-windows wheels correspond to PyTorch releases:
+    - PyTorch 2.5.x -> triton-windows 3.1.x
+    - PyTorch 2.6.x -> triton-windows 3.2.x
+    - PyTorch 2.7.x -> triton-windows 3.3.x
+    - PyTorch 2.8.x -> triton-windows 3.4.x
+    - PyTorch 2.9.x -> triton-windows 3.5.x
+    - PyTorch 2.10.x -> triton-windows 3.6.x
+
+    Args:
+        torch_version: PyTorch version string (e.g. '2.5.1').
+
+    Returns:
+        A pip-compatible requirement specifier for triton-windows.
+    """
+    match = re.match(r"^(\d+)\.(\d+)", torch_version)
+    if not match:
+        return "triton-windows"
+    major, minor = int(match.group(1)), int(match.group(2))
+    if major == 2 and minor >= 5:
+        triton_minor = minor - 4
+        return f"triton-windows>=3.{triton_minor}.0,<3.{triton_minor + 1}.0"
+    if major == 2 and minor == 4:
+        return "triton-windows<3.1.0"
+    return "triton-windows"
+
+
+def _run_pip(python_executable: str, args: list[str], environment: dict[str, str]) -> None:
+    """Execute a pip command streaming output to stdout and the setup log.
+
+    Args:
+        python_executable: Path to the Python executable.
+        args: Command-line arguments passed to pip.
+        environment: Environment variable mapping.
+
+    Raises:
+        subprocess.CalledProcessError: If the pip command exits with non-zero status.
+    """
+    command = [python_executable, "-m", "pip", *args]
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        env=environment,
+        bufsize=1,
+    )
+    assert process.stdout is not None
+    for line in process.stdout:
+        print(line, end="", file=sys.stdout, flush=True)
+        if _SETUP_LOG is not None:
+            print(line, end="", file=_SETUP_LOG, flush=True)
+    return_code = process.wait()
+    if return_code:
+        raise subprocess.CalledProcessError(return_code, command)
+
+
 def install_runtime(python_executable: str, choice: RuntimeChoice) -> None:
+    """Install hardware-specific PyTorch and optional Triton runtime wheels.
+
+    Args:
+        python_executable: Path to the Python executable.
+        choice: Selected runtime choice.
+    """
     environment = os.environ.copy()
     environment["PYTHONNOUSERSITE"] = "1"
     environment["PATH"] = os.pathsep.join(
         item for item in environment.get("PATH", "").split(os.pathsep)
         if "mingw" not in item.lower() and "scoop" not in item.lower()
     )
-    process = subprocess.Popen(
+    _run_pip(
+        python_executable,
         [
-            python_executable,
-            "-m",
-            "pip",
             "install",
             "--no-warn-script-location",
             f"torch=={TORCH_VERSION}",
             "--index-url",
             TORCH_INDEXES[choice.profile],
         ],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        env=environment,
-        bufsize=1,
+        environment,
     )
-    assert process.stdout is not None
-    for line in process.stdout:
-        print(line, end="", file=sys.stdout, flush=True)
-        print(line, end="", file=_SETUP_LOG, flush=True)
-    return_code = process.wait()
-    if return_code:
-        raise subprocess.CalledProcessError(return_code, process.args)
+
+    if platform.system() == "Windows" and choice.profile != "cpu":
+        triton_spec = get_triton_specifier(TORCH_VERSION)
+        msg = f"Installing {triton_spec} for Windows CUDA runtime...\n"
+        print(msg, end="", file=sys.stdout, flush=True)
+        if _SETUP_LOG is not None:
+            print(msg, end="", file=_SETUP_LOG, flush=True)
+        try:
+            _run_pip(
+                python_executable,
+                [
+                    "install",
+                    "--no-warn-script-location",
+                    triton_spec,
+                ],
+                environment,
+            )
+        except Exception as exc:
+            warning = f"Warning: Failed to install {triton_spec}: {exc!r}. Running without Triton.\n"
+            print(warning, end="", file=sys.stdout, flush=True)
+            if _SETUP_LOG is not None:
+                print(warning, end="", file=_SETUP_LOG, flush=True)
 
 
 def ensure_runtime(python_executable: str, root: Path) -> RuntimeChoice:
+    """Ensure the expected runtime wheels are installed and verified.
+
+    Args:
+        python_executable: Path to the Python executable.
+        root: Root directory containing PROFILE_FILE.
+
+    Returns:
+        The selected RuntimeChoice.
+    """
     choice = choose_runtime()
     marker = root / PROFILE_FILE
     if marker.exists() and marker.read_text(encoding="utf-8").strip() == choice.profile:
+        if platform.system() == "Windows" and choice.profile != "cpu":
+            triton_check = subprocess.run(
+                [python_executable, "-c", "import triton"],
+                cwd=root,
+                capture_output=True,
+                check=False,
+            )
+            if triton_check.returncode != 0:
+                install_runtime(python_executable, choice)
         return choice
+
     install_runtime(python_executable, choice)
     verification = subprocess.run(
         [python_executable, "-c", "import torch; print(torch.__version__)"],
+        cwd=root,
         capture_output=True,
         text=True,
         check=False,
@@ -140,6 +234,23 @@ def ensure_runtime(python_executable: str, root: Path) -> RuntimeChoice:
     )
     if verification.returncode != 0:
         raise RuntimeError(f"Torch verification failed: {verification.stderr.strip()}")
+
+    if platform.system() == "Windows" and choice.profile != "cpu":
+        triton_ver = subprocess.run(
+            [python_executable, "-c", "import triton; print(triton.__version__)"],
+            cwd=root,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        print(
+            f"Triton verification exit code: {triton_ver.returncode}\n"
+            f"Triton verification output: {triton_ver.stdout.strip()}\n"
+            f"Triton verification error: {triton_ver.stderr.strip()}",
+            file=_SETUP_LOG,
+            flush=True,
+        )
+
     marker.write_text(choice.profile, encoding="utf-8")
     return choice
 

requirements.txt

@@ -12,4 +12,4 @@ Pygments
 pyqtgraph
 psutil
 datasets
-cryptography
+cryptography

tests/test_runtime_setup.py

@@ -3,7 +3,7 @@
 import unittest
 from unittest.mock import patch
 
-from packaging.runtime_setup import choose_runtime
+from packaging.runtime_setup import RuntimeChoice, choose_runtime, get_triton_specifier, install_runtime
 
 
 class RuntimeSelectionTests(unittest.TestCase):
@@ -21,5 +21,74 @@ def test_macos_selects_cpu(self) -> None:
         self.assertEqual(choose_runtime("Darwin", 600).profile, "cpu")
 
 
+class TritonSpecifierTests(unittest.TestCase):
+    def test_torch_2_5_maps_to_triton_3_1(self) -> None:
+        self.assertEqual(get_triton_specifier("2.5.1"), "triton-windows>=3.1.0,<3.2.0")
+
+    def test_torch_2_6_maps_to_triton_3_2(self) -> None:
+        self.assertEqual(get_triton_specifier("2.6.0"), "triton-windows>=3.2.0,<3.3.0")
+
+    def test_torch_2_7_maps_to_triton_3_3(self) -> None:
+        self.assertEqual(get_triton_specifier("2.7.0"), "triton-windows>=3.3.0,<3.4.0")
+
+    def test_torch_2_10_maps_to_triton_3_6(self) -> None:
+        self.assertEqual(get_triton_specifier("2.10.0"), "triton-windows>=3.6.0,<3.7.0")
+
+    def test_torch_2_4_maps_to_triton_under_3_1(self) -> None:
+        self.assertEqual(get_triton_specifier("2.4.1"), "triton-windows<3.1.0")
+
+    def test_fallback_for_unrecognized_version(self) -> None:
+        self.assertEqual(get_triton_specifier("custom-build"), "triton-windows")
+
+
+class RuntimeInstallationTests(unittest.TestCase):
+    @patch("packaging.runtime_setup.platform.system", return_value="Windows")
+    @patch("packaging.runtime_setup._run_pip")
+    def test_windows_cuda_installs_both_torch_and_triton(self, mock_run_pip, _mock_platform) -> None:
+        choice = RuntimeChoice(profile="cu124", reason="GPU supported")
+        install_runtime("python.exe", choice)
+
+        self.assertEqual(mock_run_pip.call_count, 2)
+        torch_call = mock_run_pip.call_args_list[0][0]
+        self.assertIn("torch==2.5.1", torch_call[1])
+        self.assertIn("https://download.pytorch.org/whl/cu124", torch_call[1])
+
+        triton_call = mock_run_pip.call_args_list[1][0]
+        self.assertIn("triton-windows>=3.1.0,<3.2.0", triton_call[1])
+
+    @patch("packaging.runtime_setup.platform.system", return_value="Windows")
+    @patch("packaging.runtime_setup._run_pip")
+    def test_windows_cpu_skips_triton(self, mock_run_pip, _mock_platform) -> None:
+        choice = RuntimeChoice(profile="cpu", reason="No GPU detected")
+        install_runtime("python.exe", choice)
+
+        self.assertEqual(mock_run_pip.call_count, 1)
+        torch_call = mock_run_pip.call_args_list[0][0]
+        self.assertIn("torch==2.5.1", torch_call[1])
+        self.assertIn("https://download.pytorch.org/whl/cpu", torch_call[1])
+
+    @patch("packaging.runtime_setup.platform.system", return_value="Linux")
+    @patch("packaging.runtime_setup._run_pip")
+    def test_linux_skips_triton_windows(self, mock_run_pip, _mock_platform) -> None:
+        choice = RuntimeChoice(profile="cu124", reason="GPU supported")
+        install_runtime("python3", choice)
+
+        self.assertEqual(mock_run_pip.call_count, 1)
+        torch_call = mock_run_pip.call_args_list[0][0]
+        self.assertIn("torch==2.5.1", torch_call[1])
+
+    @patch("packaging.runtime_setup.platform.system", return_value="Windows")
+    @patch("packaging.runtime_setup._run_pip")
+    def test_triton_install_failure_is_non_fatal(self, mock_run_pip, _mock_platform) -> None:
+        choice = RuntimeChoice(profile="cu124", reason="GPU supported")
+        # First call (torch) succeeds, second call (triton) raises
+        mock_run_pip.side_effect = [None, RuntimeError("pip network failure")]
+
+        # install_runtime should not raise
+        install_runtime("python.exe", choice)
+        self.assertEqual(mock_run_pip.call_count, 2)
+
+
 if __name__ == "__main__":
     unittest.main()
+