Add build segmentation flags and XPU index selection

Author: Godnight1006

API.md

@@ -13,6 +13,10 @@ Or you can use the library:
 torchruntime.install(["torch", "torchvision<0.20"])
 ```
 
+Optional flags:
+- `preview=True` to allow preview builds (e.g. ROCm 6.4, nightly builds, XPU test index)
+- `unsupported=False` to forbid EOL/unsupported builds (errors instead of installing older builds)
+
 On Windows CUDA, Linux ROCm (6.x+), and Linux XPU, this also installs the appropriate Triton package to enable `torch.compile` (`triton-windows`, `pytorch-triton-rocm`, or `pytorch-triton-xpu`).
 
 ## Test torch

tests/test_installer.py

@@ -35,19 +35,26 @@ def test_cuda_platform_windows_installs_triton(monkeypatch):
 def test_cuda_nightly_platform_linux(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Linux")
     packages = ["torch", "torchvision"]
-    result = get_install_commands("nightly/cu112", packages)
+    result = get_install_commands("nightly/cu112", packages, preview=True)
     expected_url = "https://download.pytorch.org/whl/nightly/cu112"
     assert result == [packages + ["--index-url", expected_url]]
 
 
 def test_cuda_nightly_platform_windows_installs_triton(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Windows")
     packages = ["torch", "torchvision"]
-    result = get_install_commands("nightly/cu112", packages)
+    result = get_install_commands("nightly/cu112", packages, preview=True)
     expected_url = "https://download.pytorch.org/whl/nightly/cu112"
     assert result == [packages + ["--index-url", expected_url], ["triton-windows"]]
 
 
+def test_cuda_nightly_platform_requires_preview(monkeypatch):
+    monkeypatch.setattr("torchruntime.installer.os_name", "Linux")
+    packages = ["torch", "torchvision"]
+    with pytest.raises(ValueError, match="preview"):
+        get_install_commands("nightly/cu112", packages, preview=False)
+
+
 def test_rocm_4_platform_does_not_install_triton(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Linux")
     packages = ["torch", "torchvision"]
@@ -71,25 +78,43 @@ def test_rocm_6_platform_linux_installs_triton(monkeypatch):
 def test_xpu_platform_windows_with_torch_only(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Windows")
     packages = ["torch"]
-    result = get_install_commands("xpu", packages)
-    expected_url = "https://download.pytorch.org/whl/test/xpu"
+    result = get_install_commands("xpu", packages, preview=False)
+    expected_url = "https://download.pytorch.org/whl/xpu"
+    assert result == [packages + ["--index-url", expected_url]]
+
+
+def test_xpu_platform_windows_with_torchvision(monkeypatch):
+    monkeypatch.setattr("torchruntime.installer.os_name", "Windows")
+    packages = ["torch", "torchvision"]
+    result = get_install_commands("xpu", packages, preview=False)
+    expected_url = "https://download.pytorch.org/whl/xpu"
     assert result == [packages + ["--index-url", expected_url]]
 
 
-def test_xpu_platform_windows_with_torchvision(monkeypatch, capsys):
+def test_xpu_platform_windows_preview(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Windows")
     packages = ["torch", "torchvision"]
-    result = get_install_commands("xpu", packages)
-    expected_url = "https://download.pytorch.org/whl/nightly/xpu"
+    result = get_install_commands("xpu", packages, preview=True)
+    expected_url = "https://download.pytorch.org/whl/test/xpu"
     assert result == [packages + ["--index-url", expected_url]]
-    captured = capsys.readouterr()
-    assert "[WARNING]" in captured.out
 
 
 def test_xpu_platform_linux(monkeypatch):
     monkeypatch.setattr("torchruntime.installer.os_name", "Linux")
     packages = ["torch", "torchvision"]
-    result = get_install_commands("xpu", packages)
+    result = get_install_commands("xpu", packages, preview=False)
+    expected_url = "https://download.pytorch.org/whl/xpu"
+    triton_index_url = "https://download.pytorch.org/whl"
+    assert result == [
+        packages + ["--index-url", expected_url],
+        ["pytorch-triton-xpu", "--index-url", triton_index_url],
+    ]
+
+
+def test_xpu_platform_linux_preview(monkeypatch):
+    monkeypatch.setattr("torchruntime.installer.os_name", "Linux")
+    packages = ["torch", "torchvision"]
+    result = get_install_commands("xpu", packages, preview=True)
     expected_url = "https://download.pytorch.org/whl/test/xpu"
     triton_index_url = "https://download.pytorch.org/whl"
     assert result == [

tests/test_platform_detection.py

@@ -36,7 +36,7 @@ def test_amd_gpu_navi4_linux(monkeypatch):
         with pytest.raises(NotImplementedError):
             get_torch_platform(gpu_infos)
     else:
-        assert get_torch_platform(gpu_infos) == "rocm6.4"
+        assert get_torch_platform(gpu_infos) == "rocm6.2"
 
 
 def test_amd_gpu_navi3_linux(monkeypatch, capsys):
@@ -89,6 +89,14 @@ def test_amd_gpu_ellesmere_linux(monkeypatch):
     assert get_torch_platform(gpu_infos) == "rocm4.2"
 
 
+def test_amd_gpu_ellesmere_linux_unsupported_false_raises(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Linux")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "x86_64")
+    gpu_infos = [GPU(AMD, "AMD", 0x1234, "Ellesmere", True)]
+    with pytest.raises(ValueError, match="End-of-Life"):
+        get_torch_platform(gpu_infos, unsupported=False)
+
+
 def test_amd_gpu_unsupported_linux(monkeypatch, capsys):
     monkeypatch.setattr("torchruntime.platform_detection.os_name", "Linux")
     monkeypatch.setattr("torchruntime.platform_detection.arch", "x86_64")

tests/test_segmentation.py

@@ -0,0 +1,72 @@
+import pytest
+from torchruntime.device_db import GPU
+from torchruntime.platform_detection import AMD, INTEL, NVIDIA, get_torch_platform, py_version
+
+
+def test_preview_rocm_6_4_selection(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Linux")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "x86_64")
+    gpu_infos = [GPU(AMD, "AMD", 0x1234, "Navi 41", True)]
+
+    if py_version < (3, 9):
+        pytest.skip("Navi 4 requires Python 3.9+")
+
+    # Default: preview=False -> rocm6.2
+    assert get_torch_platform(gpu_infos) == "rocm6.2"
+    assert get_torch_platform(gpu_infos, preview=False) == "rocm6.2"
+
+    # preview=True -> rocm6.4
+    assert get_torch_platform(gpu_infos, preview=True) == "rocm6.4"
+
+
+def test_eol_cu118_selection(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Windows")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "amd64")
+    # Kepler architecture (e.g. GTX 780)
+    gpu_infos = [GPU(NVIDIA, "NVIDIA", "1004", "GK110 [GeForce GTX 780]", True)]
+
+    # Default: unsupported=True -> cu118
+    assert get_torch_platform(gpu_infos) == "cu118"
+    assert get_torch_platform(gpu_infos, unsupported=True) == "cu118"
+
+    # unsupported=False -> raises ValueError
+    with pytest.raises(ValueError, match="considered End-of-Life"):
+        get_torch_platform(gpu_infos, unsupported=False)
+
+
+def test_eol_rocm42_selection(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Linux")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "x86_64")
+    # Ellesmere (e.g. RX 580)
+    gpu_infos = [GPU(AMD, "AMD", "67df", "Ellesmere [Radeon RX 470/480/570/570X/580/580X/590]", True)]
+
+    # Default: unsupported=True -> rocm4.2
+    assert get_torch_platform(gpu_infos) == "rocm4.2"
+
+    # unsupported=False -> raises ValueError
+    with pytest.raises(ValueError, match="considered End-of-Life"):
+        get_torch_platform(gpu_infos, unsupported=False)
+
+
+def test_eol_directml_selection(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Windows")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "amd64")
+    gpu_infos = [GPU(AMD, "AMD", 0x1234, "Radeon", True)]
+
+    assert get_torch_platform(gpu_infos) == "directml"
+
+    with pytest.raises(ValueError, match="considered End-of-Life"):
+        get_torch_platform(gpu_infos, unsupported=False)
+
+
+def test_eol_ipex_selection(monkeypatch):
+    monkeypatch.setattr("torchruntime.platform_detection.os_name", "Linux")
+    monkeypatch.setattr("torchruntime.platform_detection.arch", "x86_64")
+    monkeypatch.setattr("torchruntime.platform_detection.py_version", (3, 8))
+    gpu_infos = [GPU(INTEL, "Intel", 0x1234, "Iris", True)]
+
+    assert get_torch_platform(gpu_infos) == "ipex"
+
+    # unsupported=False -> raises ValueError
+    with pytest.raises(ValueError, match="considered End-of-Life"):
+        get_torch_platform(gpu_infos, unsupported=False)

torchruntime/__main__.py

@@ -16,6 +16,8 @@ def print_usage(entry_command: str):
 Examples:
     {entry_command} install
     {entry_command} install --uv
+    {entry_command} install --preview
+    {entry_command} install --no-unsupported
     {entry_command} install torch==2.2.0 torchvision==0.17.0
     {entry_command} install --uv torch>=2.0.0 torchaudio
     {entry_command} install torch==2.1.* torchvision>=0.16.0 torchaudio==2.1.0
@@ -35,6 +37,8 @@ def print_usage(entry_command: str):
 
 Options:
     --uv               Use uv instead of pip for installation
+    --preview          Allow preview builds (e.g. ROCm 6.4)
+    --no-unsupported   Forbid EOL/unsupported builds (e.g. CUDA 11.8)
 
 Version specification formats (follows pip format):
     package==2.1.0     Exact version
@@ -63,14 +67,20 @@ def main():
     if command == "install":
         args = sys.argv[2:] if len(sys.argv) > 2 else []
         use_uv = "--uv" in args
-        # Remove --uv from args to get package list
-        package_versions = [arg for arg in args if arg != "--uv"] if args else None
-        install(package_versions, use_uv=use_uv)
+        preview = "--preview" in args
+        unsupported = "--no-unsupported" not in args
+        # Remove flags from args to get package list
+        package_versions = [arg for arg in args if arg not in ("--uv", "--preview", "--no-unsupported")] if args else None
+        install(package_versions, use_uv=use_uv, preview=preview, unsupported=unsupported)
     elif command == "test":
         subcommand = sys.argv[2] if len(sys.argv) > 2 else "all"
         test(subcommand)
     elif command == "info":
-        info()
+        args = sys.argv[2:] if len(sys.argv) > 2 else []
+        preview = "--preview" in args
+        unsupported = "--no-unsupported" not in args
+        from .utils import info
+        info(preview=preview, unsupported=unsupported)
     else:
         print(f"Unknown command: {command}")
         entry_path = sys.argv[0]

torchruntime/installer.py

@@ -2,20 +2,17 @@
 import sys
 import platform
 import subprocess
-
-from .consts import CONTACT_LINK
 from .device_db import get_gpus
 from .platform_detection import get_torch_platform
 
 os_name = platform.system()
 
-PIP_PREFIX = [sys.executable, "-m", "pip", "install"]
 CUDA_REGEX = re.compile(r"^(nightly/)?cu\d+$")
 ROCM_REGEX = re.compile(r"^(nightly/)?rocm\d+\.\d+$")
 ROCM_VERSION_REGEX = re.compile(r"^(?:nightly/)?rocm(?P<major>\d+)\.(?P<minor>\d+)$")
 
 
-def get_install_commands(torch_platform, packages):
+def get_install_commands(torch_platform, packages, preview=False):
     """
     Generates pip installation commands for PyTorch and related packages based on the specified platform.
 
@@ -30,6 +27,7 @@ def get_install_commands(torch_platform, packages):
         packages (list of str): List of package names (and optionally versions in pip format). Examples:
             - ["torch", "torchvision"]
             - ["torch>=2.0", "torchaudio==0.16.0"]
+        preview (bool): If True, allow preview/nightly builds. Defaults to False.
 
     Returns:
         list of list of str: Each sublist contains a pip install command (excluding the `pip install` prefix).
@@ -41,7 +39,7 @@ def get_install_commands(torch_platform, packages):
         ValueError: If an unsupported platform is provided.
 
     Notes:
-        - For "xpu" on Windows, if torchvision or torchaudio are included, the function switches to nightly builds.
+        - For "xpu", if preview is True, the function installs from the test index.
         - For "directml", the "torch-directml" package is returned as part of the installation commands.
         - For "ipex", the "intel-extension-for-pytorch" package is returned as part of the installation commands.
         - For Windows CUDA, the function also installs "triton-windows" (for torch.compile and Triton kernels).
@@ -54,6 +52,9 @@ def get_install_commands(torch_platform, packages):
     if torch_platform == "cpu":
         return [packages]
 
+    if torch_platform.startswith("nightly/") and not preview:
+        raise ValueError("preview=True is required for nightly builds")
+
     if CUDA_REGEX.match(torch_platform) or ROCM_REGEX.match(torch_platform):
         index_url = f"https://download.pytorch.org/whl/{torch_platform}"
         cmds = [packages + ["--index-url", index_url]]
@@ -69,15 +70,11 @@ def get_install_commands(torch_platform, packages):
         return cmds
 
     if torch_platform == "xpu":
-        if os_name == "Windows" and ("torchvision" in packages or "torchaudio" in packages):
-            print(
-                f"[WARNING] The preview build of 'xpu' on Windows currently only supports torch, not torchvision/torchaudio. "
-                f"torchruntime will instead use the nightly build, to get the 'xpu' version of torchaudio and torchvision as well. "
-                f"Please contact torchruntime if this is no longer accurate: {CONTACT_LINK}"
-            )
-            index_url = f"https://download.pytorch.org/whl/nightly/{torch_platform}"
-        else:
-            index_url = f"https://download.pytorch.org/whl/test/{torch_platform}"
+        index_url = (
+            f"https://download.pytorch.org/whl/test/{torch_platform}"
+            if preview
+            else f"https://download.pytorch.org/whl/{torch_platform}"
+        )
 
         cmds = [packages + ["--index-url", index_url]]
         if os_name == "Linux":
@@ -108,14 +105,16 @@ def run_commands(cmds):
         subprocess.run(cmd)
 
 
-def install(packages=[], use_uv=False):
+def install(packages=[], use_uv=False, preview=False, unsupported=True):
     """
     packages: a list of strings with package names (and optionally their versions in pip-format). e.g. ["torch", "torchvision"] or ["torch>=2.0", "torchaudio==0.16.0"]. Defaults to ["torch", "torchvision", "torchaudio"].
     use_uv: bool, whether to use uv for installation. Defaults to False.
+    preview: bool, whether to allow preview/nightly builds. Defaults to False.
+    unsupported: bool, whether to allow EOL/unsupported builds. Defaults to True.
     """
 
     gpu_infos = get_gpus()
-    torch_platform = get_torch_platform(gpu_infos, packages=packages)
-    cmds = get_install_commands(torch_platform, packages)
+    torch_platform = get_torch_platform(gpu_infos, packages=packages, preview=preview, unsupported=unsupported)
+    cmds = get_install_commands(torch_platform, packages, preview=preview)
     cmds = get_pip_commands(cmds, use_uv=use_uv)
     run_commands(cmds)

torchruntime/platform_detection.py

@@ -51,13 +51,15 @@ def _packages_require_cuda_12_4(packages):
     return False
 
 
-def get_torch_platform(gpu_infos, packages=[]):
+def get_torch_platform(gpu_infos, packages=[], preview=False, unsupported=True):
     """
     Determine the appropriate PyTorch platform to use based on the system architecture, OS, and GPU information.
 
     Args:
         gpu_infos (list of `torchruntime.device_db.GPU` instances)
         packages (list of str): Optional list of torch/torchvision/torchaudio requirement strings.
+        preview (bool): If True, allow preview/nightly builds (e.g. rocm6.4). Defaults to False.
+        unsupported (bool): If False, forbid EOL/unsupported builds (e.g. cu118). Defaults to True.
 
     Returns:
         str: A string representing the platform to use. Possible values:
@@ -70,6 +72,7 @@ def get_torch_platform(gpu_infos, packages=[]):
 
     Raises:
         NotImplementedError: For unsupported architectures, OS-GPU combinations, or multiple GPU vendors.
+        ValueError: If unsupported=False and only an EOL build is available.
         Warning: Outputs warnings for deprecated Python versions or fallback configurations.
     """
 
@@ -95,12 +98,23 @@ def get_torch_platform(gpu_infos, packages=[]):
             integrated_devices.append(device)
 
     if discrete_devices:
-        return _get_platform_for_discrete(discrete_devices, packages=packages)
+        platform = _get_platform_for_discrete(discrete_devices, packages=packages, preview=preview)
+    else:
+        platform = _get_platform_for_integrated(integrated_devices, preview=preview)
 
-    return _get_platform_for_integrated(integrated_devices)
+    # Segmentation Logic
+    EOL_PLATFORMS = {"cu118", "directml", "ipex", "rocm5.7", "rocm5.5", "rocm5.2", "rocm4.2"}
 
+    if not unsupported and platform in EOL_PLATFORMS:
+        raise ValueError(
+            f"The recommended platform '{platform}' is considered End-of-Life (EOL) and is forbidden because 'unsupported' is set to False. "
+            f"Please use a more recent GPU or set unsupported=True to allow this installation."
+        )
+
+    return platform
 
-def _get_platform_for_discrete(gpu_infos, packages=None):
+
+def _get_platform_for_discrete(gpu_infos, packages=None, preview=False):
     vendor_ids = set(gpu.vendor_id for gpu in gpu_infos)
 
     if len(vendor_ids) > 1:
@@ -124,7 +138,7 @@ def _get_platform_for_discrete(gpu_infos, packages=None):
                     raise NotImplementedError(
                         f"Torch does not support Navi 4x series of GPUs on Python 3.8. Please switch to a newer Python version to use the latest version of torch!"
                     )
-                return "rocm6.4"
+                return "rocm6.4" if preview else "rocm6.2"
             if any(device_name.startswith("Navi") for device_name in device_names) and any(
                 device_name.startswith("Vega 2") for device_name in device_names
             ):  # lowest-common denominator is rocm5.7, which works with both Navi and Vega 20
@@ -201,7 +215,7 @@ def _get_platform_for_discrete(gpu_infos, packages=None):
     return "cpu"
 
 
-def _get_platform_for_integrated(gpu_infos):
+def _get_platform_for_integrated(gpu_infos, preview=False):
     gpu = gpu_infos[0]
 
     if os_name == "Windows":
@@ -235,3 +249,4 @@ def _get_platform_for_integrated(gpu_infos):
         )
 
     return "cpu"
+