diff --git a/.github/workflows/staging-9858-macos-15.yml b/.github/workflows/staging-9858-macos-15.yml new file mode 100644 index 00000000000..354643bd9ac --- /dev/null +++ b/.github/workflows/staging-9858-macos-15.yml @@ -0,0 +1,31 @@ +name: "staging-9858 macos-15" +on: + push: + branches: ["pr-9858-xplat-ci"] + paths-ignore: + - '**/*.md' + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: macos-15 + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py tests/python/test_install_python_stack.py tests/studio/install/test_cuda_repair.py tests/studio/install/test_gpu_detection_followups.py tests/studio/install/test_pr5940_followups.py tests/studio/install/test_rocm_arch_table_parity.py tests/studio/install/test_rocm_support.py tests/studio/install/test_windows_torch_flavor_invariant.py tests/studio/test_xpu_triton_swap.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc" diff --git a/.github/workflows/staging-9858-ubuntu-latest.yml b/.github/workflows/staging-9858-ubuntu-latest.yml new file mode 100644 index 00000000000..896c1aa4af5 --- /dev/null +++ b/.github/workflows/staging-9858-ubuntu-latest.yml @@ -0,0 +1,31 @@ +name: "staging-9858 ubuntu-latest" +on: + push: + branches: ["pr-9858-xplat-ci"] + paths-ignore: + - '**/*.md' + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: ubuntu-latest + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py tests/python/test_install_python_stack.py tests/studio/install/test_cuda_repair.py tests/studio/install/test_gpu_detection_followups.py tests/studio/install/test_pr5940_followups.py tests/studio/install/test_rocm_arch_table_parity.py tests/studio/install/test_rocm_support.py tests/studio/install/test_windows_torch_flavor_invariant.py tests/studio/test_xpu_triton_swap.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc" diff --git a/.github/workflows/staging-9858-windows-latest.yml b/.github/workflows/staging-9858-windows-latest.yml new file mode 100644 index 00000000000..7ff1296aa12 --- /dev/null +++ b/.github/workflows/staging-9858-windows-latest.yml @@ -0,0 +1,31 @@ +name: "staging-9858 windows-latest" +on: + push: + branches: ["pr-9858-xplat-ci"] + paths-ignore: + - '**/*.md' + workflow_dispatch: +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read +defaults: + run: + shell: bash +jobs: + test: + runs-on: windows-latest + timeout-minutes: 30 + env: + UNSLOTH_COMPILE_DISABLE: '1' + UNSLOTH_IS_PRESENT: '1' + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - run: python -m pip install -r studio/backend/requirements/studio.txt pytest pytest-asyncio pytest-timeout python-multipart + - run: rc=0; PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py tests/python/test_install_python_stack.py tests/studio/install/test_cuda_repair.py tests/studio/install/test_gpu_detection_followups.py tests/studio/install/test_pr5940_followups.py tests/studio/install/test_rocm_arch_table_parity.py tests/studio/install/test_rocm_support.py tests/studio/install/test_windows_torch_flavor_invariant.py tests/studio/test_xpu_triton_swap.py -q --tb=short -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' > pytest_out.txt 2>&1 || rc=$?; cat pytest_out.txt; if [ "$rc" = "5" ]; then echo "no tests ran (deps absent on this runner)"; rc=0; fi; if [ "$rc" = "2" ] && grep -qE "No module named .(torch|unsloth_zoo|transformers)." pytest_out.txt && ! grep -qE "^(FAILED|ERROR) " pytest_out.txt; then echo "collection needs a dep this runner does not ship"; rc=0; fi; if [ "$rc" = "1" ] && ! grep -qE "^ERROR " pytest_out.txt; then tot=$(grep -cE "^FAILED " pytest_out.txt); dep=$(grep -cE "^FAILED .* No module named .(torch|unsloth_zoo|transformers).$" pytest_out.txt); if [ "$tot" -gt 0 ] && [ "$tot" = "$dep" ]; then echo "only tests needing a dep this runner does not ship failed"; rc=0; fi; fi; exit "$rc" diff --git a/install.sh b/install.sh index 5c07e2e6cc7..881cabf9588 100755 --- a/install.sh +++ b/install.sh @@ -4559,6 +4559,15 @@ while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_inde done _torch_index_leaf="${_torch_index_leaf##*/}" _torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]') +# Whether the caller had already STATED a backend before the assignment below overwrites it. +# setup.sh documents UNSLOTH_TORCH_BACKEND=cpu as the way to keep a deliberate CPU install, +# and on a GPU-less host the resolved value is cpu too, so without this the manifest cannot +# tell a stated choice from the automatic answer. +if [ -n "${UNSLOTH_TORCH_BACKEND:-}" ]; then + _torch_backend_was_stated=true +else + _torch_backend_was_stated=false +fi case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; @@ -4568,6 +4577,15 @@ case "$_torch_index_leaf" in *) unset UNSLOTH_TORCH_BACKEND ;; esac +# Derived from the index this script RESOLVED, which on a GPU-less machine is "cpu" whether +# or not anyone asked. Without the marker every ordinary Linux CPU install is recorded as a +# deliberate choice, and a machine that later gains a GPU is never offered the repair. +if [ -n "${UNSLOTH_TORCH_BACKEND:-}" ] && [ "$_torch_backend_was_stated" != true ]; then + export UNSLOTH_TORCH_BACKEND_SOURCE="resolved" +else + unset UNSLOTH_TORCH_BACKEND_SOURCE +fi + # Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the # ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf # merely STARTING with "rocm" isn't force-repaired from the wrong path. diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d13dd696ebb..8c3edc71684 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10610,6 +10610,10 @@ def _probe_or_none(): # only when an axis is actually quantized. _tensor_quant_kv_unsupported_binaries: set[tuple[str, int]] = set() + # Binary dirs already reported by _warn_missing_windows_cuda_runtime. The env is rebuilt + # for every launch and every --list-devices probe, so one line per binary is enough. + _missing_cuda_runtime_warned: set[str] = set() + @classmethod def _binary_key(cls, binary: Optional[str]) -> Optional[tuple[str, int]]: """(path, mtime_ns); ns mtime re-probes a same-second binary swap.""" @@ -10733,6 +10737,47 @@ def _add(path: Path) -> None: _add(site_packages / "torch" / "lib") return out + @classmethod + def _warn_missing_windows_cuda_runtime(cls, binary_dir: str, path_dirs: list[str]) -> None: + """Say so when a CUDA llama-server has no cudart to load. Diagnostic only. + + The CUDA prebuilt links ``cudart64_*.dll`` / ``cublas64_*.dll`` and takes them + from the managed venv -- ``torch/lib`` or the ``nvidia/*`` wheels, per + _windows_pip_nvidia_dll_dirs. A 2.11.0+cpu torch ships neither, so ggml cannot + load its CUDA backend and ``llama-server.exe --list-devices`` prints + ``Available devices: (none)`` while UNSLOTH_PREBUILT_INFO.json still says + ``backend cuda`` (#8473, HF discussion 87). Today that is entirely silent. + + Changes nothing about the launch: the process still starts, still falls back to + CPU, and a custom build with the DLLs somewhere else is not second-guessed. + Never raises -- a diagnostic must not be able to stop a load. + """ + try: + if binary_dir in cls._missing_cuda_runtime_warned: + return + # Same identification _installed_ggml_backends uses: the official prebuilts are + # single-backend, so the ggml CUDA lib beside llama-server IS the build. + ggml_cuda = os.path.join(binary_dir, "ggml-cuda.dll") + if not os.path.isfile(ggml_cuda): + return + for directory in path_dirs: + try: + names = os.listdir(directory) + except OSError: + continue + if any(name.lower().startswith("cudart64_") for name in names): + return + cls._missing_cuda_runtime_warned.add(binary_dir) + logger.warning( + "llama.cpp is the CUDA build (%s) but no cudart64_*.dll was found on its " + "DLL search path. The CUDA ggml backend will not load and llama-server " + "will report no devices. This is what a CPU-only PyTorch in the managed " + "environment looks like; repair the installation to restore GPU support.", + ggml_cuda, + ) + except Exception as e: + logger.debug(f"CUDA runtime DLL diagnostic failed: {e}") + @staticmethod def _build_windows_path_dirs(binary_dir: str, prefix: str, cuda_path: str) -> list[str]: """Ordered PATH entries prepended so llama-server.exe resolves cudart / @@ -10771,6 +10816,13 @@ def _llama_server_env_for_binary( ) existing_path = env.get("PATH", "") env["PATH"] = ";".join(path_dirs) + ";" + existing_path + # Warn against the FULL search path, inherited entries included: a hand-installed CUDA + # toolkit puts cudart64_*.dll on PATH without the venv or CUDA_PATH knowing, and warning + # on the prepended directories alone told working custom setups to repair a fine install. + LlamaCppBackend._warn_missing_windows_cuda_runtime( + binary_dir, + path_dirs + [d for d in existing_path.split(";") if d], + ) # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile # kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches diff --git a/studio/backend/main.py b/studio/backend/main.py index 3c17f831ae2..0014930056f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1561,11 +1561,15 @@ def _hardware_snapshot() -> Optional[tuple[bool, Optional[str], Optional[str]]]: generation = _hw_module.DETECTION_GENERATION device = _hw_module.DEVICE chat_only = bool(_hw_module.CHAT_ONLY) - reason = getattr(_hw_module, "CHAT_ONLY_REASON", None) - # Inside the guarded read, with the reason it belongs to. Read after it, a forced - # re-detect starting in between would pair this reply's reason with a detail from - # a different pass, or with none at all. - detail = getattr(_hw_module, "CHAT_ONLY_DETAIL", None) + # Refreshed, not the frozen global: the three inventory-sensitive verdicts can change + # after startup (an eGPU attached, a driver that finished restarting). Reason and detail + # come back together, or a forced re-detect starting in between would pair this reply's + # reason with a detail from a different pass. + try: + reason, detail = _hw_module.current_chat_only_verdict() + except Exception: + reason = getattr(_hw_module, "CHAT_ONLY_REASON", None) + detail = getattr(_hw_module, "CHAT_ONLY_DETAIL", None) if ( device is not None and _hw_module.DETECTION_COMPLETE.is_set() @@ -2057,8 +2061,9 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]] logger.debug(f"Could not resolve gpu_ids support: {e}") llama_uses_vulkan = False gpu_ids_supported = True - # Preserve backend/index metadata from the visibility probe: a CPU training host can expose - # a Vulkan inference GPU, and the UI must label it Vulkan, not the top-level CPU backend. + # The spread also carries `physical_devices` and `mismatch`: GPUs the OS sees that this PyTorch + # cannot open (#8473). They stay their own fields, because `devices` below is the runtime-usable + # list that model fit budgets against and the training device picker pins from. gpu_info = { **visibility_info, "available": visibility_info.get("available", False), diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index b7935300f07..dd419ee57c5 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -942,3 +942,28 @@ def real_prequant_safe_globals(monkeypatch): monkeypatch.setattr(pq, "_SAFE_GLOBALS_REGISTERED", None) monkeypatch.setattr(pq, "_RESOLVED_SAFE_GLOBALS", set()) return resolver + + +@pytest.fixture(autouse = True) +def _no_carried_over_hardware_measurements(): + """Both hardware caches start empty for every test, as they do in a fresh process. + + The torch build snapshot and the physical GPU inventory are module globals with a + 60 second TTL, so one test's host -- a suite that makes `import torch` fail, say -- + would otherwise answer for every test that ran within a minute of it. Cleared + afterwards as well, so a test that warms one deliberately does not leak either. + """ + from utils.hardware import hardware as _hw + + def _clear(): + # Under the locks: a non-blocking read hands the refresh to a daemon thread that holds + # these while it writes, so clearing without waiting lets a previous test's REAL host land + # in the cache a moment later. Torch lock FIRST, then the inventory lock, because that is + # the order the background refresh takes them in. + with _hw._torch_build_snapshot_lock, _hw._physical_gpu_inventory_lock: + _hw._torch_build_snapshot_cache = None + _hw._physical_gpu_inventory_cache = None + + _clear() + yield + _clear() diff --git a/studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py b/studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py new file mode 100644 index 00000000000..189dc88497d --- /dev/null +++ b/studio/backend/tests/test_torch_cpu_build_on_nvidia_host.py @@ -0,0 +1,2440 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A host whose GPUs PyTorch cannot use must not be reported as having no GPUs. + +Every GPU field in utils/hardware is gated on ``torch.cuda.is_available()``, so a +managed environment that ends up with a CPU-only wheel goes silent about hardware +that is plainly still there. Two Windows users hit exactly that after an in-app +update resolved torch from PyPI (whose Windows default is ``2.11.0+cpu``) over a +``cu124`` wheel: Settings > System showed ``VRAM --`` and "No visible GPU" while +``nvidia-smi`` listed both RTX A4000s, ``UNSLOTH_PREBUILT_INFO.json`` still said +``backend cuda``, and models ran on CPU (#8473, HF discussion 87). + +What these pin: + * the nvidia-smi inventory is read even though ``DEVICE`` is CPU; + * ``devices`` stays EMPTY regardless. That list is what the model-fit estimate + budgets against and what the training device picker pins from, so a card torch + cannot open must never be merged into it. The inventory travels beside it, in + ``physical_devices`` / ``mismatch``; + * ``CHAT_ONLY_REASON`` stops being ``"no_gpu"``, because on this host that is + false and the advice it implies ("get a GPU") is not the fix; + * a CPU-only wheel and an accelerator wheel whose runtime will not start stay + distinct reasons -- one is repaired by reinstalling torch, the other by the + driver; + * a host that genuinely has no GPU is unaffected, and a probe that cannot answer + (no nvidia-smi, malformed CSV, a comma inside a device name) degrades to a + structured unavailable result rather than raising out of the endpoint. + +No AMD/Windows hardware exists here, so torch and nvidia-smi are faked in the shapes +the reports describe; the Windows adapter half is exercised through its own map. +""" + +from __future__ import annotations + +import inspect +import subprocess +import types +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import main +import utils.hardware as hardware_pkg +import utils.hardware.hardware as hw +from utils.hardware import nvidia + +# nvidia-smi rows for User A's box: two A4000s, the second carrying a comma in its name. +_TWO_A4000_ROWS = "\n".join( + [ + "0, NVIDIA RTX A4000, 16376", + "1, NVIDIA RTX A4000, Founders Edition, 16376", + ] +) + + +def _fake_torch(vendor: str): + """A fake ``torch``, in the vendor shapes the reports describe. + + "cpu" -- the PyPI Windows default an unconstrained update resolves to. + "cuda_dead" -- a cu124 wheel whose runtime refuses to initialise. + "cuda" -- a healthy CUDA wheel. + """ + torch = types.ModuleType("torch") + if vendor == "cpu": + torch.version = SimpleNamespace(hip = None, cuda = None) + torch.__version__ = "2.11.0+cpu" + available = False + elif vendor == "cuda_dead": + torch.version = SimpleNamespace(hip = None, cuda = "12.4") + torch.__version__ = "2.6.0+cu124" + available = False + else: + torch.version = SimpleNamespace(hip = None, cuda = "12.8") + torch.__version__ = "2.9.1+cu128" + available = True + torch.cuda = SimpleNamespace( + is_available = lambda: available, + device_count = lambda: 2 if available else 0, + ) + return torch + + +def _smi( + monkeypatch, + stdout: str, + *, + returncode: int = 0, +): + """Pin what nvidia-smi answers, at the subprocess boundary. + + The procfs count goes with it. A CLI that cannot answer falls back to + /proc/driver/nvidia/gpus, and this suite runs on a real NVIDIA host, so a test + simulating a machine with no cards has to simulate the kernel driver's absence too. + The fallback's own tests set it back. + """ + monkeypatch.setattr( + nvidia.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(returncode = returncode, stdout = stdout), + ) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + + +@pytest.fixture(autouse = True) +def _no_inherited_visibility_mask(monkeypatch): + """An emptied mask is a deliberate CPU pin and suppresses the whole report, so a + runner that exports one (a GPU-partitioning CI job) would silently void these.""" + for var in ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"): + monkeypatch.delenv(var, raising = False) + + +@pytest.fixture(autouse = True) +def _no_carried_over_torch_measurement(monkeypatch): + """The torch snapshot is cached with a TTL, so one test's fake host would answer + for the next one. Both caches start empty here, as they do in a fresh process.""" + monkeypatch.setattr(hw, "_torch_build_snapshot_cache", None) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + + +@pytest.fixture(autouse = True) +def _no_background_inventory_refresh(monkeypatch): + """Keep the non-blocking refresh from probing the REAL host mid-test. + + get_physical_gpu_inventory(block=False) hands a stale or cold cache to a daemon + thread, which is the whole point on a request path, but in a test it races the + assertions and can drop this machine's actual nvidia-smi output into the cache. + The tests that assert on the scheduling patch Thread themselves. + """ + + class _NoThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + pass + + monkeypatch.setattr(hw.threading, "Thread", _NoThread) + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", False) + + +@pytest.fixture +def cpu_torch_on_an_nvidia_host(monkeypatch): + """torch reports no accelerator; nvidia-smi lists two A4000s.""" + monkeypatch.setitem(__import__("sys").modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", None) + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + # Cached for 60s, so a test that did not clear it would read the real host's answer. + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + + +def test_the_physical_probe_runs_without_a_cuda_device(cpu_torch_on_an_nvidia_host): + inventory = hw.get_physical_gpu_inventory() + + assert inventory["available"] is True + assert inventory["sources"] == ["nvidia-smi"] + assert [device["name"] for device in inventory["devices"]] == [ + "NVIDIA RTX A4000", + "NVIDIA RTX A4000, Founders Edition", + ] + assert [device["memory_total_gb"] for device in inventory["devices"]] == [15.99, 15.99] + assert {device["vendor"] for device in inventory["devices"]} == {"nvidia"} + + +@pytest.mark.parametrize( + ("stdout", "returncode", "failure"), + [ + ("", 0, None), + ("", 9, None), + ("nonsense\n0, only two columns\n", 0, None), + ("", 0, FileNotFoundError("nvidia-smi")), + ("", 0, subprocess.TimeoutExpired("nvidia-smi", 10)), + ], +) +def test_a_probe_that_cannot_answer_returns_a_result_rather_than_raising( + monkeypatch, stdout, returncode, failure +): + if failure is None: + _smi(monkeypatch, stdout, returncode = returncode) + else: + + def _raise(*_args, **_kwargs): + raise failure + + monkeypatch.setattr(nvidia.subprocess, "run", _raise) + # This suite runs on a real NVIDIA host, so pin procfs empty: the case under test is a + # machine with no NVIDIA driver at all. + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + + result = nvidia.get_physical_gpu_inventory() + assert result["available"] is False + assert result["devices"] == [] + + inventory = hw.get_physical_gpu_inventory() + assert inventory["available"] is False + assert inventory["devices"] == [] + assert inventory["sources"] == [] + # "The driver answered and there are no cards" is not "no probe could answer". An absent + # nvidia-smi is the exception: it is the normal state of every AMD, Intel and CPU host. + _could_not_answer = returncode != 0 or ( + failure is not None and not isinstance(failure, FileNotFoundError) + ) + assert inventory["unknown"] is _could_not_answer + + +def test_the_windows_amd_adapters_are_inventoried_too(monkeypatch): + # No vendor CLI is guaranteed on Windows AMD, so the DirectX registry map is the source. + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + monkeypatch.setattr( + hw, + "_windows_live_adapter_names", + lambda: ["AMD Radeon RX 7900 XT", "AMD Radeon(TM) Graphics"], + ) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + { + 0x24CF5: { + "name": "AMD Radeon RX 7900 XT", + "dedicated_memory_bytes": 20 * 1024**3, + }, + # No dedicated-memory value: unknown capacity, not an empty card. + 0x14CF5: {"name": "AMD Radeon(TM) Graphics"}, + } + if vendor_id == hw._AMD_PCI_VENDOR_ID + else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + + assert inventory["sources"] == ["directx-registry"] + assert [(d["name"], d["memory_total_gb"]) for d in inventory["devices"]] == [ + ("AMD Radeon(TM) Graphics", None), + ("AMD Radeon RX 7900 XT", 20.0), + ] + + +def test_a_cpu_wheel_and_a_dead_cuda_wheel_are_different_reasons(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + assert hw.classify_torch_build() == "torch_cpu_build" + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + assert hw.classify_torch_build() is None + + +def test_an_untagged_wheel_is_only_a_cpu_build_when_it_names_no_runtime(monkeypatch): + import sys + + untagged = _fake_torch("cpu") + untagged.__version__ = "2.9.0" + monkeypatch.setitem(sys.modules, "torch", untagged) + assert hw.classify_torch_build() == "torch_cpu_build" + + conda = _fake_torch("cuda_dead") + conda.__version__ = "2.6.0" + monkeypatch.setitem(sys.modules, "torch", conda) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +@pytest.mark.parametrize("var", ["CUDA_VISIBLE_DEVICES"]) +@pytest.mark.parametrize("mask", ["", " ", "-1"]) +def test_a_deliberately_emptied_mask_is_not_a_broken_install(monkeypatch, var, mask): + # Hiding the GPUs produces exactly the shape this feature keys on (torch sees none, + # nvidia-smi sees them all) without anything being broken. + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + monkeypatch.setenv(var, mask) + + assert hw.classify_torch_build() is None + assert hw._torch_gpu_mismatch_report() == {} + + monkeypatch.setenv(var, "0") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def test_a_runtime_that_raises_on_its_own_probe_counts_as_unavailable(monkeypatch): + import sys + + hostile = _fake_torch("cuda_dead") + + def _boom(): + raise RuntimeError("CUDA driver version is insufficient for CUDA runtime version") + + hostile.cuda.is_available = _boom + monkeypatch.setitem(sys.modules, "torch", hostile) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def _detect(monkeypatch): + """Run one detection pass with the globals restored afterwards.""" + for name, value in ( + ("DEVICE", None), + ("CHAT_ONLY", True), + ("CHAT_ONLY_REASON", None), + ("CHAT_ONLY_DETAIL", None), + ): + monkeypatch.setattr(hw, name, value) + with hw._DETECT_LOCK: + return hw._detect_hardware_locked() + + +def test_chat_only_stops_claiming_this_host_has_no_gpu(monkeypatch, cpu_torch_on_an_nvidia_host): + device = _detect(monkeypatch) + + assert device == hw.DeviceType.CPU + assert hw.CHAT_ONLY is True + assert hw.CHAT_ONLY_REASON == "torch_cpu_build" + assert hw.CHAT_ONLY_REASON != "no_gpu" + assert hw.CHAT_ONLY_DETAIL == "2.11.0+cpu" + + +def test_a_host_that_really_has_no_gpu_still_reads_no_gpu(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", None) + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + + _detect(monkeypatch) + + assert hw.CHAT_ONLY_REASON == "no_gpu" + assert hw.CHAT_ONLY_DETAIL is None + assert hw._torch_gpu_mismatch_report() == {} + + +def test_a_healthy_cuda_host_reports_no_mismatch_at_all(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + + assert hw._torch_gpu_mismatch_report() == {} + + +def _system_gpu_info(monkeypatch): + """(gpu, inference_gpu) from main, with the real visibility probe in place.""" + monkeypatch.setattr( + hardware_pkg, + "get_visible_gpu_utilization", + lambda: {"available": False, "backend": "cpu", "devices": []}, + ) + monkeypatch.setattr(hardware_pkg, "get_vulkan_inference_gpu_info", lambda: None) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CPU) + monkeypatch.setattr(main, "_system_gpu_cache", None) + # Warm the inventory the way startup does: the mismatch report reads it WITHOUT + # blocking, because /api/system is polled every three seconds, so a cold cache answers + # unknown. In a running backend _detect_hardware_locked has already blocked once. + hw.get_physical_gpu_inventory() + hw.torch_build_snapshot() + return main._get_cached_system_gpu_info(SimpleNamespace(debug = lambda *args: None)) + + +def test_the_system_endpoint_names_the_cards_without_offering_them( + monkeypatch, cpu_torch_on_an_nvidia_host +): + gpu, _inference_gpu = _system_gpu_info(monkeypatch) + + # The hard constraint: `devices` is the runtime-usable list that model fit budgets + # against and the training picker pins from. + assert gpu["devices"] == [] + assert gpu["available"] is False + + assert gpu["mismatch"]["reason"] == "torch_cpu_build" + assert gpu["mismatch"]["torch_version"] == "2.11.0+cpu" + assert gpu["mismatch"]["physical_count"] == 2 + assert gpu["mismatch"]["sources"] == ["nvidia-smi"] + assert [device["name"] for device in gpu["physical_devices"]] == [ + "NVIDIA RTX A4000", + "NVIDIA RTX A4000, Founders Edition", + ] + assert hw.get_backend_visible_gpu_info()["physical_devices"] == gpu["physical_devices"] + + +def test_a_cpu_host_with_no_cards_publishes_neither_field(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", None) + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + + gpu, _inference_gpu = _system_gpu_info(monkeypatch) + + assert gpu["devices"] == [] + assert "mismatch" not in gpu + assert "physical_devices" not in gpu + + +@pytest.mark.parametrize("local", ["cpu", "cpu.cxx11.abi", "cpu.cxx11abi", "CPU"]) +def test_extended_cpu_local_tags_are_still_cpu_builds(monkeypatch, local): + """PyTorch publishes CPU wheels whose local tag carries a suffix. + + An exact "cpu" match called those CUDA wheels whose runtime failed to start, and + the UI then pointed the user at a driver rather than at the reinstall that is the + actual fix. + """ + import sys + + wheel = _fake_torch("cpu") + wheel.__version__ = f"2.8.0+{local}" + monkeypatch.setitem(sys.modules, "torch", wheel) + assert hw.classify_torch_build() == "torch_cpu_build" + + +def test_a_cuda_local_tag_is_not_mistaken_for_a_cpu_one(monkeypatch): + import sys + + wheel = _fake_torch("cuda_dead") + wheel.__version__ = "2.6.0+cu124" + monkeypatch.setitem(sys.modules, "torch", wheel) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def test_export_and_video_stop_saying_no_accelerator_was_found(monkeypatch): + """Those two pages render the message verbatim, so it has to match the System tab. + + Telling a two-A4000 host that no supported accelerator was found contradicts the + inventory the same server just published, and points at hardware instead of at the + repair. + """ + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr(hw, "_has_torch", lambda: True) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + # The verdict is re-derived against the live inventory, so the cards have to be in + # it. Without them the honest answer really is that this host has no GPU. + monkeypatch.setattr( + hw, "current_chat_only_verdict", lambda: (hw.CHAT_ONLY_REASON, hw.CHAT_ONLY_DETAIL) + ) + + export = hw.export_capability() + assert export["export_supported"] is False + assert export["export_unsupported_reason"] == "torch_cpu_build" + assert "2.11.0+cpu" in export["export_unsupported_message"] + assert "No supported" not in export["export_unsupported_message"] + + video = hw.video_capability() + assert video["video_supported"] is False + assert video["video_unsupported_reason"] == "torch_cpu_build" + assert "No supported accelerator" not in video["video_unsupported_message"] + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + assert hw.export_capability()["export_unsupported_reason"] == "torch_cuda_unavailable" + assert "driver" in hw.video_capability()["video_unsupported_message"] + + +def test_a_genuinely_gpu_less_host_keeps_the_old_wording(monkeypatch): + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + # This suite runs on a real NVIDIA box, so without these the verdict refresh correctly + # notices the accelerator this test is pretending not to have. + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CPU) + monkeypatch.setattr(hw, "_torch_reports_a_usable_accelerator", lambda: False) + monkeypatch.setattr(hw, "classify_torch_build", lambda **_kw: None) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": False, "devices": [], "unknown": False}, + ) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "no_gpu") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr(hw, "_has_torch", lambda: True) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + + assert hw.export_capability()["export_unsupported_reason"] == "no_accelerator" + assert hw.video_capability()["video_unsupported_reason"] == "no_accelerator" + + +def test_nvidia_smi_is_resolved_from_the_standard_windows_locations(monkeypatch): + """A driver install can leave nvidia-smi.exe off PATH entirely. + + The bare name then raises FileNotFoundError and the inventory comes back empty on + exactly the host this feature exists for. setup.ps1 already falls back to these two + paths, so the backend has to as well. + """ + monkeypatch.setattr(nvidia.platform, "system", lambda: "Windows") + monkeypatch.setattr(nvidia.shutil, "which", lambda _name: None) + monkeypatch.setenv("ProgramFiles", r"C:\Program Files") + monkeypatch.setenv("SystemRoot", r"C:\Windows") + + nvsmi = nvidia.os.path.join(r"C:\Program Files", r"NVIDIA Corporation\NVSMI\nvidia-smi.exe") + monkeypatch.setattr(nvidia.os.path, "isfile", lambda p: p == nvsmi) + assert nvidia._nvidia_smi_executable() == nvsmi + + system32 = nvidia.os.path.join(r"C:\Windows", r"System32\nvidia-smi.exe") + monkeypatch.setattr(nvidia.os.path, "isfile", lambda p: p == system32) + assert nvidia._nvidia_smi_executable() == system32 + + monkeypatch.setattr(nvidia.os.path, "isfile", lambda _p: False) + assert nvidia._nvidia_smi_executable() == "nvidia-smi" + + +def test_path_resolution_is_a_no_op_off_windows_and_when_path_has_it(monkeypatch): + monkeypatch.setattr(nvidia.platform, "system", lambda: "Linux") + monkeypatch.setattr(nvidia.shutil, "which", lambda _name: None) + assert nvidia._nvidia_smi_executable() == "nvidia-smi" + + monkeypatch.setattr(nvidia.platform, "system", lambda: "Windows") + monkeypatch.setattr(nvidia.shutil, "which", lambda _name: "/usr/bin/nvidia-smi") + assert nvidia._nvidia_smi_executable() == "/usr/bin/nvidia-smi" + + +@pytest.mark.parametrize("mask", ["", " ", "-1"]) +@pytest.mark.parametrize("var", ["HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"]) +def test_a_hip_mask_only_counts_where_it_can_hide_something(monkeypatch, var, mask): + """A mask that cannot take effect must not silence the mismatch. + + Windows HIP has no ROCr layer, and this module's own visibility resolver already + ignores ROCR_VISIBLE_DEVICES there. A stray empty one on a Windows NVIDIA host + would otherwise restore exactly the "no GPU" verdict this feature exists to + correct. The HIP variables likewise address AMD devices, so on an NVIDIA-only + inventory they mask nothing. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + # sys.platform too, not just platform.system(): the ROCR row is gated on sys.platform + # (hardware.py's own resolver reads it), so on a Windows runner ROCR is correctly + # ignored and this row asserted the opposite of what that host should do. + monkeypatch.setattr(hw.sys, "platform", "linux") + monkeypatch.setenv(var, mask) + # Prime the cache: the mask set reads the inventory WITHOUT blocking (it is reached + # from the verdict /api/liveness reads), and a cold cache keeps every mask. + hw.get_physical_gpu_inventory() + + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": True, "devices": [{"vendor": "amd"}], "sources": ["x"]}, + ) + assert hw.classify_torch_build() is None + + +def test_rocr_is_ignored_on_windows_even_on_an_amd_host(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr(hw.sys, "platform", "win32") + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": True, "devices": [{"vendor": "amd"}], "sources": ["x"]}, + ) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + monkeypatch.delenv("ROCR_VISIBLE_DEVICES") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() is None + + +def test_an_inventory_that_answers_nothing_keeps_every_mask(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr(hw.sys, "platform", "linux") + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"available": False, "devices": []} + ) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() is None + + +def test_a_deliberate_cpu_install_is_not_reported_as_broken(monkeypatch, tmp_path): + """Pinning /cpu on a machine that has a GPU is a supported thing to do. + + No mask is empty in that case, so the classifier called the wheel the user asked + for broken, and the UI offered a repair whose only effect would be to replace it. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY", raising = False) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_URL", raising = False) + monkeypatch.setattr(hw.sys, "prefix", str(tmp_path)) + + assert hw.classify_torch_build() == "torch_cpu_build" + + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", "cpu") + assert hw.classify_torch_build() is None + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY") + + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://download.pytorch.org/whl/cpu/") + assert hw.classify_torch_build() is None + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_URL") + + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", "cu124") + assert hw.classify_torch_build() == "torch_cpu_build" + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY") + + # A recorded cpu counts only when the record says someone NAMED it: setup.ps1 selects + # /cpu automatically on a GPU-less host and records it identically. + manifest = tmp_path / "unsloth_install_manifest.json" + manifest.write_text('{"schema": 1, "expected_torch_tag": "cpu"}', encoding = "utf-8") + assert hw.classify_torch_build() == "torch_cpu_build" + + manifest.write_text( + '{"schema": 1, "expected_torch_tag": "cpu", "expected_torch_tag_pinned": true}', + encoding = "utf-8", + ) + assert hw.classify_torch_build() is None + + manifest.write_text( + '{"schema": 1, "expected_torch_tag": "cpu", "expected_torch_tag_pinned": false}', + encoding = "utf-8", + ) + assert ( + hw.classify_torch_build() == "torch_cpu_build" + ), "an automatic CPU selection is not a choice to protect" + + manifest.write_text('{"schema": 1, "expected_torch_tag": "cu124"}', encoding = "utf-8") + assert hw.classify_torch_build() == "torch_cpu_build" + + manifest.write_text('{"schema": 1}', encoding = "utf-8") + assert hw.classify_torch_build() == "torch_cpu_build" + manifest.write_text("{not json", encoding = "utf-8") + assert hw.classify_torch_build() == "torch_cpu_build" + + +def test_a_dead_accelerator_wheel_is_unaffected_by_a_cpu_record(monkeypatch, tmp_path): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr(hw.sys, "prefix", str(tmp_path)) + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", "cpu") + assert hw.classify_torch_build() is None # the pin is honoured first + + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def test_linux_amd_and_intel_cards_are_inventoried_from_sysfs(monkeypatch, tmp_path): + """The AMD/Linux shape of #8473: nvidia-smi contributes nothing. + + sysfs rather than amd-smi, because amd-smi is separate ROCm userspace and this is + the host whose ROCm install is in question. vendor 0x1002 and a byte-valued + mem_info_vram_total are both documented amdgpu interfaces. + """ + drm = tmp_path / "drm" + for name, vendor, vram in ( + ("card0", "0x1002\n", str(16 * 1024**3)), + ("card1", "0x1002\n", "0"), # an APU with no dedicated VRAM + ("card2", "0x10de\n", str(1024**3)), # NVIDIA: nvidia-smi's business, not this + ("card4", "0x8086\n", ""), # an Arc card: no vram total published + ("card3", "0x1002\n", "not a number"), + ): + device = drm / name / "device" + device.mkdir(parents = True) + (device / "vendor").write_text(vendor, encoding = "utf-8") + (device / "mem_info_vram_total").write_text(vram, encoding = "utf-8") + (drm / "card0-DP-1").mkdir() + (drm / "card9").mkdir() + + real_listdir = hw.os.listdir + monkeypatch.setattr( + hw.os, + "listdir", + lambda p: real_listdir(str(drm)) if p == "/sys/class/drm" else real_listdir(p), + ) + real_join = hw.os.path.join + monkeypatch.setattr( + hw.os.path, + "join", + lambda *parts: ( + real_join(str(drm), *parts[1:]) if parts[0] == "/sys/class/drm" else real_join(*parts) + ), + ) + + records = hw._linux_drm_sysfs_records() + assert [r["index"] for r in records] == [0, 1, 2, 3] + assert [r["vendor"] for r in records] == ["amd", "amd", "amd", "intel"] + assert all(r["source"] == "sysfs-drm" for r in records) + assert records[0]["memory_total_gb"] == 16.0 + assert records[1]["memory_total_gb"] is None + assert records[2]["memory_total_gb"] is None + assert records[3]["memory_total_gb"] is None + + +def test_the_sysfs_probe_is_silent_where_there_is_no_sysfs(monkeypatch): + monkeypatch.setattr(hw.os, "listdir", lambda _p: (_ for _ in ()).throw(OSError("no such path"))) + assert hw._linux_drm_sysfs_records() == [] + + +def test_a_token_authenticated_cpu_pin_is_still_a_cpu_pin(monkeypatch, tmp_path): + """A pinned index may carry its credential in the query, which is supported. + + A raw final-segment split sees "cpu?token=..." there, so the deliberate CPU build on + a GPU host was reported as broken and offered a repair that would replace it. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw.sys, "prefix", str(tmp_path)) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY", raising = False) + + for pinned in ( + "https://download.pytorch.org/whl/cpu?token=abc/", + "https://mirror.corp.example/whl/cpu#sha256=deadbeef", + "https://mirror.corp.example/whl/cpu//", + "https://user:pw@mirror.corp.example/whl/CPU", + ): + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", pinned) + assert hw.classify_torch_build() is None, pinned + + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://mirror.corp.example/whl/cu128?token=abc") + assert hw.classify_torch_build() == "torch_cpu_build" + + +@pytest.mark.parametrize( + "url,leaf", + [ + ("https://download.pytorch.org/whl/cpu?token=x/", "cpu"), + ("https://download.pytorch.org/whl/cu128//", "cu128"), + ("https://download.pytorch.org/whl/cpu#f", "cpu"), + ("", ""), + (" ", ""), + ], +) +def test_the_leaf_reader_matches_the_installers(url, leaf): + assert hw._torch_index_leaf(url) == leaf + + +def test_the_chat_only_verdict_follows_the_inventory(monkeypatch): + """The inventory refreshes on a 60s TTL; the startup verdict never did. + + An eGPU attached after launch, or a driver that finished restarting, left the + sidebar and the Export and Video pages saying no accelerator exists while + /api/system listed the card and published a mismatch. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "no_gpu") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [{"vendor": "nvidia"}]} + ) + # The verdict never probes torch inline; measure this fake host as detection would. + hw.torch_build_snapshot() + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr(hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": []}) + assert hw.current_chat_only_verdict() == ("no_gpu", None) + + +@pytest.mark.parametrize("frozen", ["mlx_unavailable", "detection_failed", "intel_mac", None]) +def test_the_other_verdicts_are_left_exactly_as_detection_set_them(monkeypatch, frozen): + # A 60 second probe cannot change any of these, and re-deriving them would fight + # detect_hardware() rather than follow it. + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", frozen) + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "whatever detection recorded") + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [{"vendor": "nvidia"}]} + ) + assert hw.current_chat_only_verdict() == (frozen, "whatever detection recorded") + + +def test_a_probe_that_raises_keeps_the_frozen_verdict(monkeypatch): + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr( + hw, "classify_torch_build", lambda **_kw: (_ for _ in ()).throw(RuntimeError("boom")) + ) + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + + +def test_export_and_video_read_the_refreshed_verdict(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "no_gpu") # the STALE verdict + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr(hw, "_has_torch", lambda: True) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [{"vendor": "nvidia"}]} + ) + + hw.torch_build_snapshot() + assert hw.export_capability()["export_unsupported_reason"] == "torch_cpu_build" + assert hw.video_capability()["video_unsupported_reason"] == "torch_cpu_build" + + +def test_windows_intel_adapters_are_inventoried_too(monkeypatch): + """An Arc host whose XPU wheel was replaced has exactly this shape. + + nvidia-smi contributes nothing there either, and the registry scan was filtered to + AMD, so the inventory came back empty and the mismatch was discarded. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + monkeypatch.setattr( + hw, "_windows_live_adapter_names", lambda: ["Intel(R) Arc(TM) A770 Graphics"] + ) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + {0x1: {"name": "Intel(R) Arc(TM) A770", "dedicated_memory_bytes": 16 * 1024**3}} + if vendor_id == hw._INTEL_PCI_VENDOR_ID + else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + + assert [d["vendor"] for d in inventory["devices"]] == ["intel"] + assert inventory["devices"][0]["memory_total_gb"] == 16.0 + assert inventory["available"] is True + assert inventory["unknown"] is True + + _smi(monkeypatch, "", returncode = 0) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + inventory = hw.get_physical_gpu_inventory() + assert [d["vendor"] for d in inventory["devices"]] == ["intel"] + assert inventory["unknown"] is False + + +def test_a_transient_probe_failure_does_not_retire_a_settled_mismatch(monkeypatch): + """nvidia-smi timing out is not the GPU going away. + + The aggregate probe returns a structured empty result rather than raising, so the + refreshed verdict read it as "no cards" and handed the user the opposite advice for + a whole cache interval. + """ + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr(hw, "classify_torch_build", lambda **_kw: "torch_cpu_build") + monkeypatch.setattr(hw, "_torch_reports_a_usable_accelerator", lambda: False) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": False, "devices": [], "unknown": True}, + ) + hw.torch_build_snapshot() + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": False, "devices": [], "unknown": False}, + ) + assert hw.current_chat_only_verdict() == ("no_gpu", None) + + +def test_the_health_path_never_waits_on_the_gpu_probe(monkeypatch): + """/api/health and /api/liveness both read the chat-only verdict. + + The NVIDIA half shells out with a 10 second timeout and the lock makes concurrent + callers queue, so a hung driver -- the exact host this feature is for -- would stall + the event loop every time the TTL expired, against a 2 second desktop timeout. + """ + import sys + + calls = {"blocking": 0, "threads": 0} + + def _probe(): + calls["blocking"] += 1 + return { + "available": True, + "devices": [{"vendor": "nvidia"}], + "sources": ["x"], + "unknown": False, + } + + class _Thread: + def __init__(self, *a, **k): + calls["threads"] += 1 + + def start(self): + pass + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "_probe_physical_gpu_inventory", _probe) + monkeypatch.setattr(hw.threading, "Thread", _Thread) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", False) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + + hw.torch_build_snapshot() + # That warm-up IS the blocking pass, standing in for detection. What the request path + # may not do is shell out again, so count from here. + calls["blocking"] = 0 + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", False) + + # Cold inventory cache: no probe inline, the refresh goes to a thread, and the explicit + # unknown keeps the frozen verdict. + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + assert calls["blocking"] == 0, "nothing may shell out on the request path" + assert calls["threads"] >= 1 + + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", True) + before = calls["threads"] + hw.get_physical_gpu_inventory(block = False) + assert calls["threads"] == before + + +def test_a_stale_cache_is_served_rather_than_re_probed_off_the_request_path(monkeypatch): + warm = { + "available": True, + "devices": [{"vendor": "nvidia"}], + "sources": ["x"], + "unknown": False, + } + calls = {"n": 0} + + def _probe(): + calls["n"] += 1 + return warm + + class _Thread: + def __init__(self, *a, **k): + pass + + def start(self): + pass + + monkeypatch.setattr(hw, "_probe_physical_gpu_inventory", _probe) + monkeypatch.setattr(hw.threading, "Thread", _Thread) + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", False) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", (hw.time.monotonic() - 3600, warm)) + + assert hw.get_physical_gpu_inventory(block = False) is warm + assert calls["n"] == 0, "a stale answer beats a subprocess on the request path" + + assert hw.get_physical_gpu_inventory() is warm + assert calls["n"] == 1 + + +def test_a_process_that_cannot_start_a_thread_keeps_the_stale_answer(monkeypatch): + warm = { + "available": True, + "devices": [{"vendor": "nvidia"}], + "sources": ["x"], + "unknown": False, + } + monkeypatch.setattr(hw, "_physical_gpu_inventory_refreshing", False) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", (hw.time.monotonic() - 3600, warm)) + + def _boom(*_a, **_k): + raise RuntimeError("can't start new thread") + + monkeypatch.setattr(hw.threading, "Thread", _boom) + assert hw.get_physical_gpu_inventory(block = False) is warm + assert hw._physical_gpu_inventory_refreshing is False + + +def test_a_stale_registry_record_is_not_reported_as_a_gpu(monkeypatch): + """The DirectX registry outlives the hardware. + + setup.ps1 says so and uses these records only to RE-LABEL an adapter its live WMI + scan also returned, never to add one. A CPU-only machine with a driver record left + behind would otherwise be told it has an unusable GPU and offered a repair that + cannot restore absent hardware. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + monkeypatch.setattr(hw, "_windows_live_adapter_names", lambda: ["Microsoft Basic Display"]) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + {0x1: {"name": "AMD Radeon RX 6800", "dedicated_memory_bytes": 16 * 1024**3}} + if vendor_id == hw._AMD_PCI_VENDOR_ID + else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + assert inventory["devices"] == [] + assert inventory["sources"] == [] + + +def test_a_live_scan_that_cannot_answer_reports_unknown_rather_than_guessing(monkeypatch): + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + monkeypatch.setattr(hw, "_windows_live_adapter_names", lambda: None) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + {0x1: {"name": "AMD Radeon RX 6800"}} if vendor_id == hw._AMD_PCI_VENDOR_ID else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + assert inventory["devices"] == [] + assert inventory["unknown"] is True + + +@pytest.mark.parametrize( + "registry,live,matched", + [ + ("AMD Radeon RX 7900 XT", "AMD Radeon RX 7900 XT", True), + ("AMD Radeon RX 7900 XT", "AMD Radeon RX 7900 XT Graphics", True), + ("Intel(R) Arc(TM) A770 Graphics", "Intel(R) Arc(TM) A770", True), + ("AMD Radeon RX 6800", "NVIDIA GeForce RTX 4090", False), + ("", "AMD Radeon RX 6800", False), + ("AMD Radeon RX 6800", "", False), + ], +) +def test_the_registry_to_live_join_matches_setup_ps1(registry, live, matched): + assert hw._adapter_name_is_live(registry, [live]) is matched + + +def test_an_ordinary_intel_igpu_does_not_establish_a_mismatch(monkeypatch, tmp_path): + """setup.sh does not autodetect Linux XPU, and setup.ps1 limits it to Arc. + + An Intel UHD iGPU beside a CPU wheel is the correct state of that machine, so + counting it would report a broken install and offer a repair that reinstalls the + very CPU build it just replaced. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw.sys, "prefix", str(tmp_path)) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY", raising = False) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_URL", raising = False) + + igpu = [{"vendor": "intel", "name": None, "index": 0}] + assert hw._devices_that_can_establish_a_mismatch(igpu) == [] + + arc = [{"vendor": "intel", "name": "Intel(R) Arc(TM) A770 Graphics", "index": 0}] + assert hw._devices_that_can_establish_a_mismatch(arc) == arc + others = [ + {"vendor": "nvidia", "index": 0}, + {"vendor": "amd", "index": 0, "gfx_candidates": ["gfx1100"]}, + ] + assert hw._devices_that_can_establish_a_mismatch(others) == others + + +def test_a_nameless_intel_card_counts_once_xpu_was_actually_chosen(monkeypatch, tmp_path): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw.sys, "prefix", str(tmp_path)) + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_URL", raising = False) + nameless = [{"vendor": "intel", "name": None, "index": 0}] + + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", "xpu") + assert hw._devices_that_can_establish_a_mismatch(nameless) == nameless + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_FAMILY") + + manifest = tmp_path / "unsloth_install_manifest.json" + manifest.write_text('{"schema": 1, "expected_torch_tag": "xpu"}', encoding = "utf-8") + assert hw._devices_that_can_establish_a_mismatch(nameless) == nameless + + manifest.write_text('{"schema": 1, "expected_torch_tag": "cpu"}', encoding = "utf-8") + assert hw._devices_that_can_establish_a_mismatch(nameless) == [] + + xpu_torch = _fake_torch("cpu") + xpu_torch.__version__ = "2.9.0+xpu" + monkeypatch.setitem(sys.modules, "torch", xpu_torch) + assert hw._devices_that_can_establish_a_mismatch(nameless) == nameless + + +def test_an_accelerator_that_came_back_retires_the_cached_verdict(monkeypatch): + """Only reason and detail are refreshed here; DEVICE and CHAT_ONLY are not. + + So a driver that finished restarting left Train and Export disabled AND the UI + saying there is no GPU, which is worse than the stale mismatch it replaced. + """ + import sys + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + monkeypatch.setattr(hw, "_REDETECTION_REQUESTED", False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + calls = {"n": 0} + monkeypatch.setattr(hw, "invalidate_detection", lambda: calls.__setitem__("n", calls["n"] + 1)) + # The recovery is measured, not probed inline. The recovery path drops the snapshot + # again, so the fresh pass re-reads this fake host rather than the pre-recovery answer. + hw.torch_build_snapshot() + + assert hw.current_chat_only_verdict() == ("torch_cuda_unavailable", "2.6.0+cu124") + assert calls["n"] == 1 + + hw.current_chat_only_verdict() + hw.current_chat_only_verdict() + assert calls["n"] == 1 + + +def test_a_host_whose_accelerator_never_came_back_is_not_re_detected(monkeypatch): + import sys + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr(hw, "_REDETECTION_REQUESTED", False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [{"vendor": "nvidia"}]} + ) + calls = {"n": 0} + monkeypatch.setattr(hw, "invalidate_detection", lambda: calls.__setitem__("n", calls["n"] + 1)) + + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + assert calls["n"] == 0 + + +def test_a_healthy_xpu_wheel_is_not_called_unavailable(monkeypatch): + """An Intel host that started while its runtime was down and later recovered. + + Asking only about CUDA classified a perfectly healthy +xpu wheel as + torch_cuda_unavailable forever, so the recovery branch was never reached and the + process stayed chat-only until restart. + """ + import sys + + xpu = _fake_torch("cuda_dead") + xpu.__version__ = "2.9.0+xpu" + xpu.xpu = types.SimpleNamespace(is_available = lambda: True) + monkeypatch.setitem(sys.modules, "torch", xpu) + assert hw.classify_torch_build() is None + + xpu.xpu = types.SimpleNamespace(is_available = lambda: False) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + def _boom(): + raise RuntimeError("Level Zero not initialised") + + xpu.xpu = types.SimpleNamespace(is_available = _boom) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +@pytest.mark.parametrize("mask", ["", " ", "-1"]) +def test_an_emptied_xpu_mask_is_a_deliberate_hide(monkeypatch, mask): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": True, "devices": [{"vendor": "intel"}]}, + ) + monkeypatch.setenv("ZE_AFFINITY_MASK", mask) + assert hw.classify_torch_build() is None + + monkeypatch.setenv("ZE_AFFINITY_MASK", "0") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def test_the_xpu_mask_is_ignored_on_an_nvidia_only_host(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda_dead")) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"available": True, "devices": [{"vendor": "nvidia"}]}, + ) + monkeypatch.setenv("ZE_AFFINITY_MASK", "") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + +def test_duplicate_registry_records_are_claimed_one_to_one(monkeypatch): + """Two identical cards leave two identical registry records. + + Remove one and WMI reports a single live adapter, which a reusable predicate matched + to BOTH: a GPU reported that is gone, and its VRAM counted twice. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + # nvidia-smi ANSWERS (absent), so only the registry vendors go unanswered here. + monkeypatch.setattr( + nvidia.subprocess, + "run", + lambda *_a, **_k: (_ for _ in ()).throw(FileNotFoundError("nvidia-smi")), + ) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + monkeypatch.setattr(hw, "_windows_live_adapter_names", lambda: ["AMD Radeon RX 7900 XT"]) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + { + 0x1: {"name": "AMD Radeon RX 7900 XT", "dedicated_memory_bytes": 20 * 1024**3}, + 0x2: {"name": "AMD Radeon RX 7900 XT", "dedicated_memory_bytes": 20 * 1024**3}, + } + if vendor_id == hw._AMD_PCI_VENDOR_ID + else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + assert len(inventory["devices"]) == 1, "one live adapter corroborates one record" + assert inventory["devices"][0]["memory_total_gb"] == 20.0 + + +def test_both_records_survive_when_both_cards_are_live(monkeypatch): + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, "", returncode = 9) + monkeypatch.setattr( + hw, + "_windows_live_adapter_names", + lambda: ["AMD Radeon RX 7900 XT", "AMD Radeon RX 7900 XT"], + ) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **_kw: ( + { + 0x1: {"name": "AMD Radeon RX 7900 XT"}, + 0x2: {"name": "AMD Radeon RX 7900 XT"}, + } + if vendor_id == hw._AMD_PCI_VENDOR_ID + else {} + ), + ) + + assert len(hw.get_physical_gpu_inventory()["devices"]) == 2 + + +def test_the_live_scan_stops_on_a_cim_failure_rather_than_reporting_none(monkeypatch): + """-ErrorAction SilentlyContinue exits 0 with empty stdout. + + That is indistinguishable from "no adapters", so the corroboration would have + dropped every real card on a host with damaged WMI, recreating the false no-GPU + result this whole change exists to remove. + """ + source = inspect.getsource(hw._windows_live_adapter_names) + command = "".join(line for line in source.splitlines() if not line.strip().startswith("#")) + assert "SilentlyContinue" not in command + assert "$ErrorActionPreference='Stop'" in command + + +def test_a_missing_nvidia_smi_does_not_warn_every_refresh(monkeypatch, capsys): + """The normal state of every CPU-only, AMD and Intel host. + + The inventory calls this on a 60 second refresh reached from the health and system + polls, so a warning here is a line a minute on a machine that is working correctly. + """ + + def _missing(*_a, **_k): + raise FileNotFoundError("nvidia-smi") + + monkeypatch.setattr(nvidia.subprocess, "run", _missing) + capsys.readouterr() + assert ( + nvidia._query_gpu_inventory("test") is nvidia.NVIDIA_SMI_ABSENT + ), "an absent CLI is its own answer, not the None that means a probe failed" + assert '"level": "warning"' not in capsys.readouterr().out + + def _hang(*_a, **_k): + raise subprocess.TimeoutExpired("nvidia-smi", 10) + + monkeypatch.setattr(nvidia.subprocess, "run", _hang) + capsys.readouterr() + assert nvidia._query_gpu_inventory("test") is None + assert '"level": "warning"' in capsys.readouterr().out + + +def test_the_recovery_actually_starts_a_detection_pass(monkeypatch): + """Retiring the epoch alone did nothing, so the round-six fix never took effect. + + invalidate_detection leaves DEVICE set and DETECTION_COMPLETE raised, so /api/health + kept reading the settled snapshot and start_background_detection returned at once + because DEVICE was not None. The process stayed chat-only until restart, which is + exactly what the fix claimed to solve. + """ + import sys + + calls = [] + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + monkeypatch.setattr(hw, "_REDETECTION_REQUESTED", False) + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + monkeypatch.setattr(hw, "invalidate_detection", lambda: calls.append("epoch") or 1) + monkeypatch.setattr(hw, "_discard_detection_locked", lambda: calls.append("discard")) + monkeypatch.setattr(hw, "start_background_detection", lambda: calls.append("start")) + + # detect_hardware() warms the snapshot on the real path; the verdict only ever reads it + # non-blocking, so measure this fake host the same way first. + hw.torch_build_snapshot() + hw.current_chat_only_verdict() + + # Order matters: discarding after the pass started would race it, and starting before + # the discard is the no-op this fixes. + assert calls == ["epoch", "discard", "start"] + + +def test_the_system_mismatch_report_does_not_block(monkeypatch): + # GET /api/system is polled every three seconds with _system_gpu_cache_lock held for the + # whole call, so a hung nvidia-smi would queue every concurrent read behind it. + source = inspect.getsource(hw._torch_gpu_mismatch_report) + assert "get_physical_gpu_inventory(block=False)" in source.replace(" ", "") + + +def test_an_unimportable_torch_still_reports_the_cards(monkeypatch, tmp_path): + """A CUDA wheel whose native runtime will not load. + + Returning None meant the mismatch report never ran, so a GPU host published no + physical_devices and the System tab said no visible GPU while nvidia-smi could + enumerate the card: the central failure, for the runtime-failure case. + """ + import importlib.abc + import importlib.util + import sys + + # A real import that raises, not an object in sys.modules whose attribute access does. + # Reading __spec__ off an existing sys.modules entry is import-machinery behaviour that + # changed in 3.13: on 3.10-3.12 `import torch` hands back the stub untouched, so the + # premise below silently inverted and this test passed for the wrong reason. + class _WillNotLoad(importlib.abc.Loader): + def create_module(self, spec): + return None + + def exec_module(self, module): + raise OSError("[WinError 126] cudart64_12.dll could not be loaded") + + class _Finder(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname, + path = None, + target = None, + ): + if fullname == "torch": + return importlib.util.spec_from_loader("torch", _WillNotLoad()) + return None + + monkeypatch.delitem(sys.modules, "torch", raising = False) + monkeypatch.setattr(sys, "meta_path", [_Finder(), *sys.meta_path]) + # _has_torch() is NOT forced here: it reports False for a wheel that will not import, + # and the early return on it used to keep this host from the on-disk fallback below. + assert hw._has_torch() is False, "the premise: an unimportable torch reads as absent" + + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.6.0+cu124") + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + for label in ("2.11.0+rocm7.2", "2.9.1+xpu"): + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda label = label: label) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.11.0+cpu") + assert hw.classify_torch_build() == "torch_cpu_build" + + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "") + assert hw.classify_torch_build() is None + + +def test_the_disk_label_reader_needs_no_interpreter(tmp_path): + pkg = tmp_path / "torch" + pkg.mkdir() + (pkg / "version.py").write_text( + '__version__ = "2.6.0+cu124"\ncuda = "12.4"\n', encoding = "utf-8" + ) + with patch.object( + hw.importlib.util, + "find_spec", + return_value = SimpleNamespace(submodule_search_locations = [str(pkg)]), + ): + assert hw._installed_torch_label_on_disk() == "2.6.0+cu124" + + with patch.object(hw.importlib.util, "find_spec", return_value = None): + assert hw._installed_torch_label_on_disk() == "" + + with patch.object(hw.importlib.util, "find_spec", side_effect = ValueError("boom")): + assert hw._installed_torch_label_on_disk() == "" + + +def test_an_untagged_xpu_build_is_not_called_a_cpu_wheel(monkeypatch): + """torch.version.xpu is the marker, and an untagged wheel can carry it. + + An Arc host whose Level Zero runtime is down has no local version tag and neither + a cuda nor a hip one, so it read as a CPU wheel and the UI offered a reinstall, + which is the one remedy that cannot help. _torch_reports_an_xpu_runtime() already + knew better, so the card established a mismatch and then got the wrong advice. + """ + import sys + + torch = types.ModuleType("torch") + torch.__version__ = "2.9.1" # untagged, as a source or vendor build is + torch.version = SimpleNamespace(cuda = None, hip = None, xpu = "20250101") + torch.cuda = SimpleNamespace(is_available = lambda: False) + torch.xpu = SimpleNamespace(is_available = lambda: False) + monkeypatch.setitem(sys.modules, "torch", torch) + assert hw.classify_torch_build() == "torch_cuda_unavailable" + + torch.version = SimpleNamespace(cuda = None, hip = None, xpu = None) + assert hw.classify_torch_build() == "torch_cpu_build" + + +def test_the_health_path_never_imports_torch_inline(monkeypatch): + """classify_torch_build() imports torch and asks CUDA and XPU if they are available. + + Both block while a driver is wedged or restarting, which is the state + torch_cuda_unavailable names, and /api/health and /api/liveness reach them through + the chat-only verdict against a two second desktop timeout. Making only the + inventory lookup non-blocking still left this on the request thread. + """ + calls = {"probes": 0, "threads": 0} + + def _classify(**_kw): + calls["probes"] += 1 + return "torch_cpu_build" + + class _Thread: + def __init__(self, *a, **k): + calls["threads"] += 1 + + def start(self): + pass + + monkeypatch.setattr(hw, "classify_torch_build", _classify) + monkeypatch.setattr(hw, "_torch_reports_a_usable_accelerator", lambda: False) + monkeypatch.setattr(hw.threading, "Thread", _Thread) + monkeypatch.setattr(hw, "_torch_build_snapshot_refreshing", False) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [], "unknown": True} + ) + + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + assert calls["probes"] == 0, "no torch probe may run on the health path" + assert calls["threads"] == 1 + + hw.torch_build_snapshot() + assert calls["probes"] == 1 + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + assert calls["probes"] == 1 + + +def test_a_recovery_re_measures_torch_rather_than_reusing_the_old_answer(monkeypatch): + """The snapshot that said the accelerator was missing must not outlive it. + + Left in place, the fresh detection pass the recovery starts would read the same + stale measurement for the rest of its TTL and settle on the verdict it just retired. + """ + import sys + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + monkeypatch.setattr(hw, "_REDETECTION_REQUESTED", False) + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + monkeypatch.setattr(hw, "invalidate_detection", lambda: 1) + monkeypatch.setattr(hw, "_discard_detection_locked", lambda: None) + monkeypatch.setattr(hw, "start_background_detection", lambda: None) + + hw.torch_build_snapshot() + assert hw._torch_build_snapshot_cache is not None + hw.current_chat_only_verdict() + assert ( + hw._torch_build_snapshot_cache is None + ), "the recovery must drop the measurement it was taken before" + + +def test_a_vendor_mask_does_not_hide_another_vendors_card(monkeypatch): + """A hybrid NVIDIA + Arc host with ZE_AFFINITY_MASK="". + + The mask hides the Arc and nothing else, but it was cancelling the whole + classification, so a CPU-only wheel went unreported for the NVIDIA card the user + never masked. The Arc still has to drop out of the mismatch inventory: it is + hidden on purpose, and counting it would offer a repair for a chosen configuration. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setenv("ZE_AFFINITY_MASK", "") + hybrid = [ + {"vendor": "nvidia", "name": "NVIDIA RTX A4000"}, + {"vendor": "intel", "name": "Intel(R) Arc(TM) A770"}, + ] + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": hybrid, "unknown": False} + ) + + assert hw.classify_torch_build() == "torch_cpu_build" + kept = hw._devices_that_can_establish_a_mismatch(hybrid) + assert [d["vendor"] for d in kept] == ["nvidia"] + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() is None + assert hw._devices_that_can_establish_a_mismatch(hybrid) == [] + + +def test_an_intel_only_host_still_reads_an_emptied_ze_mask_as_deliberate(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setenv("ZE_AFFINITY_MASK", "-1") + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"devices": [{"vendor": "intel", "name": "Arc A770"}], "unknown": False}, + ) + assert hw.classify_torch_build() is None + + +def test_an_untagged_conda_cuda_build_that_will_not_import_is_not_a_cpu_wheel(monkeypatch): + """torch/version.py records the runtime even when the wheel carries no local tag. + + The importable path already reads exactly these attributes, so reading only + __version__ on the failure path gave the same installation the opposite diagnosis + and sent the user to reinstall a GPU wheel it already has. + """ + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.9.1") + monkeypatch.setattr( + hw, "_installed_torch_markers_on_disk", lambda: {"cuda": "12.8", "hip": None, "xpu": None} + ) + assert hw._classification_from_disk_label() == "torch_cuda_unavailable" + + for marker in ("hip", "xpu"): + monkeypatch.setattr( + hw, + "_installed_torch_markers_on_disk", + lambda marker = marker: {"cuda": None, "hip": None, "xpu": None} | {marker: "1.0"}, + ) + assert hw._classification_from_disk_label() == "torch_cuda_unavailable" + + monkeypatch.setattr( + hw, "_installed_torch_markers_on_disk", lambda: {"cuda": None, "hip": None, "xpu": None} + ) + assert hw._classification_from_disk_label() == "torch_cpu_build" + + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "") + assert hw._classification_from_disk_label() is None + + +def test_the_marker_reader_parses_both_shapes_torch_has_shipped(tmp_path): + """Parsed, not executed: this runs when importing torch is the thing that fails.""" + package = tmp_path / "torch" + package.mkdir() + (package / "version.py").write_text( + "from typing import Optional\n" + "__version__ = '2.9.1'\n" + "debug = False\n" + "cuda: Optional[str] = '12.8'\n" + "hip = None\n" + "xpu: Optional[str] = None\n", + encoding = "utf-8", + ) + + class _Spec: + submodule_search_locations = [str(package)] + + import types as _types + + original = hw.importlib.util.find_spec + hw.importlib.util.find_spec = lambda name: _Spec() if name == "torch" else original(name) + try: + assert hw._installed_torch_markers_on_disk() == { + "cuda": "12.8", + "hip": None, + "xpu": None, + } + finally: + hw.importlib.util.find_spec = original + assert isinstance(_types, _types.ModuleType) + + +def test_an_unimportable_gpu_wheel_is_a_mismatch_rather_than_a_detection_failure(): + """detection_failed sends the user to the server log; the mismatch offers the repair. + + _has_torch() is False for a wheel whose runtime will not load, so detection took + the TORCH_IMPORT_ERROR arm and never reached the classification -- and the verdict + refresh deliberately freezes detection_failed, so nothing revisited it later. + """ + import pathlib + + source = pathlib.Path(hw.__file__).read_text(encoding = "utf-8") + branch = source[source.index("elif TORCH_IMPORT_ERROR is not None:") :] + branch = branch[: branch.index('elif platform.system() == "Darwin":')] + assert ( + "_classification_from_disk_label()" in branch + ), "the broken-runtime host must still be classified from the wheel on disk" + assert "torch_build_snapshot()" not in branch, ( + "and classified WITHOUT probing: importing torch is what failed here, it takes " + "seconds to fail on a real broken wheel, and a retry re-runs torch/__init__ " + "against the partial module tree the first attempt left behind" + ) + assert branch.index("detection_failed") < branch.index("_mismatch_verdict_for_this_host"), ( + "detection_failed stays the default; the mismatch only replaces it when the " + "inventory actually found a card" + ) + + +def test_the_mismatch_verdict_names_the_wheel_it_could_not_import(monkeypatch): + monkeypatch.setattr( + hw, + "torch_build_snapshot", + lambda **_kw: { + "reason": "torch_cuda_unavailable", + "usable": False, + "unknown": False, + }, + ) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: {"devices": [{"vendor": "nvidia", "name": "A4000"}], "unknown": False}, + ) + monkeypatch.setattr(hw, "_torch_version_label", lambda: None) + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.6.0+cu124") + + assert hw._mismatch_verdict_for_this_host() == ("torch_cuda_unavailable", "2.6.0+cu124") + + monkeypatch.setattr(hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": []}) + assert hw._mismatch_verdict_for_this_host() == (None, None) + + +def test_cuda_visible_devices_masks_an_amd_host_too(monkeypatch): + """HIP honours CUDA_VISIBLE_DEVICES alongside its own variables. + + This module's own visibility resolver reads all three together on an AMD host, so + mapping the variable to NVIDIA alone had an AMD-only box launched with + CUDA_VISIBLE_DEVICES="" reported as broken and offered a repair for a mask its + owner set on purpose. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + amd = [{"vendor": "amd", "name": "Radeon RX 7900 XT", "gfx_candidates": ["gfx1100"]}] + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": amd, "unknown": False} + ) + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() is None + assert hw._devices_that_can_establish_a_mismatch(amd) == [] + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + assert hw.classify_torch_build() == "torch_cpu_build" + assert hw._devices_that_can_establish_a_mismatch(amd) == amd + + +def test_the_health_path_never_retries_a_failed_torch_import(monkeypatch): + """The label was read with _torch_version_label(), which imports torch. + + On the broken-runtime host that import is what fails: it can take seconds, and + _has_torch() purges the partial module afterwards, so each call genuinely re-runs + the native load. /api/health, /api/liveness and GET /api/system all reach it. + """ + calls = {"n": 0} + + def _label(): + calls["n"] += 1 + return "should not be reached" + + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", "OSError('libcudart.so.12')") + monkeypatch.setattr(hw, "_torch_version_label", _label) + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.6.0+cu124") + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + monkeypatch.setattr(hw, "classify_torch_build", lambda **_kw: "torch_cuda_unavailable") + monkeypatch.setattr(hw, "_torch_reports_a_usable_accelerator", lambda: False) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: { + "devices": [{"vendor": "nvidia", "name": "A4000"}], + "sources": ["nvidia-smi"], + "unknown": False, + }, + ) + hw.torch_build_snapshot() + + assert hw.current_chat_only_verdict() == ("torch_cuda_unavailable", "2.6.0+cu124") + report = hw._torch_gpu_mismatch_report() + assert report["mismatch"]["torch_version"] == "2.6.0+cu124" + assert calls["n"] == 0, "the failed import must not be retried on a request path" + + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", None) + monkeypatch.setattr(hw, "_torch_version_label", lambda: "2.11.0+cpu") + assert hw._reported_torch_label() == "2.11.0+cpu" + + +def test_an_amd_card_the_installers_decline_is_not_a_broken_install(monkeypatch): + """setup.sh ships a ROCm wheel only for RDNA 2 and newer. + + A Polaris or RDNA 1 card is left on CPU torch on purpose, and the DRM sysfs walk + reports it as vendor amd like any other, so it was being counted as evidence of a + broken install and offered a repair that cannot make that GPU usable. + """ + monkeypatch.setattr(hw, "_expected_rocm_flavor_was_chosen", lambda: False) + monkeypatch.setattr(hw, "_torch_reports_a_hip_runtime", lambda: False) + + polaris = [{"vendor": "amd", "gfx_candidates": ["gfx803"]}] + assert hw._devices_that_can_establish_a_mismatch(polaris) == [] + + supported = [{"vendor": "amd", "gfx_candidates": ["gfx1100"]}] + assert hw._devices_that_can_establish_a_mismatch(supported) == supported + + # A host with no ROCm userspace names no arch at all, and setup.sh detects AMD only + # through rocminfo and amd-smi, so that machine was never going to get a ROCm wheel. + unnamed = [{"vendor": "amd"}] + assert hw._devices_that_can_establish_a_mismatch(unnamed) == [] + + assert hw._devices_that_can_establish_a_mismatch([{"vendor": "amd", "gfx": "gfx1201"}]) + assert hw._devices_that_can_establish_a_mismatch([{"vendor": "amd", "gfx": "gfx900"}]) == [] + + +def test_a_rocm_expectation_outranks_the_arch_table(monkeypatch): + """If this install asked for a ROCm wheel, its absence is a fault either way.""" + monkeypatch.setattr(hw, "_torch_reports_a_hip_runtime", lambda: False) + monkeypatch.setattr(hw, "_expected_rocm_flavor_was_chosen", lambda: True) + old = [{"vendor": "amd", "gfx_candidates": ["gfx803"]}] + assert hw._devices_that_can_establish_a_mismatch(old) == old + + monkeypatch.setattr(hw, "_expected_rocm_flavor_was_chosen", lambda: False) + monkeypatch.setattr(hw, "_torch_reports_a_hip_runtime", lambda: True) + assert hw._devices_that_can_establish_a_mismatch(old) == old + + +def test_the_supported_arch_set_matches_the_installer(monkeypatch): + """One table in two places drifts, and it drifts BOTH ways. + + install.sh's _amd_arch_index_family_for_gfx is the map that decides whether this + stack ships a ROCm wheel for a card (it mirrors install.ps1's $archFamilyMap). A gfx + it lists but the backend does not is a supported card reported as no_gpu with no + repair; one the backend lists but it does not is a repair that cannot help. + """ + import pathlib + import re + + root = pathlib.Path(hw.__file__).resolve().parents[4] + install_sh = (root / "install.sh").read_text(encoding = "utf-8") + block = install_sh[install_sh.index("_amd_arch_index_family_for_gfx()") :] + block = block[: block.index("esac")] + # Only the case LABELS: the values on the right are index families (gfx103X-all), whose + # prefixes a bare regex over the block would collect. + shipped = { + gfx + for line in block.splitlines() + for gfx in re.findall(r"gfx[0-9a-f]+", line.split(")")[0]) + } + assert len(shipped) > 10, "the installer's arch map moved" + + extra = {"gfx906"} + assert set(hw._ROCM_SUPPORTED_GFX) == shipped | extra, ( + "the backend must accept exactly the architectures the installers ship a wheel " + f"for; installer has {sorted(shipped | extra)}" + ) + stack = (root / "studio" / "install_python_stack.py").read_text(encoding = "utf-8") + assert "gfx906" in stack, "the gfx906 path this set carries has gone" + + +def test_the_gfx_probe_answers_nothing_without_a_rocm_userspace(monkeypatch): + """No rocminfo and no amd-smi is the common case on the host in question.""" + + def _missing(*_a, **_k): + raise FileNotFoundError("rocminfo") + + monkeypatch.setattr(hw.subprocess, "run", _missing) + assert hw._linux_amd_gfx_candidates() == [] + + class _Result: + stdout = " Name: gfx1030\n Name: gfx1030\n" + + monkeypatch.setattr(hw.subprocess, "run", lambda *_a, **_k: _Result()) + assert hw._linux_amd_gfx_candidates() == ["gfx1030"] + + +def test_a_cold_start_measures_the_inventory_before_honouring_a_mask(monkeypatch): + """An irrelevant empty mask must not decide the verdict from a cold cache. + + HIP_VISIBLE_DEVICES="" on an NVIDIA-only host is not a statement about the NVIDIA + card, but on the first pass the non-blocking read answers unknown, every mask counts + as relevant, and the classification was suppressed and cached as "torch is fine" for + a whole TTL -- withholding the repair from a host nvidia-smi describes a moment + later. + """ + import sys + + probes = {"n": 0} + + def _probe(): + probes["n"] += 1 + return { + "available": True, + "devices": [{"vendor": "nvidia", "name": "A4000"}], + "sources": ["nvidia-smi"], + "unknown": False, + } + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + monkeypatch.setattr(hw, "_probe_physical_gpu_inventory", _probe) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + + snapshot = hw.torch_build_snapshot() + + assert probes["n"] >= 1, "the cached pass must measure rather than read a cold cache" + assert ( + snapshot["reason"] == "torch_cpu_build" + ), "a HIP mask says nothing about an NVIDIA card and must not suppress the repair" + + +def test_an_absent_nvidia_smi_is_an_answer_not_a_failed_probe(monkeypatch): + """An AMD-only host has no nvidia-smi by design. + + Marking its inventory unknown kept a settled mismatch alive for good once the AMD + card that established it was detached: the sysfs walk correctly found nothing, and + the verdict refresh read the unknown as "the probe declined" and preserved the old + reason, while /api/system had already dropped the device rows. + """ + + def _missing(*_a, **_k): + raise FileNotFoundError("nvidia-smi") + + monkeypatch.setattr(nvidia.subprocess, "run", _missing) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_linux_drm_sysfs_records", lambda: []) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + + inventory = hw.get_physical_gpu_inventory() + assert inventory["devices"] == [] + assert inventory["unknown"] is False, ( + "no nvidia-smi is the normal state of an AMD or CPU-only host, not a probe " + "that could not answer" + ) + + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cpu_build") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.11.0+cpu") + monkeypatch.setattr(hw, "classify_torch_build", lambda **_kw: "torch_cpu_build") + monkeypatch.setattr(hw, "_torch_reports_a_usable_accelerator", lambda: False) + hw.torch_build_snapshot() + assert hw.current_chat_only_verdict() == ("no_gpu", None) + + def _hang(*_a, **_k): + raise subprocess.TimeoutExpired("nvidia-smi", 10) + + monkeypatch.setattr(nvidia.subprocess, "run", _hang) + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + assert hw.get_physical_gpu_inventory()["unknown"] is True + + +def test_a_second_recovery_can_still_ask_for_a_pass(monkeypatch): + """The guard is one request per recovery, not one per process. + + It was cleared only by detect_hardware(), while the recovery starts its pass through + start_background_detection(), which runs ensure_hardware_detected(). So the first + recovery raised the guard for good: a driver that flapped, or any later lifecycle + that published an inventory-sensitive CPU verdict, could never get another pass. + """ + import sys + + monkeypatch.setattr(hw, "_REDETECTION_REQUESTED", False) + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "torch_cuda_unavailable") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "2.6.0+cu124") + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cuda")) + starts = {"n": 0} + monkeypatch.setattr(hw, "invalidate_detection", lambda: 1) + monkeypatch.setattr(hw, "_discard_detection_locked", lambda: None) + monkeypatch.setattr( + hw, "start_background_detection", lambda: starts.__setitem__("n", starts["n"] + 1) + ) + + hw.torch_build_snapshot() + hw.current_chat_only_verdict() + assert starts["n"] == 1 + assert hw._REDETECTION_REQUESTED is True, "one request per recovery, while it runs" + + source = _hardware_source() + ensure = source[source.index("def ensure_hardware_detected(") :] + ensure = ensure[: ensure.index("def _detect_hardware_locked(")] + assert ( + "_REDETECTION_REQUESTED = False" in ensure + ), "a settled background pass must release the guard, or there is never a second one" + assert ensure.index("DETECTION_COMPLETE.set()") < ensure.index( + "_REDETECTION_REQUESTED = False" + ), "released only once the pass has actually published" + + +def test_one_capability_response_describes_one_host(monkeypatch): + """Three reads of a verdict that refreshes on a TTL can disagree. + + An eGPU attached or removed mid-response could give the reason from one host and the + message from another, or a reason with no message at all. + """ + verdicts = [ + ("torch_cpu_build", "2.11.0+cpu"), + ("no_gpu", None), + (None, None), + ] + monkeypatch.setattr(hw, "DEVICE", hw.DeviceType.CPU) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CPU) + monkeypatch.setattr(hw, "_has_torch", lambda: True) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr( + hw, "current_chat_only_verdict", lambda: verdicts.pop(0) if verdicts else (None, None) + ) + + export = hw.export_capability() + assert export["export_unsupported_reason"] == "torch_cpu_build" + assert "CPU-only build" in export["export_unsupported_message"] + assert len(verdicts) == 2, "the response must be built from ONE reading of the verdict" + + +def _hardware_source() -> str: + import pathlib + return pathlib.Path(hw.__file__).resolve().read_text(encoding = "utf-8") + + +def test_the_hardware_module_loads_without_the_rest_of_the_package(tmp_path): + """tests/python/test_e2e_no_torch_sandbox.py executes this module on its own. + + It builds a minimal stub tree -- loggers, structlog, utils/hardware -- and nothing + else, so a top-level import of anything further inside utils makes the module + unloadable there, which is how a Windows-only console-hiding helper broke hardware + detection on a host with no torch at all. + """ + import pathlib + import re as _re + + source = pathlib.Path(hw.__file__).resolve().read_text(encoding = "utf-8") + header = source[: source.index("logger = get_logger(__name__)")] + imports = _re.findall(r"^\s*(?:from|import)\s+([\w.]+)", header, _re.MULTILINE) + for module in imports: + if module == "utils.hardware" or module.startswith("utils.hardware."): + continue + assert not module.startswith("utils"), ( + f"{module} is imported at module scope; the no-torch sandbox stubs only " + "utils.hardware, so import it inside the function that needs it" + ) + + +def test_a_driver_without_the_cli_still_reports_its_cards(monkeypatch): + """/proc/driver/nvidia/gpus is published whatever nvidia-smi's state is. + + install_python_stack._has_usable_nvidia_gpu() falls back to it for the same reason, + so without this the installer could select or repair a CUDA wheel on a host where + the backend insisted there was no card, and the user got no_gpu with no repair. + """ + + def _missing(*_a, **_k): + raise FileNotFoundError("nvidia-smi") + + monkeypatch.setattr(nvidia.subprocess, "run", _missing) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 2) + + result = nvidia.get_physical_gpu_inventory() + assert result["available"] is True + assert result["absent"] is False + assert [d["vendor"] for d in result["devices"]] == ["nvidia", "nvidia"] + assert [d["name"] for d in result["devices"]] == [None, None] + assert [d["memory_total_gb"] for d in result["devices"]] == [None, None] + assert result["source"] == "proc-driver-nvidia" + + +def test_a_kfd_confirmed_amd_card_stays_eligible(monkeypatch): + """A minimal AMD host with /dev/kfd but no rocminfo or amd-smi. + + install_python_stack._has_rocm_gpu() uses the KFD topology as its fallback exactly + so `studio update` can repair CPU-only torch there, so a card the installer would + repair cannot be dropped from the evidence here. + """ + monkeypatch.setattr(hw, "_expected_rocm_flavor_was_chosen", lambda: False) + monkeypatch.setattr(hw, "_torch_reports_a_hip_runtime", lambda: False) + unnamed = [{"vendor": "amd"}] + + monkeypatch.setattr(hw, "_linux_kfd_reports_an_amd_gpu", lambda: True) + assert hw._devices_that_can_establish_a_mismatch(unnamed) == unnamed + + monkeypatch.setattr(hw, "_linux_kfd_reports_an_amd_gpu", lambda: False) + assert hw._devices_that_can_establish_a_mismatch(unnamed) == [] + + monkeypatch.setattr(hw, "_linux_kfd_reports_an_amd_gpu", lambda: True) + assert ( + hw._devices_that_can_establish_a_mismatch([{"vendor": "amd", "gfx_candidates": ["gfx803"]}]) + == [] + ) + + +def test_the_kfd_probe_rejects_a_non_amd_node(monkeypatch, tmp_path): + """The NVIDIA open kernel module registers KFD nodes of its own. + + Their vendor_id is 4318, not AMD's 4098, and the installer guards on exactly that, + so an NVIDIA-only host must not read as AMD here either. + """ + nodes = tmp_path / "nodes" + for name, gpu_id, vendor in (("0", "0", "4098"), ("1", "5555", "4318")): + node = nodes / name + node.mkdir(parents = True) + (node / "gpu_id").write_text(gpu_id, encoding = "utf-8") + (node / "properties").write_text(f"vendor_id {vendor}\n", encoding = "utf-8") + + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + real_listdir, real_open = hw.os.listdir, open + monkeypatch.setattr( + hw.os, + "listdir", + lambda path: real_listdir(nodes) if "kfd" in str(path) else real_listdir(path), + ) + monkeypatch.setattr( + hw.os.path, + "join", + lambda *parts: str(nodes.joinpath(*parts[1:])) if "kfd" in parts[0] else "/".join(parts), + ) + assert ( + hw._linux_kfd_reports_an_amd_gpu() is False + ), "a CPU node and an NVIDIA-owned node are not an AMD GPU" + + (nodes / "1" / "properties").write_text("vendor_id 4098\n", encoding = "utf-8") + assert hw._linux_kfd_reports_an_amd_gpu() is True + assert real_open is open + + +def test_the_installer_records_who_named_the_flavor(tmp_path, monkeypatch): + """The manifest has to carry the provenance, or the backend cannot ask. + + setup.ps1 publishes UNSLOTH_EXPECTED_TORCH_TAG for an automatic /cpu choice exactly + as it does for a pinned one, so the handover variable alone is not evidence. + """ + import importlib.util + import pathlib + + root = pathlib.Path(hw.__file__).resolve().parents[3] + spec = importlib.util.spec_from_file_location( + "_install_manifest_probe", root / "install_manifest.py" + ) + manifest_mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(manifest_mod) + + monkeypatch.setattr(manifest_mod, "venv_root", lambda: tmp_path) + monkeypatch.setattr(manifest_mod, "manifest_path", lambda root = None: tmp_path / "m.json") + monkeypatch.setattr(manifest_mod, "requirement_digests", lambda *_a, **_k: {}) + monkeypatch.setattr(manifest_mod, "installed_requirements_root", lambda *_a, **_k: None) + monkeypatch.setattr(manifest_mod, "_installed_version", lambda *_a, **_k: "0") + + import json as _json + + manifest_mod.write_manifest(expected_torch_tag = "cpu", expected_torch_tag_pinned = False) + written = _json.loads((tmp_path / "m.json").read_text(encoding = "utf-8")) + assert written["expected_torch_tag"] == "cpu" + assert written["expected_torch_tag_pinned"] is False + + manifest_mod.write_manifest(expected_torch_tag = "cpu", expected_torch_tag_pinned = True) + written = _json.loads((tmp_path / "m.json").read_text(encoding = "utf-8")) + assert written["expected_torch_tag_pinned"] is True + + manifest_mod.write_manifest(expected_torch_tag = "cpu") + written = _json.loads((tmp_path / "m.json").read_text(encoding = "utf-8")) + assert "expected_torch_tag_pinned" not in written + monkeypatch.setattr(manifest_mod, "read_manifest", lambda root = None: written) + assert manifest_mod.recorded_torch_flavor_was_pinned() is False + + +def test_a_vendor_that_did_not_answer_keeps_the_inventory_unknown(monkeypatch): + """A hybrid Intel iGPU plus NVIDIA dGPU host with nvidia-smi timing out. + + The Intel row cancelled the unknown, and _devices_that_can_establish_a_mismatch then + discarded the iGPU as ineligible, so the verdict saw neither an eligible card nor an + unanswered probe and downgraded a settled mismatch to no_gpu for a cache interval -- + hiding the repair while the NVIDIA probe was merely unavailable. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + monkeypatch.setattr( + hw, "_linux_drm_sysfs_records", lambda: [{"vendor": "intel", "name": None, "index": 0}] + ) + + def _timeout(*_a, **_k): + raise subprocess.TimeoutExpired("nvidia-smi", 10) + + monkeypatch.setattr(nvidia.subprocess, "run", _timeout) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + + inventory = hw.get_physical_gpu_inventory() + assert [d["vendor"] for d in inventory["devices"]] == ["intel"] + assert inventory["unknown"] is True + + _smi(monkeypatch, "0, NVIDIA RTX A4000, 16376\n") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + assert hw.get_physical_gpu_inventory()["unknown"] is False + + +def test_a_stale_registry_record_cannot_claim_a_longer_named_live_card(monkeypatch): + """ "RX 7900 XT" is a prefix of a live "RX 7900 XTX". + + Records are walked in LUID order and consume the first live name they match, so the + stale XT could claim the XTX that is really installed and the inventory would + publish the removed card's name and its VRAM. + """ + live = ["AMD Radeon RX 7900 XTX"] + assert hw._claim_live_adapter("AMD Radeon RX 7900 XTX", live) == 0 + + both = ["AMD Radeon RX 7900 XTX", "AMD Radeon RX 7900 XT"] + assert both[hw._claim_live_adapter("AMD Radeon RX 7900 XT", both)] == "AMD Radeon RX 7900 XT" + assert both[hw._claim_live_adapter("AMD Radeon RX 7900 XTX", both)] == "AMD Radeon RX 7900 XTX" + + # The prefix rule still has to work: the registry description and the WMI display name + # spell the same card differently, and one is routinely a prefix of the other. + assert hw._claim_live_adapter("AMD Radeon RX 7900 XTX", ["AMD Radeon RX 7900 XTX 24GB"]) == 0 + assert hw._claim_live_adapter("Something Else", live) is None + + +def test_a_broken_nvidia_smi_still_reports_the_kernel_driver_cards(monkeypatch): + """Absent is not the only way the CLI fails to answer. + + A nvidia-smi that hangs past its timeout, or exits non-zero, leaves the kernel + driver enumerating cards regardless. On a cold start there is no settled verdict for + the resulting unknown to protect, so the host was reported as having no GPU at all, + with no repair, for as long as the CLI stayed broken. + """ + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 1) + + def _hang(*_a, **_k): + raise subprocess.TimeoutExpired("nvidia-smi", 10) + + monkeypatch.setattr(nvidia.subprocess, "run", _hang) + result = nvidia.get_physical_gpu_inventory() + assert [d["vendor"] for d in result["devices"]] == ["nvidia"] + assert result["source"] == "proc-driver-nvidia" + + monkeypatch.setattr( + nvidia.subprocess, + "run", + lambda *_a, **_k: SimpleNamespace(returncode = 9, stdout = ""), + ) + assert nvidia.get_physical_gpu_inventory()["available"] is True + + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + assert nvidia.get_physical_gpu_inventory()["error"] is not None + + +def test_hip_visible_devices_outranks_the_cuda_alias(monkeypatch): + """HIP reads its own variables first; the alias is only a fallback. + + An AMD host that NAMES its devices in HIP_VISIBLE_DEVICES while inheriting an empty + CUDA_VISIBLE_DEVICES has not hidden anything, and _get_parent_visible_gpu_spec in + this same module already applies that precedence. + """ + import sys + + monkeypatch.setitem(sys.modules, "torch", _fake_torch("cpu")) + amd = [{"vendor": "amd", "gfx_candidates": ["gfx1100"]}] + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": amd, "unknown": False} + ) + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + assert hw.classify_torch_build() == "torch_cpu_build" + assert hw._devices_that_can_establish_a_mismatch(amd) == amd + + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + assert hw.classify_torch_build() is None + assert hw._devices_that_can_establish_a_mismatch(amd) == [] + + monkeypatch.delenv("HIP_VISIBLE_DEVICES") + assert hw.classify_torch_build() is None + + +# =============================================== a refresh that cannot answer is not news + + +def test_an_unanswerable_refresh_keeps_the_cards_it_already_found(monkeypatch): + """nvidia-smi timing out must not read as "the GPUs were removed". + + Overwriting the cache with the empty failure holds for a whole TTL, and the two + consumers then disagree inside one response: current_chat_only_verdict() keeps its + frozen mismatch on ``unknown`` while _torch_gpu_mismatch_report() reads the devices, + finds none, and drops physical_devices and mismatch from /api/system. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + good = hw.get_physical_gpu_inventory() + assert len(good["devices"]) == 2 and good["unknown"] is False + + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", (0.0, good)) + _smi(monkeypatch, "", returncode = 9) + after = hw.get_physical_gpu_inventory() + + assert [d["name"] for d in after["devices"]] == [d["name"] for d in good["devices"]] + assert after["available"] is True + # Still unknown: these rows describe the host as it was, not as this pass measured it. + assert after["unknown"] is True + assert after["unanswered"] == ["nvidia"] + + +def test_a_vendor_that_answered_none_is_allowed_to_lose_its_cards(monkeypatch): + """The carry-forward is per unanswered vendor only, so a real removal still lands.""" + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + _smi(monkeypatch, _TWO_A4000_ROWS) + good = hw.get_physical_gpu_inventory() + + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", (0.0, good)) + monkeypatch.setattr( + nvidia.subprocess, + "run", + lambda *_a, **_k: (_ for _ in ()).throw(FileNotFoundError("nvidia-smi")), + ) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + after = hw.get_physical_gpu_inventory() + + assert after["devices"] == [] + assert after["unknown"] is False + + +def test_a_registry_that_cannot_be_read_is_not_a_host_without_adapters(monkeypatch): + """`{}` from the DirectX helper covers both, so the inventory has to ask which.""" + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_physical_gpu_inventory_cache", None) + # nvidia-smi ANSWERS (absent), so only the registry vendors go unanswered here. + monkeypatch.setattr( + nvidia.subprocess, + "run", + lambda *_a, **_k: (_ for _ in ()).throw(FileNotFoundError("nvidia-smi")), + ) + monkeypatch.setattr(nvidia, "_linux_nvidia_procfs_gpu_count", lambda: 0) + monkeypatch.setattr(hw, "_windows_live_adapter_names", lambda: ["AMD Radeon RX 7900 XT"]) + monkeypatch.setattr( + hw, + "_windows_amd_adapter_records_by_luid", + lambda vendor_id = hw._AMD_PCI_VENDOR_ID, **kw: ( + None if kw.get("distinguish_failure") else {} + ), + ) + + inventory = hw.get_physical_gpu_inventory() + + assert inventory["devices"] == [] + assert inventory["unknown"] is True + assert inventory["unanswered"] == ["amd", "intel"] + + +def test_the_ranking_callers_still_see_an_empty_map(monkeypatch): + """Only the inventory asked for the distinction; nothing else changed shape.""" + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_windows_amd_adapter_records_or_none", lambda *a, **k: None) + assert hw._windows_amd_adapter_records_by_luid() == {} + assert hw._windows_amd_adapter_records_by_luid(distinguish_failure = True) is None + + +# ============================================ a Windows AMD card the registry did not name + + +@pytest.mark.parametrize( + "name,eligible", + [ + ("AMD Radeon RX 7900 XT", True), + ("AMD Radeon RX 9070 XT", True), + ("AMD Radeon 780M Graphics", True), + # Polaris and RDNA 1: no wheel family covers them, so a repair could not change + # anything and this host is on CPU torch on purpose. + ("AMD Radeon RX 580", False), + ("AMD Radeon RX 5700 XT", False), + ("", False), + ], +) +def test_a_windows_adapter_with_no_adapter_family_is_read_from_its_name( + monkeypatch, name, eligible +): + """AdapterFamily is written by the driver, so the reported RX 7900 XT arrives bare. + + Falling through to the KFD probe discards it: that probe is Linux-only and returns + False on Windows, which made Windows accidentally stricter than Linux for the exact + card this feature was written for. + """ + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw, "_linux_kfd_reports_an_amd_gpu", lambda: False) + device = {"vendor": "amd", "name": name, "source": "directx-registry"} + assert hw._amd_device_can_establish_a_mismatch(device) is eligible + + +def test_a_named_arch_still_wins_over_the_marketing_name(monkeypatch): + """The name table only fills a gap; it never overrides what the driver reported.""" + monkeypatch.setattr(hw, "_linux_kfd_reports_an_amd_gpu", lambda: True) + assert ( + hw._amd_device_can_establish_a_mismatch( + {"vendor": "amd", "name": "AMD Radeon RX 7900 XT", "gfx": "gfx803"} + ) + is False + ) + + +# ================================================== only a real ROCm family is a ROCm pin + + +@pytest.mark.parametrize( + "leaf,chosen", + [ + ("rocm6.4", True), + ("rocm7", True), + ("gfx1151", True), + ("gfx120X-all", True), + # Suffixed leaves are custom pins the installer routes verbatim. Reading one as ROCm + # waives the supported-architecture filter and calls a deliberate install broken. + ("rocm-rel-7.2.1", False), + ("rocm7.2-private", False), + ("gfx-mirror", False), + ("cpu", False), + ], +) +def test_only_the_installers_rocm_families_count_as_a_rocm_choice(monkeypatch, leaf, chosen): + monkeypatch.delenv("UNSLOTH_TORCH_INDEX_URL", raising = False) + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", f"https://example.invalid/whl/{leaf}") + monkeypatch.setattr(hw.sys, "prefix", "/nonexistent-prefix-for-this-test") + assert hw._expected_rocm_flavor_was_chosen() is chosen + + +def test_the_backend_and_the_installer_agree_on_the_family_predicate(): + """One vocabulary, two files: drift here is a silent behaviour split.""" + for leaf in ( + "rocm6.4", + "rocm7", + "gfx1151", + "gfx120X-all", + "rocm-rel-7.2.1", + "rocm7.2-private", + "gfx-mirror", + "cpu", + "cu124", + "", + ): + assert hw._is_pip_rocm_family_leaf(leaf) == _installer_rocm_family(leaf), leaf + + +def _installer_rocm_family(leaf: str) -> bool: + """install_python_stack._is_pip_rocm_family_leaf, loaded without importing the module.""" + import ast + import pathlib + import re as _re + + source = (pathlib.Path(__file__).resolve().parents[2] / "install_python_stack.py").read_text( + encoding = "utf-8" + ) + tree = ast.parse(source) + fn = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_is_pip_rocm_family_leaf" + ) + namespace: dict = {"re": _re} + exec(compile(ast.Module(body = [fn], type_ignores = []), "", "exec"), namespace) + return namespace["_is_pip_rocm_family_leaf"](leaf) + + +# ======================================= a detection failure the disk already explained + + +def test_a_disk_classified_detection_failure_transitions_when_the_probe_recovers(monkeypatch): + """torch will not import AND the OS probe timed out: only the inventory was missing. + + detect_hardware() classifies the wheel from disk for exactly this host, then publishes + detection_failed when the inventory cannot corroborate it. Freezing that forever leaves + /api/system reporting the recovered mismatch while the sidebar, Export and Video keep + saying detection failed and never offer the repair. + """ + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "detection_failed") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", OSError("[WinError 126] cudart64_12.dll")) + monkeypatch.setattr(hw, "_installed_torch_label_on_disk", lambda: "2.11.0+cpu") + monkeypatch.setattr( + hw, + "torch_build_snapshot", + lambda **_kw: {"reason": "torch_cpu_build", "usable": False, "unknown": False}, + ) + monkeypatch.setattr( + hw, + "get_physical_gpu_inventory", + lambda **_kw: { + "devices": [{"vendor": "nvidia", "name": "NVIDIA RTX A4000"}], + "unknown": False, + }, + ) + + assert hw.current_chat_only_verdict() == ("torch_cpu_build", "2.11.0+cpu") + + +def test_a_detection_failure_never_degrades_into_no_gpu(monkeypatch): + """A host that never measured cannot be told it has no GPU; that IS the claim + detect_hardware() refused to make for it.""" + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "detection_failed") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", OSError("boom")) + monkeypatch.setattr( + hw, + "torch_build_snapshot", + lambda **_kw: {"reason": "torch_cpu_build", "usable": False, "unknown": False}, + ) + monkeypatch.setattr( + hw, "get_physical_gpu_inventory", lambda **_kw: {"devices": [], "unknown": False} + ) + + assert hw.current_chat_only_verdict() == ("detection_failed", None) + + +@pytest.mark.parametrize("reason", ["mlx_unavailable", "intel_mac"]) +def test_the_other_frozen_verdicts_stay_frozen(monkeypatch, reason): + """Only the disk-classified failure moved; a 60 second probe still cannot speak to these.""" + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", reason) + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", "detail") + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", OSError("boom")) + assert hw.current_chat_only_verdict() == (reason, "detail") + + +def test_a_detection_failure_with_an_importable_torch_stays_frozen(monkeypatch): + """Nothing classified this host, so there is nothing for the inventory to confirm.""" + monkeypatch.setattr(hw, "CHAT_ONLY_REASON", "detection_failed") + monkeypatch.setattr(hw, "CHAT_ONLY_DETAIL", None) + monkeypatch.setattr(hw, "TORCH_IMPORT_ERROR", None) + assert hw.current_chat_only_verdict() == ("detection_failed", None) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index d57ff38f750..79984f49a39 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,15 +16,19 @@ ... """ +import ast import copy import gc import glob +import importlib.util +import json import os import platform import re import subprocess import sys import threading +import time import types from contextlib import contextmanager from importlib.metadata import PackageNotFoundError, version as pkg_version @@ -76,7 +80,10 @@ class DeviceType(str, Enum): # Why CHAT_ONLY is True (Train/Export disabled). None when training is enabled. # "mlx_unavailable": Apple Silicon but the MLX stack is missing, too old, or broken # (the usual cause of "Train/Export greyed out" on Macs after a reinstall dropped MLX); -# "intel_mac": Intel Mac (no PyTorch/MLX); "no_gpu": CPU-only non-Mac host. +# "intel_mac": Intel Mac (no PyTorch/MLX); "no_gpu": CPU-only non-Mac host; +# "torch_cpu_build" / "torch_cuda_unavailable": the host HAS GPUs, this PyTorch cannot +# use them -- see classify_torch_build(). Those two must not read as "no_gpu": the fix +# is repairing the install, not buying a GPU. CHAT_ONLY_REASON: Optional[str] = None # What exactly blocked the reason above, when there is something specific to say. Only # "mlx_unavailable" sets it today: the gate is all-or-nothing across mlx, mlx-lm and @@ -331,6 +338,1158 @@ def _mlx_stack_detail() -> Optional[str]: return "; ".join(blockers[:3]) +# Every other GPU probe here starts from get_device(), so all go quiet the moment +# torch.cuda.is_available() is False (issue #8473). This one asks what the OS sees. +# Display-only: nothing here may be merged into the runtime device lists, which feed the +# model-fit estimate and the training device pin. TTL cache, since an eGPU changes it. +_PHYSICAL_GPU_INVENTORY_TTL_SECONDS = 60.0 +_physical_gpu_inventory_lock = threading.Lock() +_physical_gpu_inventory_refresh_lock = threading.Lock() +_physical_gpu_inventory_refreshing = False +# One re-detection request per recovery, or a poll every few seconds starves detection. +_REDETECTION_REQUESTED = False +_physical_gpu_inventory_cache: Optional[tuple[float, Dict[str, Any]]] = None +# Same treatment for the torch side: classify_torch_build() calls availability probes that +# block on a wedged driver, and /api/health and /api/liveness reach it through the verdict. +_TORCH_BUILD_SNAPSHOT_TTL_SECONDS = 60.0 +_torch_build_snapshot_lock = threading.Lock() +_torch_build_snapshot_refresh_lock = threading.Lock() +_torch_build_snapshot_refreshing = False +_torch_build_snapshot_cache: Optional[tuple[float, Dict[str, Any]]] = None +# "No measurement yet", distinct from a measured {"reason": None} meaning torch is fine. +_UNKNOWN_TORCH_BUILD_SNAPSHOT: Dict[str, Any] = { + "reason": None, + "usable": False, + "unknown": True, +} + + +def _probe_physical_gpu_inventory() -> Dict[str, Any]: + """One uncached pass over the vendor probes. Never raises. + + ``index`` is the probe's own row number and is vendor-local, so it identifies a + device only together with ``vendor``. It is not a pin: see the note above. + """ + devices: list[Dict[str, Any]] = [] + sources: list[str] = [] + # "No devices" and "no probe could answer" are different facts: collapsing them let a + # transient nvidia-smi timeout read as "the GPU disappeared". PER VENDOR, so an Intel + # sysfs row cannot answer for an NVIDIA probe that timed out. + unanswered: set = set() + + try: + from . import nvidia + result = nvidia.get_physical_gpu_inventory() + except Exception as e: + logger.debug("NVIDIA physical inventory probe failed: %s", e) + unanswered.add("nvidia") + else: + # An absent nvidia-smi is an answer: the normal state of every AMD, Intel and CPU host. + if result.get("error") and not result.get("absent"): + unanswered.add("nvidia") + nvidia_devices = result.get("devices") or [] + if nvidia_devices: + devices.extend(nvidia_devices) + sources.append(result.get("source") or "nvidia-smi") + + # Windows AMD has no vendor CLI guaranteed present (amd-smi elevates a child without a + # HIP SDK), but the DirectX registry records every adapter with its dedicated memory. + if platform.system() == "Windows": + # The registry outlives the hardware, so a record only RE-LABELS an adapter the live WMI + # scan also returned, never adds one (setup.ps1 has the same rule). + _live = _windows_live_adapter_names() + for _vendor, _vendor_id in ( + ("amd", _AMD_PCI_VENDOR_ID), + # Same registry, one more vendor id: an Arc host whose XPU wheel became a CPU one. + ("intel", _INTEL_PCI_VENDOR_ID), + ): + try: + records = _windows_amd_adapter_records_by_luid(_vendor_id, distinguish_failure = True) + except Exception as e: + logger.debug("Windows %s adapter inventory probe failed: %s", _vendor, e) + records = None + if records is None: + # A registry key that could not be read is not "this vendor has no adapters": + # the aggregate would publish unknown=False with nothing in it, and the settled + # mismatch would collapse to no_gpu for a whole TTL. + unanswered.add(_vendor) + continue + if not records: + continue + if _live is None: + # The live scan could not answer, so contribute nothing and say so. + unanswered.add(_vendor) + continue + # Consumed one-to-one, or a reusable predicate matches one surviving adapter to both. + _unclaimed = list(_live) + _corroborated = [] + for luid in sorted(records): + _match = _claim_live_adapter(records[luid].get("name"), _unclaimed) + if _match is not None: + _unclaimed.pop(_match) + _corroborated.append(luid) + for ordinal, luid in enumerate(_corroborated): + record = records[luid] + dedicated = record.get("dedicated_memory_bytes") + devices.append( + { + "vendor": _vendor, + "index": ordinal, + "name": record.get("name"), + "memory_total_gb": (round(dedicated / 1024**3, 2) if dedicated else None), + **({"gfx": record["gfx"]} if record.get("gfx") else {}), + "source": "directx-registry", + } + ) + if _corroborated and "directx-registry" not in sources: + sources.append("directx-registry") + + # Linux AMD, the shape #8473 actually reported. sysfs rather than amd-smi: amd-smi is a + # separate ROCm userspace package, and the ROCm install is exactly what is in question + # here; the amdgpu kernel driver publishes these files with no subprocess at all. + if platform.system() == "Linux": + try: + sysfs_devices = _linux_drm_sysfs_records() + except Exception as e: + logger.debug("Linux DRM sysfs inventory probe failed: %s", e) + sysfs_devices = [] + unanswered.update(("amd", "intel")) + if sysfs_devices: + if any(device.get("vendor") == "amd" for device in sysfs_devices): + # The installers only ship a ROCm wheel for _ROCM_SUPPORTED_GFX, so an older card left on + # CPU torch on purpose must not read as broken. Attached to every AMD record, not matched + # one-to-one: setup.sh says the orderings do not have to agree. + _gfx = _linux_amd_gfx_candidates() + if _gfx: + for device in sysfs_devices: + if device.get("vendor") == "amd": + device["gfx_candidates"] = _gfx + devices.extend(sysfs_devices) + sources.append("sysfs-drm") + + _unanswered = sorted(unanswered - {device.get("vendor") for device in devices}) + return { + "available": bool(devices), + "devices": devices, + "sources": sources, + "unknown": bool(_unanswered), + # Which vendors, not just that one exists: the cache carries the previous answer + # forward per vendor, and only a vendor named here may be carried. + "unanswered": _unanswered, + } + + +def _linux_amd_gfx_candidates() -> list[str]: + """gfx targets the ROCm userspace reports for this host, sorted, or []. + + rocminfo is what studio/setup.sh itself reads to pick an arch, and amd-smi is its + fallback there too. Neither is required to exist: a host with no ROCm userspace at + all reports nothing here, and the caller then keeps the card rather than guessing + an arch for it. Bounded and never raises; this runs inside the inventory probe, + which is cached and refreshed off the request path. + """ + for command, pattern in ( + (["rocminfo"], r"\bgfx[0-9a-f]+\b"), + (["amd-smi", "static", "--asic"], r"\bgfx[0-9a-f]+\b"), + ): + try: + result = subprocess.run( + command, + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 10, + check = False, + ) + except FileNotFoundError: + continue + except (OSError, subprocess.TimeoutExpired) as e: + logger.debug("%s could not report a gfx target: %s", command[0], e) + continue + found = {match.group(0).lower() for match in re.finditer(pattern, result.stdout or "")} + if found: + return sorted(found) + return [] + + +def _hidden_console_kwargs() -> Dict[str, Any]: + """subprocess kwargs that keep a console window from flashing on Windows. + + Imported here rather than at module scope on purpose: this module is loaded very + early, and tests/python/test_e2e_no_torch_sandbox.py executes it against a minimal + stub tree with no utils.subprocess_compat in it, so a top-level import would make + the module unloadable on a host without the rest of the package. Empty kwargs are + the correct answer everywhere but Windows anyway. + """ + try: + from utils.subprocess_compat import windows_hidden_subprocess_kwargs + except Exception: + return {} + return windows_hidden_subprocess_kwargs() + + +def _windows_live_adapter_names() -> Optional[list[str]]: + """Display adapter names Windows reports as PRESENT, or None when it could not say. + + Win32_VideoController is the same live source setup.ps1 scans. None is the important + third answer: "the scan failed" must not read as "no adapters", or a corroboration + step would silently drop every real card on a host whose WMI is unavailable. + """ + if platform.system() != "Windows": + return None + try: + ps = ( + # -ErrorAction Stop, NOT SilentlyContinue: the suppressed form exits 0 with empty + # stdout, indistinguishable from "no adapters" on a host with damaged WMI. + "$ErrorActionPreference='Stop';" + "(Get-CimInstance Win32_VideoController" + ' | Select-Object -ExpandProperty Name) -join "`n"' + ) + r = subprocess.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 10, + **_hidden_console_kwargs(), + ) + except (OSError, subprocess.SubprocessError) as e: + logger.debug("Live Windows adapter scan failed: %s", e) + return None + if r.returncode != 0: + return None + return [line.strip() for line in (r.stdout or "").splitlines() if line.strip()] + + +def _claim_live_adapter(name: Optional[str], live_names: list[str]) -> Optional[int]: + """Index of the first unclaimed live adapter this record matches, or None. + + Returned rather than a bool so the caller can consume the entry: see the + same-named-duplicates note at the call site. + """ + candidate = (name or "").strip().lower() + if not candidate: + return None + # Exact first: a stale "RX 7900 XT" record would otherwise claim a live "RX 7900 XTX". The + # prefix rule still has to exist, because the registry carries the driver description and + # WMI the display name. + for index, live in enumerate(live_names): + if live.strip().lower() == candidate: + return index + for index, live in enumerate(live_names): + if _adapter_name_is_live(name, [live]): + return index + return None + + +def _adapter_name_is_live(name: Optional[str], live_names: list[str]) -> bool: + """Whether a registry adapter name corresponds to one the live scan returned. + + Substring in either direction, because the two sources spell the same card + differently: the registry carries the driver's own description while WMI carries the + display name, and one is often a prefix of the other. setup.ps1 joins them the same + way, with Contains. + """ + candidate = (name or "").strip().lower() + if not candidate: + return False + return any( + candidate in live.lower() or live.lower() in candidate + for live in live_names + if live.strip() + ) + + +def _linux_drm_sysfs_records() -> list[Dict[str, Any]]: + """Every AMD or Intel card the DRM drivers have bound, from /sys/class/drm. + + ``vendor`` is 0x1002 for AMD/ATI and 0x8086 for Intel, and amdgpu's + ``mem_info_vram_total`` is a byte count, both documented kernel interfaces. Intel + publishes no equivalent total on the discrete path, so an Arc card is reported with + an unknown capacity rather than left out: knowing the GPU exists is the point. + Only ``cardN`` is walked: the connector entries (``card0-DP-1``) and render nodes + point at the same device and would double-count. + + No name is reported. The kernel publishes no marketing string here, and inventing + one, or borrowing an amd-smi row whose ordering is not guaranteed to match this + walk, would attach the wrong name to a card. A card with no name still tells the + user their GPU exists, which is the whole point. + + Never raises: an unreadable or malformed file skips that card. + """ + root = "/sys/class/drm" + try: + entries = sorted(os.listdir(root)) + except OSError: + return [] + records: list[Dict[str, Any]] = [] + for entry in entries: + if not re.fullmatch(r"card\d+", entry): + continue + device = os.path.join(root, entry, "device") + try: + with open(os.path.join(device, "vendor"), encoding = "utf-8") as fh: + vendor = fh.read().strip().lower() + except OSError: + continue + vendors = {"0x1002": "amd", "0x8086": "intel"} + if vendor not in vendors: + continue + total_gb = None + try: + with open(os.path.join(device, "mem_info_vram_total"), encoding = "utf-8") as fh: + total_bytes = int(fh.read().strip()) + total_gb = round(total_bytes / 1024**3, 2) if total_bytes > 0 else None + except (OSError, ValueError): + pass + records.append( + { + "vendor": vendors[vendor], + "index": len(records), + "name": None, + "memory_total_gb": total_gb, + "source": "sysfs-drm", + } + ) + return records + + +_UNKNOWN_PHYSICAL_GPU_INVENTORY: Dict[str, Any] = { + "available": False, + "devices": [], + "sources": [], + "unknown": True, + "unanswered": [], +} + + +def _carry_unanswered_vendors_forward( + inventory: Dict[str, Any], previous: Optional[Dict[str, Any]] +) -> Dict[str, Any]: + """Re-add the devices only a vendor this pass could not ask had reported. + + A refresh that cannot reach nvidia-smi is not news that the cards left, but caching + it as the new inventory says exactly that for a whole TTL: current_chat_only_verdict() + holds its frozen reason on ``unknown``, while _torch_gpu_mismatch_report() reads the + devices and finds none, so the Resources tab says no GPU on the same response where + the sidebar says PyTorch cannot use one. + + Per vendor, and only for vendors in ``unanswered``: a vendor that answered "none" gave + an answer, and a card that really was removed has to be allowed to disappear. + """ + stale_vendors = set(inventory.get("unanswered") or ()) + if not stale_vendors or not previous: + return inventory + carried = [ + device + for device in (previous.get("devices") or []) + if device.get("vendor") in stale_vendors + ] + if not carried: + return inventory + merged = dict(inventory) + merged["devices"] = list(inventory.get("devices") or []) + carried + merged["available"] = True + merged["sources"] = list(inventory.get("sources") or []) + [ + source + for source in (previous.get("sources") or []) + if source not in (inventory.get("sources") or []) + ] + # Still unknown: these rows describe the host as it was, not as this pass measured it. + return merged + + +def _run_physical_gpu_inventory_probe() -> Dict[str, Any]: + """One probe pass, stored in the cache. Never raises.""" + global _physical_gpu_inventory_cache + previous = _physical_gpu_inventory_cache + try: + inventory = _probe_physical_gpu_inventory() + except Exception as e: + logger.debug("Physical GPU inventory probe failed: %s", e) + inventory = dict(_UNKNOWN_PHYSICAL_GPU_INVENTORY) + # A pass that raised asked nobody, so every vendor is unanswered. + inventory["unanswered"] = ["amd", "intel", "nvidia"] + inventory = _carry_unanswered_vendors_forward( + inventory, previous[1] if previous is not None else None + ) + _physical_gpu_inventory_cache = (time.monotonic(), inventory) + return inventory + + +def get_physical_gpu_inventory(*, block: bool = True) -> Dict[str, Any]: + """GPUs the OS enumerates, whether or not PyTorch can use them. + + Returns ``{available, devices, sources, unknown}``. Runs regardless of ``DEVICE`` + and never raises: a probe that cannot answer contributes nothing rather than + failing the caller, so an endpoint reaching here still returns a body. + + ``block=False`` for anything on a request path. The NVIDIA half shells out with a + 10 second timeout and the lock makes concurrent callers queue behind it, so a hung + or restarting driver -- the exact host this inventory exists for -- would stall the + event loop for ten seconds every time the TTL expired. /api/health and + /api/liveness both reach here through the chat-only verdict, and the desktop gives + a local request two seconds. A non-blocking caller takes the last known answer, + kicks the refresh onto a daemon thread, and gets the new value on a later call; + with nothing cached at all it gets the explicit unknown, which every consumer + already reads as "keep what you had". + """ + now = time.monotonic() + cached_entry = _physical_gpu_inventory_cache + if cached_entry is not None and now - cached_entry[0] < _PHYSICAL_GPU_INVENTORY_TTL_SECONDS: + return cached_entry[1] + if not block: + _schedule_physical_gpu_inventory_refresh() + return ( + cached_entry[1] if cached_entry is not None else dict(_UNKNOWN_PHYSICAL_GPU_INVENTORY) + ) + with _physical_gpu_inventory_lock: + cached_entry = _physical_gpu_inventory_cache + if ( + cached_entry is not None + and time.monotonic() - cached_entry[0] < _PHYSICAL_GPU_INVENTORY_TTL_SECONDS + ): + return cached_entry[1] + return _run_physical_gpu_inventory_probe() + + +def _schedule_physical_gpu_inventory_refresh() -> None: + """Refresh the inventory off the caller's thread, one pass at a time. + + Single-flight through the same lock the blocking path uses, tried without waiting: + a refresh already running is exactly what a second caller wants, so it returns + rather than queueing a duplicate subprocess behind it. + """ + global _physical_gpu_inventory_refreshing + with _physical_gpu_inventory_refresh_lock: + if _physical_gpu_inventory_refreshing: + return + _physical_gpu_inventory_refreshing = True + + def _refresh() -> None: + global _physical_gpu_inventory_refreshing + try: + with _physical_gpu_inventory_lock: + _run_physical_gpu_inventory_probe() + except Exception as e: + logger.debug("Background inventory refresh failed: %s", e) + finally: + with _physical_gpu_inventory_refresh_lock: + _physical_gpu_inventory_refreshing = False + + try: + threading.Thread(target = _refresh, name = "gpu-inventory-refresh", daemon = True).start() + except Exception as e: + logger.debug("Could not start the inventory refresh thread: %s", e) + with _physical_gpu_inventory_refresh_lock: + _physical_gpu_inventory_refreshing = False + + +def _reported_torch_label(published: Optional[str] = None) -> Optional[str]: + """The version to show for this torch, without ever retrying a failed import. + + ``_torch_version_label()`` imports torch, and on a host whose native runtime will + not load that import is the thing that fails: it can take seconds, and _has_torch() + purges the partial module afterwards so the next attempt genuinely re-runs the + native load. /api/health and /api/liveness reach the label through the chat-only + verdict, so on exactly the host this feature exists for it would block the event + loop on every call. Whatever detection already published wins, then the label the + wheel carries on disk. + """ + if TORCH_IMPORT_ERROR is not None: + return published or _installed_torch_label_on_disk() or None + return _torch_version_label() or _installed_torch_label_on_disk() or None + + +def _torch_version_label() -> Optional[str]: + """``torch.__version__`` when it can be read, else None. Never raises.""" + try: + import torch + return str(torch.__version__) + except Exception: + return None + + +# A mask hides that vendor's devices and nothing else. +_VISIBILITY_MASK_VENDORS: Dict[str, frozenset] = { + # HIP honours CUDA_VISIBLE_DEVICES too, so an AMD-only box launched with it emptied is a + # deliberately masked host, not a broken one. + "CUDA_VISIBLE_DEVICES": frozenset({"nvidia", "amd"}), + "HIP_VISIBLE_DEVICES": frozenset({"amd"}), + "ROCR_VISIBLE_DEVICES": frozenset({"amd"}), + "ZE_AFFINITY_MASK": frozenset({"intel"}), +} + + +def _mask_is_emptied(var: str) -> bool: + """True when ``var`` is set to a value that hides every device it addresses. + + Set-but-empty and "-1" are the two spellings; a mask NAMING devices is not this, + since that host expects those devices to work. + """ + value = os.environ.get(var) + return value is not None and value.strip() in ("", "-1") + + +def _masks_hide_every_accelerator(*, block_inventory: bool = False) -> bool: + """True when the masks account for every accelerator this host has. + + Then torch reporting none is the configuration working, not a broken install, and + the whole classification is suppressed. A mask that covers only SOME of the cards + is not this: an emptied ZE_AFFINITY_MASK beside an unmasked NVIDIA GPU hides the + Arc and nothing else, and cancelling the classification there would let a CPU-only + wheel go unreported for a card the user never masked. Those cards are dropped from + the mismatch inventory instead. + + An inventory that found nothing, or could not answer, stays conservative: a mask + may well be hiding the only accelerator, and that is the case this existed for. + """ + masked = _vendors_masked_off(block_inventory = block_inventory) + if not masked: + return False + try: + inventory = get_physical_gpu_inventory(block = block_inventory) + except Exception: + inventory = dict(_UNKNOWN_PHYSICAL_GPU_INVENTORY) + devices = inventory.get("devices") or [] + if devices: + return all(device.get("vendor") in masked for device in devices) + # An inventory that could not ANSWER stays conservative; one that answered nothing has no + # card to hide. A cold non-blocking read is unknown, and an irrelevant empty mask would + # then cache "torch is fine" for a TTL, so the caching path passes block_inventory = True. + return True + + +def _vendors_masked_off(*, block_inventory: bool = False) -> set: + """Vendors whose devices are all hidden by a mask that can take effect here.""" + relevant = _relevant_visibility_masks(block_inventory = block_inventory) + masked: set = set() + for var in relevant: + if _mask_is_emptied(var): + masked |= _VISIBILITY_MASK_VENDORS.get(var, frozenset()) + # HIP reads its own variables FIRST and falls back to CUDA_VISIBLE_DEVICES only when + # neither is set (the precedence _get_parent_visible_gpu_spec applies). A host that + # NAMES devices in HIP_VISIBLE_DEVICES has not hidden them. + if any( + os.environ.get(var) is not None and not _mask_is_emptied(var) + for var in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES") + if var in relevant + ): + masked.discard("amd") + return masked + + +def _relevant_visibility_masks(*, block_inventory: bool = False) -> tuple[str, ...]: + """The visibility variables that can actually hide a GPU on THIS host. + + A mask that cannot take effect must not silence the mismatch. ROCR_VISIBLE_DEVICES + is the concrete case: Windows HIP has no ROCr layer, and this module's own + visibility resolver already ignores the variable there, so a stray empty one would + otherwise restore the "no GPU" verdict on a Windows NVIDIA host with a CPU wheel. + + The HIP-layer variables are likewise only consulted when an AMD card is actually + present, so a leftover HIP_VISIBLE_DEVICES on an NVIDIA-only box cannot mask + anything it does not address. An inventory that found nothing, or could not answer, + keeps every variable: unknown must stay conservative rather than start ignoring + masks that may well be real. + """ + masks = ["CUDA_VISIBLE_DEVICES"] + hip_masks = ["HIP_VISIBLE_DEVICES"] + if sys.platform != "win32": + hip_masks.append("ROCR_VISIBLE_DEVICES") + try: + # block=False: this is reached from the chat-only verdict, which /api/health and + # /api/liveness read. A vendor set does not change between refreshes. + devices = get_physical_gpu_inventory(block = block_inventory).get("devices") or [] + except Exception: + devices = [] + if not devices or any(d.get("vendor") == "amd" for d in devices): + masks.extend(hip_masks) + # ZE_AFFINITY_MASK is the XPU equivalent; the rest of this module already reads an + # emptied one as hiding every Intel device. + if not devices or any(d.get("vendor") == "intel" for d in devices): + masks.append("ZE_AFFINITY_MASK") + return tuple(masks) + + +def _torch_index_leaf(url: str) -> str: + """Final path segment of a torch index URL, lowercased, query and fragment removed. + + A token-authenticated pin is supported and common (.../whl/cpu?token=...), and a raw + final-segment split sees "cpu?token=..." there, so a deliberate CPU install on a GPU + host would be reported as broken and offered a repair. Trailing slashes come off the + PATH only, or a token ending in "/" is corrupted. Mirrors the installer's + _torch_index_leaf / _trim_index_path_slashes, which the same variable is read by. + """ + value = str(url).strip().lower() + if not value: + return "" + path = re.fullmatch(r"([^?#]*)([?#].*)?", value) + path = path.group(1) if path else value + return path.rstrip("/").rsplit("/", 1)[-1] + + +def _expected_cpu_flavor_was_chosen() -> bool: + """Whether THIS install deliberately selected a CPU wheel. + + Two sources, both meaning "the user named cpu": an explicit index pin set for the + running process, and the flavor the last completed install recorded in the venv's + manifest, which install_python_stack.py writes from the same expectation the + Windows flavor invariant enforces. + + Only "cpu" is acted on. An unknown or absent record means nothing was recorded and + must not be read as a choice, which is the rule install_manifest.recorded_torch_flavor + documents. Read straight off disk rather than by importing that module, which lives + outside the backend package. Never raises. + """ + for var in ("UNSLOTH_TORCH_INDEX_FAMILY", "UNSLOTH_TORCH_INDEX_URL"): + if _torch_index_leaf(os.environ.get(var) or "") == "cpu": + return True + try: + path = os.path.join(sys.prefix, "unsloth_install_manifest.json") + with open(path, encoding = "utf-8") as fh: + manifest = json.load(fh) + recorded = manifest.get("expected_torch_tag") + pinned = manifest.get("expected_torch_tag_pinned") + except (OSError, ValueError, AttributeError): + return False + if not isinstance(recorded, str) or recorded.strip().lower() != "cpu": + return False + # A recorded cpu is not by itself a choice: setup.ps1 selects /cpu automatically on a + # GPU-less host and records it exactly as it records a pinned one. + return bool(pinned) + + +def classify_torch_build(*, block_inventory: bool = False) -> Optional[str]: + """Why this PyTorch exposes no accelerator, when the build itself is the reason. + + "torch_cpu_build" -- a CPU-only wheel: ``2.11.0+cpu``, or an untagged build + with neither ``torch.version.cuda`` nor ``.hip``. No + driver update fixes it; only reinstalling torch from + the right index does. + "torch_cuda_unavailable" -- an accelerator wheel (``+cu124``, ``+rocm6.4``, + ``+xpu``) whose runtime refuses to initialise: driver + too old, no permission on the device nodes, a cudart + that will not load. The wheel is right, the + environment is not. + None -- torch is missing, unimportable, or has a working + accelerator. Nothing to say. + + Kept as two reasons rather than one flag because the advice differs, and telling + someone with a healthy cu124 wheel to reinstall torch sends them the wrong way. + Never raises. + """ + # An emptied visibility mask has the exact shape of a broken install without anything + # being broken. A mask NAMING devices is not this: that host expects those to work. + if _masks_hide_every_accelerator(block_inventory = block_inventory): + return None + if _expected_cpu_flavor_was_chosen(): + return None + if not _has_torch(): + # Absent and unimportable differ, and the second is what this exists for. _has_torch() + # collapses both to False. + return _classification_from_disk_label() + try: + import torch + + # XPU as well as CUDA, or a recovered Intel host stays torch_cuda_unavailable forever. + for _available in ( + getattr(getattr(torch, "cuda", None), "is_available", None), + getattr(getattr(torch, "xpu", None), "is_available", None), + ): + try: + if callable(_available) and _available(): + return None + except Exception: + continue + version = str(getattr(torch, "__version__", "")) + # The installers' vocabulary: "2.11.0+cpu" -> "cpu", "2.6.0+cu124" -> "cu124". + local = version.partition("+")[2].strip().lower() + cuda_tag = getattr(getattr(torch, "version", None), "cuda", None) + hip_tag = getattr(getattr(torch, "version", None), "hip", None) + # An untagged wheel that sets torch.version.xpu is a GPU build whose runtime is down. + xpu_tag = getattr(getattr(torch, "version", None), "xpu", None) + if local == "cpu" or local.startswith("cpu."): + # PyTorch publishes extended CPU local tags such as "2.8.0+cpu.cxx11.abi". + return "torch_cpu_build" + if not local and cuda_tag is None and hip_tag is None and xpu_tag is None: + # Untagged and built against no GPU runtime: the PyPI macOS/CPU wheel shape. An untagged + # wheel that DOES set version.cuda (conda) belongs in the second case. + return "torch_cpu_build" + return "torch_cuda_unavailable" + except Exception as e: + # torch is INSTALLED but will not import. Classify from the wheel on disk instead. + logger.debug("torch build classification fell back to the on-disk label: %s", e) + return _classification_from_disk_label() + + +def _classification_from_disk_label() -> Optional[str]: + """Classify from the wheel's own version label, with no interpreter started. + + ``None`` when nothing is installed to read: an absent torch is not a mismatch. + """ + label = _installed_torch_label_on_disk() + markers = _installed_torch_markers_on_disk() + if not label and not any(markers.values()): + return None + if "+cu" in label or "+rocm" in label or "+xpu" in label: + return "torch_cuda_unavailable" + if any(markers.values()): + return "torch_cuda_unavailable" + return "torch_cpu_build" + + +# setup.ps1's rule: only Arc and Data Center GPU autodetect as XPU. An ordinary Intel +# iGPU is a display adapter, not a training device. +_XPU_ADAPTER_NAME_RE = re.compile(r"intel.*(arc|data center gpu)", re.IGNORECASE) + + +def _devices_that_can_establish_a_mismatch(devices: list[Dict[str, Any]]) -> list[Dict[str, Any]]: + """The subset of the inventory whose presence means PyTorch OUGHT to have a GPU. + + NVIDIA and AMD qualify outright. Intel does not, by itself: setup.sh does not + autodetect Linux XPU at all, and setup.ps1 limits automatic XPU detection to Arc and + Data Center GPU by name. An Intel UHD iGPU beside a CPU wheel is the expected, + correct state of that machine, so counting it would report a mismatch and offer a + repair that reinstalls the very CPU build it just replaced. + + An Intel card still counts when it is one this stack would have installed XPU torch + for: a name that matches setup.ps1's rule, or an XPU expectation recorded for this + venv, or a torch that carries an XPU runtime. The Linux sysfs walk publishes no + name, which is exactly why the expectation and the runtime are consulted too. + """ + xpu_expected = _expected_xpu_flavor_was_chosen() or _torch_reports_an_xpu_runtime() + # Per vendor, not host-wide: an emptied ZE_AFFINITY_MASK beside an unmasked NVIDIA + # card hides the Arc and nothing else. + masked_off = _vendors_masked_off() + # A ROCm expectation, or a torch carrying a HIP runtime, settles it for every AMD + # card: this stack asked for that wheel, whatever the arch table says. + rocm_expected = _expected_rocm_flavor_was_chosen() or _torch_reports_a_hip_runtime() + keep: list[Dict[str, Any]] = [] + for device in devices: + if device.get("vendor") in masked_off: + continue + if device.get("vendor") == "amd": + if rocm_expected or _amd_device_can_establish_a_mismatch(device): + keep.append(device) + continue + if device.get("vendor") != "intel": + keep.append(device) + continue + if xpu_expected or _XPU_ADAPTER_NAME_RE.search(str(device.get("name") or "")): + keep.append(device) + return keep + + +# The gfx targets this stack will actually install a ROCm wheel for: install.sh's +# _amd_arch_index_family_for_gfx, plus gfx906 from the ROCm 6.3 path. A card outside this +# set (Polaris gfx803, RDNA 1 gfx101x) is left on CPU torch ON PURPOSE. +_ROCM_SUPPORTED_GFX = frozenset( + { + "gfx906", + "gfx908", + "gfx90a", + "gfx1030", + "gfx1031", + "gfx1032", + "gfx1033", + "gfx1034", + "gfx1035", + "gfx1036", + "gfx1100", + "gfx1101", + "gfx1102", + "gfx1103", + "gfx1150", + "gfx1151", + "gfx1152", + "gfx1200", + "gfx1201", + } +) + + +def _linux_kfd_reports_an_amd_gpu() -> bool: + """Whether the KFD topology enumerates an AMD GPU node. Never raises. + + The same probe, and the same vendor guard, install_python_stack._has_rocm_gpu() + uses: gpu_id 0 is a CPU node, and the NVIDIA open kernel module registers KFD nodes + of its own with vendor_id 4318, so AMD ownership has to be confirmed rather than + assumed. A node whose properties cannot be read is skipped for that reason. + """ + if platform.system() != "Linux": + return False + nodes = "/sys/class/kfd/kfd/topology/nodes" + try: + entries = os.listdir(nodes) + except OSError: + return False + for entry in entries: + try: + with open(os.path.join(nodes, entry, "gpu_id"), encoding = "utf-8") as fh: + gpu_id = fh.read().strip() + except (OSError, UnicodeDecodeError): + continue + if not gpu_id or gpu_id == "0": + continue + try: + with open(os.path.join(nodes, entry, "properties"), encoding = "utf-8") as fh: + properties = fh.read() + except (OSError, UnicodeDecodeError): + continue + if re.search(r"\bvendor_id\s+4098\b", properties): + return True + return False + + +def _is_pip_rocm_family_leaf(leaf: str) -> bool: + """True when a lowercased leaf names a pip ROCm family: EXACTLY rocm[.] + or gfx. install_python_stack._is_pip_rocm_family_leaf, kept in step with it. + + A suffixed leaf (rocm-rel-7.2.1, gfx-mirror) starts with the same letters but is a + custom pin the installer routes verbatim and never treats as a ROCm choice. Reading + one as ROCm here waives the supported-architecture filter, so a gfx803 host that was + deliberately left on CPU torch gets told its own install is broken. + """ + return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf)) + + +def _expected_rocm_flavor_was_chosen() -> bool: + """Whether this install selected a ROCm wheel, by pin or by recorded flavor.""" + for var in ("UNSLOTH_TORCH_INDEX_FAMILY", "UNSLOTH_TORCH_INDEX_URL"): + if _is_pip_rocm_family_leaf(_torch_index_leaf(os.environ.get(var) or "")): + return True + try: + path = os.path.join(sys.prefix, "unsloth_install_manifest.json") + with open(path, encoding = "utf-8") as fh: + recorded = json.load(fh).get("expected_torch_tag") + except (OSError, ValueError, AttributeError): + return False + return isinstance(recorded, str) and recorded.strip().lower().startswith("rocm") + + +def _torch_reports_a_hip_runtime() -> bool: + """Whether the installed torch is a ROCm build, however unusable it currently is.""" + if TORCH_IMPORT_ERROR is not None: + return "+rocm" in _installed_torch_label_on_disk() or bool( + _installed_torch_markers_on_disk()["hip"] + ) + try: + import torch + if "+rocm" in str(getattr(torch, "__version__", "")).lower(): + return True + return getattr(getattr(torch, "version", None), "hip", None) is not None + except Exception: + return False + + +# Marketing name -> gfx, mirroring setup.ps1's $nameArchTable and +# install_python_stack._WIN_GPU_NAME_ARCH_TABLE. Only names those two route to a wheel +# family: this decides whether a repair could change anything, so a card no index covers +# must not match. Most specific first. +_GPU_NAME_GFX_TABLE: "list[tuple[str, str]]" = [ + (r"9070|9080|R9700", "gfx1201"), + (r"9060", "gfx1200"), + (r"8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"), + (r"890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]", "gfx1150"), + (r"860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", "gfx1152"), + (r"RX 7900|PRO W7900|PRO W7800", "gfx1100"), + (r"RX 7800|RX 7700(?!S)|PRO W7700|PRO V710", "gfx1101"), + (r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500", "gfx1102"), + (r"780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme", "gfx1103"), + (r"RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900", "gfx1030"), + (r"RX 6650|RX 6600|PRO W6600|PRO W6650", "gfx1032"), + (r"RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500", "gfx1034"), +] + + +def _rocm_supported_gfx_from_gpu_name(name: str) -> Optional[str]: + """The gfx arch this marketing name maps to, when the ROCm wheels cover it.""" + if not name: + return None + for pattern, arch in _GPU_NAME_GFX_TABLE: + if re.search(pattern, name, re.IGNORECASE) and arch in _ROCM_SUPPORTED_GFX: + return arch + return None + + +def _amd_device_can_establish_a_mismatch(device: Dict[str, Any]) -> bool: + """Whether this AMD adapter is one the installers would have given a ROCm wheel. + + A stack that deliberately declines to support a card cannot then call the CPU wheel + beside it a fault. The arch is consulted only when something could name it: the DRM + sysfs walk publishes no name, so a host with no ROCm userspace at all reports no + candidates and stays counted, which is the conservative direction and the shape the + original report had (an RX 7900 XT, which IS supported). + """ + candidates = [ + str(gfx).lower() + for gfx in ( + device.get("gfx_candidates") or ([device.get("gfx")] if device.get("gfx") else []) + ) + if gfx + ] + if not candidates: + # The DirectX registry publishes AdapterFamily only when the driver wrote one, so a + # supported Windows card (the reported RX 7900 XT) can arrive with no arch at all. + # setup.ps1 answers that from the marketing name, so answer it the same way before + # falling through to a probe that can only speak for Linux. + _named = _rocm_supported_gfx_from_gpu_name(device.get("name") or "") + if _named: + return True + # Nothing NAMES the card, so ask what the installer asks: _has_rocm_gpu() falls back to + # the KFD topology, and a card `studio update` would repair is eligible here. + return _linux_kfd_reports_an_amd_gpu() + return any(gfx in _ROCM_SUPPORTED_GFX for gfx in candidates) + + +def _expected_xpu_flavor_was_chosen() -> bool: + """Whether this install selected an XPU wheel, by pin or by recorded flavor.""" + for var in ("UNSLOTH_TORCH_INDEX_FAMILY", "UNSLOTH_TORCH_INDEX_URL"): + if _torch_index_leaf(os.environ.get(var) or "") == "xpu": + return True + try: + path = os.path.join(sys.prefix, "unsloth_install_manifest.json") + with open(path, encoding = "utf-8") as fh: + recorded = json.load(fh).get("expected_torch_tag") + except (OSError, ValueError, AttributeError): + return False + return isinstance(recorded, str) and recorded.strip().lower() == "xpu" + + +def _torch_reports_an_xpu_runtime() -> bool: + """Whether the installed torch is an XPU build, however unusable it currently is.""" + if TORCH_IMPORT_ERROR is not None: + # Asking costs a full torch/__init__ against the partial module tree the import left. + return "+xpu" in _installed_torch_label_on_disk() or bool( + _installed_torch_markers_on_disk()["xpu"] + ) + try: + import torch + if "+xpu" in str(getattr(torch, "__version__", "")).lower(): + return True + return getattr(getattr(torch, "version", None), "xpu", None) is not None + except Exception: + return False + + +def _installed_torch_label_on_disk() -> str: + """``torch.__version__`` read out of the installed torch/version.py, or "". + + No interpreter is started, which is the point: this is reached when importing torch + is the thing that fails. The installers read the same file for the same reason. + Never raises. + """ + try: + spec = importlib.util.find_spec("torch") + except Exception: + return "" + locations = list(getattr(spec, "submodule_search_locations", None) or []) if spec else [] + for location in locations: + try: + with open(os.path.join(location, "version.py"), encoding = "utf-8") as fh: + for line in fh: + if line.startswith("__version__"): + return line.partition("=")[2].strip().strip("\"'").lower() + except OSError: + continue + return "" + + +def _installed_torch_markers_on_disk() -> Dict[str, Optional[str]]: + """``{cuda, hip, xpu}`` as recorded in the installed torch/version.py. + + The version LABEL is not the whole story: a conda or source CUDA build is untagged + and records its runtime here instead, and the importable path already reads exactly + these three attributes. Without them the failure path gave the same installation the + opposite diagnosis, telling the user to reinstall a GPU wheel it already has rather + than to fix the driver. Parsed, not executed, for the same reason as the label. + Never raises; a value the file does not set reads as None. + """ + markers: Dict[str, Optional[str]] = {"cuda": None, "hip": None, "xpu": None} + try: + spec = importlib.util.find_spec("torch") + except Exception: + return markers + locations = list(getattr(spec, "submodule_search_locations", None) or []) if spec else [] + for location in locations: + try: + with open(os.path.join(location, "version.py"), encoding = "utf-8") as fh: + source = fh.read() + except OSError: + continue + try: + tree = ast.parse(source) + except SyntaxError: + continue + for node in tree.body: + # torch has shipped both `cuda = '12.8'` and `cuda: Optional[str] = '12.8'`. + if isinstance(node, ast.Assign): + names = [n.id for n in node.targets if isinstance(n, ast.Name)] + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + names = [node.target.id] + else: + continue + value = getattr(node, "value", None) + if not isinstance(value, ast.Constant) or not isinstance(value.value, str): + continue + for name in names: + if name in markers: + markers[name] = value.value + return markers + return markers + + +def _run_torch_build_snapshot() -> Dict[str, Any]: + """One uncached pass over torch, cached on the way out. Never raises.""" + global _torch_build_snapshot_cache + snapshot = { + # block_inventory: never a request path, so masks compare against a MEASURED inventory. + "reason": classify_torch_build(block_inventory = True), + "usable": _torch_reports_a_usable_accelerator(), + "unknown": False, + } + _torch_build_snapshot_cache = (time.monotonic(), snapshot) + return snapshot + + +def torch_build_snapshot(*, block: bool = True) -> Dict[str, Any]: + """``{reason, usable, unknown}`` for this venv's torch, cached with a TTL. + + ``block=False`` for anything on a request path: both probes inside import torch and + ask the CUDA and XPU runtimes whether they are available, and a wedged or + restarting driver -- the very state ``torch_cuda_unavailable`` names -- can hold + those for as long as the driver takes. /api/health and /api/liveness reach here + through the chat-only verdict, against a two second desktop timeout, so a + non-blocking caller takes the last measurement, kicks the refresh onto a daemon + thread, and reads the new one later. + + With nothing measured yet it gets ``unknown``, never a guess: detect_hardware() + takes the blocking path on exactly the hosts whose verdict can be re-derived, so + the cache is warm before any request can consult it. + """ + now = time.monotonic() + cached_entry = _torch_build_snapshot_cache + if cached_entry is not None and now - cached_entry[0] < _TORCH_BUILD_SNAPSHOT_TTL_SECONDS: + return cached_entry[1] + if not block: + _schedule_torch_build_snapshot_refresh() + return cached_entry[1] if cached_entry is not None else dict(_UNKNOWN_TORCH_BUILD_SNAPSHOT) + with _torch_build_snapshot_lock: + cached_entry = _torch_build_snapshot_cache + if ( + cached_entry is not None + and time.monotonic() - cached_entry[0] < _TORCH_BUILD_SNAPSHOT_TTL_SECONDS + ): + return cached_entry[1] + return _run_torch_build_snapshot() + + +def _schedule_torch_build_snapshot_refresh() -> None: + """Refresh the torch snapshot off the caller's thread, one pass at a time.""" + global _torch_build_snapshot_refreshing + with _torch_build_snapshot_refresh_lock: + if _torch_build_snapshot_refreshing: + return + _torch_build_snapshot_refreshing = True + + def _refresh() -> None: + global _torch_build_snapshot_refreshing + try: + with _torch_build_snapshot_lock: + _run_torch_build_snapshot() + except Exception as e: + logger.debug("Background torch build refresh failed: %s", e) + finally: + with _torch_build_snapshot_refresh_lock: + _torch_build_snapshot_refreshing = False + + try: + threading.Thread(target = _refresh, name = "torch-build-refresh", daemon = True).start() + except Exception as e: + logger.debug("Could not start the torch build refresh thread: %s", e) + with _torch_build_snapshot_refresh_lock: + _torch_build_snapshot_refreshing = False + + +def _seed_torch_build_snapshot(reason: Optional[str]) -> None: + """Record a classification that was reached without probing torch. + + The broken-runtime host is classified from the wheel on disk, and that answer is + exactly what the request paths need; letting them re-measure would put the import + that already failed back on the health thread. + """ + global _torch_build_snapshot_cache + _torch_build_snapshot_cache = ( + time.monotonic(), + {"reason": reason, "usable": False, "unknown": False}, + ) + + +def invalidate_torch_build_snapshot() -> None: + """Drop the cached measurement so the next blocking caller re-probes.""" + global _torch_build_snapshot_cache + _torch_build_snapshot_cache = None + + +def _mismatch_verdict_for_this_host( + reason: Optional[str] = None, +) -> tuple[Optional[str], Optional[str]]: + """``(reason, detail)`` when this host's GPUs are real but PyTorch cannot use them. + + ``(None, None)`` otherwise. Blocking, and only detection calls it: it runs off the + request path, and the blocking pass is what warms the caches that /api/health then + reads without probing anything itself. + """ + if reason is None: + reason = torch_build_snapshot()["reason"] + if reason is None: + return None, None + if not _devices_that_can_establish_a_mismatch( + get_physical_gpu_inventory().get("devices") or [] + ): + return None, None + detail = _reported_torch_label() + logger.warning( + "GPUs are present on this host but PyTorch cannot use them (%s%s); " + "Train/Export disabled (chat-only). Repair the installation to restore GPU " + "support.", + reason, + f", installed {detail}" if detail else "", + ) + return reason, detail + + +def _torch_gpu_mismatch_report() -> Dict[str, Any]: + """``physical_devices`` + ``mismatch`` for a host whose GPUs PyTorch cannot use. + + ``{}`` when there is nothing to report: torch is absent, unimportable or healthy, + or no probe found a physical card. Both keys sit BESIDE ``devices`` in the + visibility payload and never inside it -- ``devices`` is the runtime-usable list + that model fit budgets against and that the training device picker pins from. + """ + # block=False: request path, and the torch probes hold as long as a wedged driver. + reason = torch_build_snapshot(block = False)["reason"] + if reason is None: + return {} + # block=False: GET /api/system holds _system_gpu_cache_lock for the whole call. + inventory = get_physical_gpu_inventory(block = False) + physical = _devices_that_can_establish_a_mismatch(inventory.get("devices") or []) + if not physical: + return {} + return { + "physical_devices": physical, + "mismatch": { + "reason": reason, + "torch_version": _reported_torch_label(CHAT_ONLY_DETAIL), + "physical_count": len(physical), + "sources": inventory.get("sources") or [], + }, + } + + def verdict_pending_mlx_repair(chat_only: bool, reason: Optional[str]) -> bool: """True when this settled verdict is one the MLX self-heal is about to overturn. @@ -479,6 +1638,9 @@ def detect_hardware() -> DeviceType: return device DETECTION_GENERATION += 1 DETECTION_COMPLETE.set() + # A pass has settled, so a later recovery may ask for another one. + global _REDETECTION_REQUESTED + _REDETECTION_REQUESTED = False return device @@ -496,6 +1658,7 @@ def ensure_hardware_detected(epoch: Optional[int] = None) -> DeviceType: reading it here would bind the pass to the retirement it must lose to. Direct callers pass nothing and own the current epoch.""" global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, CHAT_ONLY_DETAIL, DETECTION_GENERATION + global _REDETECTION_REQUESTED with _DETECT_LOCK: if epoch is None: # A nested read inside an owning scope belongs to that pass, not to whatever @@ -540,6 +1703,11 @@ def ensure_hardware_detected(epoch: Optional[int] = None) -> DeviceType: # waiter trusting it could publish training-enabled for a CPU/chat-only host. # Unconditional, unlike the counter: re-setting is a no-op and a late waiter needs it. DETECTION_COMPLETE.set() + if produced_here: + # A pass has settled, so a later recovery may ask for another one. Cleared here as well as + # in detect_hardware(), because recovery starts its pass through this function. After the + # epoch check, so a retired pass does not release a guard it did not satisfy. + _REDETECTION_REQUESTED = False return DEVICE @@ -663,10 +1831,25 @@ def _detect_hardware_locked() -> DeviceType: elif TORCH_IMPORT_ERROR is not None: # torch installed but broken, so this host was never measured. "no_gpu" would lie. CHAT_ONLY_REASON = "detection_failed" + # Still measurable from the wheel on disk, and this host needs it most: detection_failed + # otherwise sends the user to the server log instead of offering the repair, and the + # verdict refresh deliberately freezes it. From DISK, because the import is what failed. + _disk_reason = _classification_from_disk_label() + _seed_torch_build_snapshot(_disk_reason) + _build_reason, _build_detail = _mismatch_verdict_for_this_host(_disk_reason) + if _build_reason is not None: + CHAT_ONLY_REASON, CHAT_ONLY_DETAIL = _build_reason, _build_detail elif platform.system() == "Darwin": CHAT_ONLY_REASON = "intel_mac" # Intel Mac: no PyTorch/MLX -> GGUF-only by design. else: + # torch imported cleanly and reported no accelerator, which is NOT "this host has no GPU": + # a Windows update installs PyPI's 2.11.0+cpu over cu124 while nvidia-smi lists every + # card (#8473). Ask the OS before blaming hardware. CHAT_ONLY_REASON = "no_gpu" + if torch_ok: + _build_reason, _build_detail = _mismatch_verdict_for_this_host() + if _build_reason is not None: + CHAT_ONLY_REASON, CHAT_ONLY_DETAIL = _build_reason, _build_detail print("Hardware detected: CPU training backend (no PyTorch/MLX GPU backend available)") return DEVICE @@ -682,6 +1865,139 @@ def get_device() -> DeviceType: return ensure_hardware_detected() +def _torch_reports_a_usable_accelerator() -> bool: + """Whether torch can open a GPU right now. Never raises.""" + try: + import torch + for probe in ( + getattr(getattr(torch, "cuda", None), "is_available", None), + getattr(getattr(torch, "xpu", None), "is_available", None), + ): + try: + if callable(probe) and probe(): + return True + except Exception: + continue + except Exception: + return False + return False + + +def _request_hardware_redetection() -> None: + """Ask for a fresh detection pass, at most one per recovery. Never raises. + + Retiring the epoch is what the rest of this module already uses to mean "the + published verdict is stale"; recomputing DEVICE here instead would publish from a + request thread and race the detection lock. + """ + global _REDETECTION_REQUESTED + if _REDETECTION_REQUESTED: + return + try: + _REDETECTION_REQUESTED = True + # Retiring the epoch alone is NOT enough: invalidate_detection leaves DEVICE set and + # DETECTION_COMPLETE raised, so start_background_detection returns immediately. Both under + # the lock, so a running pass cannot publish over the reset. + invalidate_detection() + invalidate_torch_build_snapshot() + with _DETECT_LOCK: + _discard_detection_locked() + start_background_detection() + logger.info( + "An accelerator became usable after startup; discarded the cached hardware " + "verdict and started a fresh detection pass." + ) + except Exception as e: + _REDETECTION_REQUESTED = False + logger.debug("Could not request hardware re-detection: %s", e) + + +def current_chat_only_verdict() -> tuple[Optional[str], Optional[str]]: + """``(reason, detail)``, re-derived when the physical inventory can still change it. + + detect_hardware() runs once at startup, but the inventory it consulted refreshes on + a 60 second TTL. An eGPU attached after launch, or a driver that finished restarting + after the first probe, flips the answer while the frozen verdict keeps saying + ``no_gpu``: /api/system would list the card and publish a mismatch while the sidebar + and the Export and Video pages went on insisting no accelerator exists. The reverse + is the same bug -- a card that goes away leaves a mismatch nobody can act on. + + Only the three inventory-sensitive verdicts are re-derived, plus the one + detection_failed that is not really unmeasured: a torch that will not import was + classified from its wheel on disk at startup, and if the OS probe had also not + answered yet then the inventory is the ONLY thing still missing. Freezing that host + is the same split this function exists to close, with the repair guidance as the + thing it withholds. mlx_unavailable, intel_mac and a detection_failed with no + importable torch and no readable wheel describe things a 60 second probe cannot + change, and re-deriving those would fight detect_hardware() rather than follow it. + + Never raises: a probe that cannot answer keeps the frozen verdict. + """ + reason, detail = CHAT_ONLY_REASON, CHAT_ONLY_DETAIL + frozen_but_measurable = reason == "detection_failed" and TORCH_IMPORT_ERROR is not None + if reason not in ("no_gpu", "torch_cpu_build", "torch_cuda_unavailable"): + if not frozen_but_measurable: + return reason, detail + try: + snapshot = torch_build_snapshot(block = False) + if snapshot["unknown"]: + # Nothing measured yet is not "torch is fine"; keep what detection published. + return reason, detail + build_reason = snapshot["reason"] + if build_reason is None and snapshot["usable"]: + # The accelerator came BACK. Only reason and detail refresh here, so DEVICE and CHAT_ONLY + # would stay frozen at CPU while this reported no_gpu, until a restart. Retire the + # detection instead and let the next pass publish the real answer. + _request_hardware_redetection() + return reason, detail + # block=False: /api/health and /api/liveness reach here, and the NVIDIA half shells out + # with a 10 second timeout on exactly the hung-driver host this exists for. + inventory = get_physical_gpu_inventory(block = False) + if build_reason is not None and _devices_that_can_establish_a_mismatch( + inventory.get("devices") or [] + ): + return build_reason, _reported_torch_label(detail) + if inventory.get("unknown"): + return reason, detail + except Exception as e: + logger.debug("chat-only verdict refresh failed: %s", e) + return reason, detail + # "No GPU here" is a measurement, and a host whose torch will not import never made + # one: detect_hardware() refused to say it, so this cannot say it either. + return (reason, detail) if frozen_but_measurable else ("no_gpu", None) + + +def _gpu_present_but_unusable_message( + feature: str, verdict: Optional[tuple[Optional[str], Optional[str]]] = None +) -> Optional[str]: + """The capability message for a host whose GPUs are real but unreachable by torch. + + ``None`` when this host is not in that state. detect_hardware() records + ``torch_cpu_build`` / ``torch_cuda_unavailable`` only after the OS inventory has + actually found a card, so reaching this point means the "no supported accelerator + was found" wording below would contradict the System tab and send the user after + hardware they already own. Both reasons are surfaced verbatim by the Export and + Video pages and by rejected export API calls. + """ + # ``verdict`` lets one response read it ONCE: two reads across a TTL boundary can describe + # different hosts. + reason, detail = verdict if verdict is not None else current_chat_only_verdict() + if reason not in ("torch_cpu_build", "torch_cuda_unavailable"): + return None + installed = f" (installed {detail})" if detail else "" + if reason == "torch_cpu_build": + return ( + f"This host has a GPU, but the installed PyTorch is a CPU-only build{installed}, " + f"so {feature} cannot use it. Repair the installation from Settings to reinstall " + f"the GPU build." + ) + return ( + f"This host has a GPU, but the installed PyTorch{installed} cannot initialise it, so " + f"{feature} cannot use it. This is usually a driver or runtime mismatch; repairing the " + f"installation from Settings reinstalls a matching PyTorch build." + ) + + def export_capability() -> dict: """Whether model export can run here, with a torch-aware reason when it cannot. @@ -697,9 +2013,10 @@ def export_capability() -> dict: "export_unsupported_reason": None, "export_unsupported_message": None, } + verdict = current_chat_only_verdict() # No accelerator: name the blocker. Detection failure first -- the branches below all # describe a measured host, so a broken probe would tell a GPU box to install PyTorch. - if CHAT_ONLY_REASON == "detection_failed": + if verdict[0] == "detection_failed": reason = "detection_failed" message = ( "Hardware detection failed on this host, so export is disabled. The server log records " @@ -711,6 +2028,11 @@ def export_capability() -> dict: "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " "`unsloth studio update` to restore MLX and enable export." ) + elif _gpu_present_but_unusable_message("export", verdict) is not None: + # BEFORE _has_torch(), which reports an unimportable wheel as absent while re-running the + # seconds-long import that already failed. + reason = verdict[0] + message = _gpu_present_but_unusable_message("export", verdict) elif not _has_torch(): reason = "pytorch_not_installed" message = ( @@ -749,7 +2071,8 @@ def video_capability() -> dict: } # Detection failure first, as in export_capability: the branches below all describe a # measured host, so a broken probe would tell a GPU box to go buy a GPU. - if CHAT_ONLY_REASON == "detection_failed": + verdict = current_chat_only_verdict() + if verdict[0] == "detection_failed": reason = "detection_failed" message = ( "Hardware detection failed on this host, so video generation is disabled. The server " @@ -792,6 +2115,9 @@ def video_capability() -> dict: "Video generation requires Apple Silicon. This Intel Mac has no Metal (MPS) device " "for the video pipelines to run on." ) + elif _gpu_present_but_unusable_message("video generation", verdict) is not None: + reason = verdict[0] + message = _gpu_present_but_unusable_message("video generation", verdict) elif not _has_torch(): reason = "pytorch_not_installed" message = ( @@ -2015,6 +3341,7 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: # 0x1002. NVIDIA's open kernel module also registers KFD nodes (vendor_id 0x10DE); # a non-AMD node is not a HIP device and must never take an ordinal. _AMD_PCI_VENDOR_ID = 4098 +_INTEL_PCI_VENDOR_ID = 0x8086 def _rocm_kfd_gpu_pci_ids() -> list[str]: @@ -2259,8 +3586,14 @@ def _parse_adapter_family_gfx(family: str) -> str: return "" -def _windows_amd_adapter_records_by_luid() -> dict[int, Dict[str, Any]]: - """DirectX registry metadata for AMD adapters, keyed by LUID. +def _windows_amd_adapter_records_by_luid( + vendor_id_filter: int = _AMD_PCI_VENDOR_ID, *, distinguish_failure: bool = False +) -> "dict[int, Dict[str, Any]] | None": + """DirectX registry metadata for one vendor's adapters, keyed by LUID. + + ``vendor_id_filter`` defaults to AMD, which is every pre-existing caller. The + physical inventory passes Intel's id as well, so an Arc host whose XPU wheel was + replaced is reported rather than going silent. ``gfx`` is absent when the driver wrote no ``AdapterFamily``. ``dedicated_memory_bytes`` is absent when neither dedicated-memory value is available. @@ -2268,16 +3601,31 @@ def _windows_amd_adapter_records_by_luid() -> dict[int, Dict[str, Any]]: All or nothing: a record this cannot read makes the map incomplete, and an incomplete map is indistinguishable from a complete one at the join, which would then pair a visible card with a hidden same-named card's counter. So - any failure past the point where a subkey is known to be an adapter returns - ``{}``, which drops the caller back to capacity ranking. Same for off - Windows or without the key. + any failure past the point where a subkey is known to be an adapter gives up + on the whole map, which drops the caller back to capacity ranking. Same for + off Windows or without the key. + + The ranking callers only need "no usable map", so they get ``{}``. + ``distinguish_failure`` returns ``None`` for that case instead: the inventory has + to tell a vendor with no adapters from a vendor it could not ask, because + publishing the first for the second erases a settled mismatch. """ - if platform.system() != "Windows": + records = _windows_amd_adapter_records_or_none(vendor_id_filter) + if records is None and not distinguish_failure: return {} + return records + + +def _windows_amd_adapter_records_or_none( + vendor_id_filter: int = _AMD_PCI_VENDOR_ID, +) -> "dict[int, Dict[str, Any]] | None": + """The read itself. ``None`` whenever the registry could not answer.""" + if platform.system() != "Windows": + return None try: import winreg except ImportError: - return {} + return None by_luid: dict[int, Dict[str, Any]] = {} try: with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DIRECTX_KEY) as dx_key: @@ -2289,7 +3637,7 @@ def _windows_amd_adapter_records_by_luid() -> dict[int, Dict[str, Any]]: continue with winreg.OpenKey(dx_key, subkey) as adapter_key: vendor_id, _ = winreg.QueryValueEx(adapter_key, "VendorId") - if int(vendor_id) != _AMD_PCI_VENDOR_ID: + if int(vendor_id) != vendor_id_filter: continue luid, _ = winreg.QueryValueEx(adapter_key, "AdapterLuid") description, _ = winreg.QueryValueEx(adapter_key, "Description") @@ -2311,8 +3659,7 @@ def _windows_amd_adapter_records_by_luid() -> dict[int, Dict[str, Any]]: pass name = str(description).strip() if not name: - # An AMD adapter this cannot name: see the all-or-nothing note. - return {} + return None record = {"name": name} gfx = _parse_adapter_family_gfx(str(family)) if gfx: @@ -2322,7 +3669,7 @@ def _windows_amd_adapter_records_by_luid() -> dict[int, Dict[str, Any]]: by_luid[int(luid)] = record except Exception as e: logger.debug("DirectX adapter registry read declined: %s", e) - return {} + return None return by_luid @@ -4749,6 +6096,9 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: "parent_visible_gpu_ids": [], "devices": [], "index_kind": "vulkan", + # Physically present cards this PyTorch cannot open, reported ALONGSIDE the empty + # `devices` above and never merged into it. + **_torch_gpu_mismatch_report(), } diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py index adefc121e37..b3d2e6e15eb 100644 --- a/studio/backend/utils/hardware/nvidia.py +++ b/studio/backend/utils/hardware/nvidia.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import os +import platform +import shutil import subprocess from typing import Any, Optional @@ -236,24 +239,58 @@ def get_visible_gpu_utilization( } -def get_backend_visible_gpu_info( - parent_visible_ids: Optional[list[int]], backend_cuda_visible_devices: Optional[str] -) -> dict[str, Any]: - # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to - # visible devices. - if parent_visible_ids is None: - return { - "available": False, - "backend_cuda_visible_devices": backend_cuda_visible_devices, - "parent_visible_gpu_ids": [], - "devices": [], - "index_kind": "unresolved", - } - visible_ordinals = _visible_ordinal_map(parent_visible_ids) +def _nvidia_smi_executable() -> str: + """The nvidia-smi to run, resolving the standard Windows locations off PATH. + + A driver install can leave nvidia-smi.exe in the NVSMI directory or the driver + store without putting either on PATH, and a bare "nvidia-smi" then raises + FileNotFoundError. That would leave the physical inventory empty on exactly the + host this inventory exists for: real GPUs, a PyTorch that cannot see them. Same + two locations setup.ps1 already falls back to. Returns the bare name when nothing + better is found, so the caller's existing OSError handling still applies. + """ + found = shutil.which("nvidia-smi") + if found: + return found + if platform.system() != "Windows": + return "nvidia-smi" + for base, tail in ( + (os.environ.get("ProgramFiles"), r"NVIDIA Corporation\NVSMI\nvidia-smi.exe"), + (os.environ.get("SystemRoot"), r"System32\nvidia-smi.exe"), + ): + if not base: + continue + candidate = os.path.join(base, tail) + if os.path.isfile(candidate): + return candidate + return "nvidia-smi" + + +# "nvidia-smi is not on this machine" is a conclusive answer, not a failed probe: it is the +# normal state of every CPU-only, AMD and Intel host, and the installers read the same +# absence the same way. Distinct from None, which means a probe that WAS found could not +# answer -- a hung driver, a permission fault, a non-zero exit. +NVIDIA_SMI_ABSENT = object() + + +def _query_gpu_inventory(caller: str) -> Any: + """``[{index, name, memory_total_gb}]`` for every GPU nvidia-smi enumerates. + + ``None`` when the query could not be answered at all -- no nvidia-smi on PATH, a + driver that hung past the timeout, a non-zero exit. Callers report that as + "unknown", which is not the same as the empty list a working driver with no + cards returns. Never raises. + + Split out of get_backend_visible_gpu_info so the same rows can be read WITHOUT a + ``DeviceType.CUDA`` precondition: get_physical_gpu_inventory below is reached on + exactly the host where torch reports no CUDA device, and that host still has its + GPUs. Rows a caller cannot make sense of are dropped, not raised on -- a name + holding commas is rejoined, and a malformed index or memory column skips the row. + """ try: result = subprocess.run( [ - "nvidia-smi", + _nvidia_smi_executable(), "--query-gpu=index,name,memory.total", "--format=csv,noheader,nounits", ], @@ -265,25 +302,20 @@ def get_backend_visible_gpu_info( env = child_env_without_native_path_secret(), **_windows_hidden_subprocess_kwargs(), ) + except FileNotFoundError as e: + # No nvidia-smi at all, the NORMAL state of every CPU-only, AMD and Intel host. This is + # called on a 60 second refresh reached from the health and system polls, so warning here + # would log a line every minute on machines that are working correctly. + logger.debug("nvidia-smi is not installed (%s): %s", caller, e) + return NVIDIA_SMI_ABSENT except (OSError, subprocess.TimeoutExpired) as e: - logger.warning("nvidia-smi query failed in get_backend_visible_gpu_info: %s", e) - return { - "available": False, - "backend_cuda_visible_devices": backend_cuda_visible_devices, - "parent_visible_gpu_ids": parent_visible_ids or [], - "devices": [], - "index_kind": "physical", - } + # Past this point an nvidia-smi WAS found, so a failure is a real fault on this host. + logger.warning("nvidia-smi query failed in %s: %s", caller, e) + return None if result.returncode != 0: - return { - "available": False, - "backend_cuda_visible_devices": backend_cuda_visible_devices, - "parent_visible_gpu_ids": parent_visible_ids or [], - "devices": [], - "index_kind": "physical", - } + return None - devices = [] + rows: list[dict[str, Any]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] if len(parts) < 3: @@ -292,14 +324,125 @@ def get_backend_visible_gpu_info( idx = int(parts[0]) except (ValueError, TypeError): continue - if visible_ordinals is not None and idx not in visible_ordinals: - continue # Rejoin in case the GPU name contains commas name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1]) try: mem_total_mb = int(parts[-1]) except (ValueError, TypeError): continue + rows.append( + { + "index": idx, + "name": name, + "memory_total_gb": round(mem_total_mb / 1024, 2), + } + ) + return rows + + +def _linux_nvidia_procfs_gpu_count() -> int: + """How many GPUs the NVIDIA kernel driver enumerates under /proc, or 0. + + One subdirectory per GPU, published whatever nvidia-smi's state is, which is why the + installer falls back to it too. Never raises; 0 on any platform without it. + """ + if platform.system() != "Linux": + return 0 + try: + entries = os.listdir("/proc/driver/nvidia/gpus") + except OSError: + return 0 + return len(entries) + + +def get_physical_gpu_inventory() -> dict[str, Any]: + """Every NVIDIA GPU the driver enumerates, with no visibility mask and no torch. + + Display-only inventory: ``index`` is nvidia-smi's own row number, which is a + physical id and NOT something a caller may pin, because the whole point of this + probe is that PyTorch cannot open these devices. A failed probe comes back as a + structured unavailable result, so this never raises out of an endpoint. + """ + rows = _query_gpu_inventory("get_physical_gpu_inventory") + # Either way the CLI could not answer. The kernel driver publishes its cards regardless, + # and on a cold start there is no settled verdict for the resulting unknown to protect. + if (rows is NVIDIA_SMI_ABSENT or rows is None) and _linux_nvidia_procfs_gpu_count(): + # The kernel driver is loaded and enumerating cards; only the CLI is missing. + # _has_usable_nvidia_gpu() reads the same directory, so without this the installer can + # repair a CUDA wheel on a host the backend insists has no card. No name and no capacity: + # procfs gives neither, and an invented one would be worse than an honest blank. + return { + "available": True, + "source": "proc-driver-nvidia", + "devices": [ + { + "vendor": "nvidia", + "index": ordinal, + "name": None, + "memory_total_gb": None, + "source": "proc-driver-nvidia", + } + for ordinal in range(_linux_nvidia_procfs_gpu_count()) + ], + "error": None, + "absent": False, + } + if rows is NVIDIA_SMI_ABSENT: + # An answer, and the caller must not read it as "some probe failed": an AMD-only host has + # no nvidia-smi by design. + return { + "available": False, + "source": "nvidia-smi", + "devices": [], + "error": "nvidia-smi is not installed", + "absent": True, + } + if rows is None: + return { + "available": False, + "source": "nvidia-smi", + "devices": [], + "error": "nvidia-smi did not answer", + "absent": False, + } + return { + "available": bool(rows), + "source": "nvidia-smi", + "devices": [{**row, "vendor": "nvidia", "source": "nvidia-smi"} for row in rows], + "error": None, + "absent": False, + } + + +def get_backend_visible_gpu_info( + parent_visible_ids: Optional[list[int]], backend_cuda_visible_devices: Optional[str] +) -> dict[str, Any]: + # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to + # visible devices. + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + visible_ordinals = _visible_ordinal_map(parent_visible_ids) + rows = _query_gpu_inventory("get_backend_visible_gpu_info") + if rows is None or rows is NVIDIA_SMI_ABSENT: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + devices = [] + for row in rows: + idx = row["index"] + if visible_ordinals is not None and idx not in visible_ordinals: + continue devices.append( { "index": idx, @@ -307,8 +450,8 @@ def get_backend_visible_gpu_info( "visible_ordinal": ( visible_ordinals[idx] if visible_ordinals is not None else len(devices) ), - "name": name, - "memory_total_gb": round(mem_total_mb / 1024, 2), + "name": row["name"], + "memory_total_gb": row["memory_total_gb"], } ) diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index d51a639a925..abc02112923 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -31,6 +31,7 @@ import { useTheme, } from "@/features/settings"; import { SttDownloadPrompt } from "@/features/settings/components/stt-download-prompt"; +import { TauriRepairContext } from "@/hooks/tauri-repair-context"; import { TauriUpdateContext } from "@/hooks/tauri-update-context"; import { type BackendStatus, useTauriBackend } from "@/hooks/use-tauri-backend"; import { useTauriUpdate } from "@/hooks/use-tauri-update"; @@ -43,6 +44,7 @@ import { type CSSProperties, type ReactNode, useEffect, + useMemo, useRef, useState, } from "react"; @@ -555,8 +557,26 @@ function TauriWrapper({ children }: { children: ReactNode }) { retryInstall, approveElevation, copyDiagnostics, + startRepair, } = useTauriBackend(); + // Settings' manual repair reruns the INSTALLER, not `studio update`: an update reuses the + // environment it finds, so a venv whose PyTorch was replaced by a CPU-only wheel comes back + // from a successful update still CPU-only. + // + // Through a ref, not a dependency: startRepair is a plain function declaration rebuilt on + // every render, so listing it would give the context a new identity on each status tick and + // pinning it with [] would freeze the first render's closure. + const startRepairRef = useRef(startRepair); + startRepairRef.current = startRepair; + const repairController = useMemo( + () => ({ + repairInstall: () => startRepairRef.current({ forceInstaller: true }), + isExternalServer, + }), + [isExternalServer], + ); + const appliedWindowModeRef = useRef(null); const hasEnteredAppModeRef = useRef(false); const windowLayoutGenerationRef = useRef(0); @@ -781,10 +801,10 @@ function TauriWrapper({ children }: { children: ReactNode }) { // alike, and a declined quit puts the user back where they were rather than remounting // the tree under them. const content = ( - <> + {shell} {closing && } - + ); const chromeVars = ( diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 7b2022658e4..28face8770d 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -588,6 +588,18 @@ const SIDEBAR_SELECTOR = '[data-slot="sidebar"]'; const VERDICT_UNKNOWN_POLL_MS = 3000; const SELF_HEAL_POLL_MS = 15000; const VERDICT_POLL_STALL_MS = 30000; +// The backend refreshes its physical GPU inventory on a 60s TTL and can reclassify a host +// without a restart: attach an eGPU to a CPU-torch machine and no_gpu becomes +// torch_cpu_build. Nothing else re-reads the verdict. Matched to that TTL, because polling +// faster than the backend can change its answer is pure request traffic. +const INVENTORY_POLL_MS = 60000; +// The verdicts the inventory can still move. Everything else describes something a probe +// cannot change (an Intel Mac stays an Intel Mac). +const INVENTORY_SENSITIVE_REASONS = new Set([ + "no_gpu", + "torch_cpu_build", + "torch_cuda_unavailable", +]); /** One workflow in the list under the Images row. */ function WorkflowChoice({ @@ -810,9 +822,16 @@ export function AppSidebar() { : "Training needs MLX. Run `unsloth studio update` to enable Train." : chatOnlyReason === "intel_mac" ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." - : chatOnlyReason === "no_gpu" - ? "Training needs an NVIDIA or AMD GPU." - : undefined; + : chatOnlyReason === "torch_cpu_build" || + chatOnlyReason === "torch_cuda_unavailable" + ? // The host HAS GPUs; this PyTorch cannot open them. "Get a GPU" is both wrong + // and unactionable here, so name the installed build and point at the repair. + chatOnlyDetail + ? `Training needs a working PyTorch GPU build. This machine's GPUs were detected but PyTorch ${chatOnlyDetail} cannot use them; repair the installation.` + : "Training needs a working PyTorch GPU build. This machine's GPUs were detected but PyTorch cannot use them; repair the installation." + : chatOnlyReason === "no_gpu" + ? "Training needs an NVIDIA or AMD GPU." + : undefined; // Everything without a hint reaches VideoPage, which answers from the backend's video verdict. const videoDisabledHint = videoNavHint(chatOnlyMeasured, chatOnlyReason); const videoDisabled = videoDisabledHint !== undefined; @@ -832,7 +851,9 @@ export function AppSidebar() { // recovery poll in the app, and the sidebar is mounted on every route that gates on the // verdict (studio-page reads the same store, so it recovers with it; video-page reads the // backend's video verdict instead and needs nothing from here). - if (selfHealSettled && !capabilitiesUnknown) return; + const inventorySensitive = + chatOnly && INVENTORY_SENSITIVE_REASONS.has(chatOnlyReason ?? ""); + if (selfHealSettled && !capabilitiesUnknown && !inventorySensitive) return; let pollingSince = 0; // Which read currently owns the guard. A read that outlived the stall window is replaced, // and the replacement takes the guard with it; without an owner the abandoned read's @@ -851,7 +872,11 @@ export function AppSidebar() { .finally(() => { if (owned === pollOwner) pollingSince = 0; }); - }, capabilitiesUnknown ? VERDICT_UNKNOWN_POLL_MS : SELF_HEAL_POLL_MS); + }, capabilitiesUnknown + ? VERDICT_UNKNOWN_POLL_MS + : selfHealSettled + ? INVENTORY_POLL_MS + : SELF_HEAL_POLL_MS); return () => window.clearInterval(id); }, [capabilitiesUnknown, chatOnly, chatOnlyReason, detectionDeferred]); diff --git a/studio/frontend/src/config/hardware-verdict.ts b/studio/frontend/src/config/hardware-verdict.ts index 6ab06253574..57638dc6465 100644 --- a/studio/frontend/src/config/hardware-verdict.ts +++ b/studio/frontend/src/config/hardware-verdict.ts @@ -50,6 +50,13 @@ export function videoNavHint( // that reads as a fix that would not work. Mirrors the backend's message for this host. if (chatOnlyReason === "intel_mac") return "Video generation requires Apple Silicon. This Intel Mac has no Metal device to run it."; + // The GPU is there; PyTorch is what cannot reach it (a Windows update that resolved torch + // from PyPI leaves a +cpu wheel beside two working cards), so "get a GPU" is a lie. + if ( + chatOnlyReason === "torch_cpu_build" || + chatOnlyReason === "torch_cuda_unavailable" + ) + return "Video generation needs a working PyTorch GPU build. This machine's GPUs were detected but PyTorch cannot use them; repair the installation."; if (chatOnlyReason === "no_gpu") return "Video generation needs an NVIDIA or AMD GPU."; return undefined; } diff --git a/studio/frontend/src/features/settings/components/desktop-repair-control.tsx b/studio/frontend/src/features/settings/components/desktop-repair-control.tsx new file mode 100644 index 00000000000..289236d9a84 --- /dev/null +++ b/studio/frontend/src/features/settings/components/desktop-repair-control.tsx @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { useTauriRepairController } from "@/hooks/tauri-repair-context"; +import { useT } from "@/i18n"; +import { type ReactElement, useState } from "react"; +import { SettingsRow } from "./settings-row"; + +/** + * Reruns the bundled installer over the managed environment. + * + * Exists because an update cannot repair everything an update can break. `studio update` + * reuses the environment it finds; only the installer re-selects the PyTorch index and + * force-reinstalls the trio. A managed venv that ended up with a CPU-only PyTorch wheel + * therefore survives any number of successful updates, and until now the only way out was + * to paste the install one-liner into a terminal. + * + * Desktop-only: outside Tauri there is no controller in context and the row renders nothing. + * Hidden for an externally started backend too: start_managed_repair refuses to mutate an + * environment the desktop does not manage, but the refusal arrives after startRepair has + * already cleared the external-server state and swapped the shell to the repairing screen, + * so a connected user would end up on the repair-error screen instead of on their server. + * Confirmed before it runs, because it stops the backend and rewrites the environment. + */ +export function DesktopRepairControl(): ReactElement | null { + const t = useT(); + const repair = useTauriRepairController(); + const [confirmOpen, setConfirmOpen] = useState(false); + if (!repair || repair.isExternalServer) return null; + + return ( + <> + + + + + + + + + {t("settings.general.repairInstall.confirmTitle")} + + + {t("settings.general.repairInstall.confirmDescription")} + + + + + + + + + + ); +} diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 9faf44af18e..a26c760edc5 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -59,6 +59,7 @@ import { import { loadCloseToTray, updateCloseToTray } from "../api/close-to-tray"; import { loadLaunchAtLogin, updateLaunchAtLogin } from "../api/launch-at-login"; import { ChangePasswordDialog } from "../components/change-password-dialog"; +import { DesktopRepairControl } from "../components/desktop-repair-control"; import { DesktopUpdateControl, DesktopUpdateNote, @@ -742,6 +743,10 @@ export function GeneralTab() { {t("settings.general.resetPreferences.action")} + {/* Same section as the reset row: both rewrite state the user cannot easily put + back, and the desktop-only repair renders nothing on the web build, which + would leave a section header with no rows under it if it had its own. */} + diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx index f3dfb7f54ba..ce40eec7e8a 100644 --- a/studio/frontend/src/features/settings/tabs/resources-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -19,6 +19,7 @@ import { aggregateGpuMemoryTotalGb, useSystemInfo, type GpuDevice, + type SystemGpuInfo, } from "@/hooks/use-system"; import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; @@ -193,6 +194,33 @@ function deviceOrdinal(device: GpuDevice): number | undefined { return device.visible_ordinal ?? device.index; } +/** A GPU the OS enumerates that this PyTorch cannot open, from /api/system's + * `gpu.physical_devices`. `index` is the probe's own row number, vendor-local and + * NOT a pin. Declared here instead of widened onto SystemGpuInfo on purpose: these + * are display-only, and that shared type is what model fit budgets against and what + * the training device picker pins from, where an unusable card must never appear. */ +interface PhysicalGpuDevice { + vendor?: string; + index?: number; + name?: string | null; + memory_total_gb?: number | null; + source?: string; +} + +/** Why the devices above are unusable, from /api/system's `gpu.mismatch`. Absent on + * a healthy host and on one that genuinely has no GPU, so its presence is the whole + * signal. `reason` is "torch_cpu_build" or "torch_cuda_unavailable". */ +interface GpuTorchMismatch { + reason?: string; + torch_version?: string | null; + physical_count?: number; +} + +type GpuPhysicalInventory = SystemGpuInfo & { + physical_devices?: PhysicalGpuDevice[]; + mismatch?: GpuTorchMismatch | null; +}; + export function ResourcesTab() { const t = useT(); const liveUpdates = useSettingsPanelPrefsStore((s) => s.resourcesLiveUpdates); @@ -401,6 +429,25 @@ export function ResourcesTab() { : formatGiB(metrics.vramTotal); // The placeholder's cpu backend and empty package list are not facts about the host either. const hostReading = (value: string) => (hostUnread ? unknownLabel : value); + // systemInfo.gpu rather than displayedGpu: this describes the TRAINING view of the host, + // and a Vulkan llama.cpp makes displayedGpu fall back to the inference inventory, which is + // precisely the host that must be told. Gated on the read having settled. + const gpuInventory = hostUnread + ? null + : ((systemInfo.gpu ?? null) as GpuPhysicalInventory | null); + const gpuMismatch = gpuInventory?.mismatch ?? null; + const physicalDevices = gpuMismatch + ? (gpuInventory?.physical_devices ?? []) + : []; + // A CPU-only wheel is fixed by reinstalling torch, a dead runtime by the driver. + const gpuMismatchMessage = gpuMismatch + ? t( + gpuMismatch.reason === "torch_cpu_build" + ? "settings.resources.gpu.mismatchCpuBuild" + : "settings.resources.gpu.mismatchUnavailable", + { version: gpuMismatch.torch_version ?? unknownLabel }, + ) + : null; return (
@@ -489,7 +536,9 @@ export function ResourcesTab() { ? metrics.vramUsageKnown ? `${formatGiB(metrics.vramUsed)} / ${vramCapacityLabel}` : `${unknownLabel} / ${vramCapacityLabel}` - : t("settings.resources.liveMonitor.noGpu") + : gpuMismatch + ? t("settings.resources.liveMonitor.gpuUnusable") + : t("settings.resources.liveMonitor.noGpu") } detail={ gpuUnknown @@ -500,7 +549,9 @@ export function ResourcesTab() { value: formatGiB(metrics.vramFree), }) : unknownLabel - : backendLabel + : gpuMismatch + ? t("settings.resources.liveMonitor.gpuUnusableDetail") + : backendLabel } percent={metrics.vramUsageKnown ? metrics.vramPercent : null} /> @@ -508,6 +559,33 @@ export function ResourcesTab() { + {/* Physically present, torch-unusable cards. Their own block, above the + device rows and visually separate from them, because nothing here is + selectable: the rows below are what a model can be loaded onto. */} + {gpuMismatch ? ( +
+

+ {gpuMismatchMessage} +

+ {physicalDevices.map((device, index) => ( +
+ + {device.name ?? t("settings.resources.gpu.unknownDevice")} + + + {`${ + isFiniteNumber(device.memory_total_gb) + ? formatGiB(device.memory_total_gb) + : unknownLabel + } · ${t("settings.resources.gpu.unusableDevice")}`} + +
+ ))} +
+ ) : null} {separateInferenceGpu && (
@@ -617,6 +695,12 @@ export function ResourcesTab() {
); }) + ) : gpuMismatch ? ( + // Not "no visible GPU": the cards are listed directly above. Its own branch, so the + // CPU-only host's line below stays exactly as it was. +
+ {t("settings.resources.gpu.noUsableGpu")} +
) : (
{gpuUnknown ? gpuUnknownLabel : t("settings.resources.gpu.noGpu")} diff --git a/studio/frontend/src/hooks/tauri-repair-context.ts b/studio/frontend/src/hooks/tauri-repair-context.ts new file mode 100644 index 00000000000..baf86048b44 --- /dev/null +++ b/studio/frontend/src/hooks/tauri-repair-context.ts @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createContext, useContext } from "react"; + +/** What Settings needs from the backend controller to offer a manual repair. */ +export type TauriRepairController = { + /** Reruns the bundled installer over the managed environment, then restarts the backend. */ + repairInstall: () => Promise; + /** + * True when the desktop is attached to a backend someone started in a terminal. + * + * start_managed_repair refuses to touch that environment, but startRepair switches the + * whole shell to the repairing screen before the refusal comes back, so an offered + * action would drop a connected user on the repair-error screen. The consumer hides + * itself instead. The settings update row already suppresses its action this way. + */ + isExternalServer: boolean; +}; + +// Null outside Tauri and on the startup screen, where the app shell has not mounted. The +// consumer renders nothing rather than offering an action that cannot run. +export const TauriRepairContext = createContext( + null, +); + +export function useTauriRepairController(): TauriRepairController | null { + return useContext(TauriRepairContext); +} diff --git a/studio/frontend/src/hooks/use-tauri-backend.ts b/studio/frontend/src/hooks/use-tauri-backend.ts index ab7800c10ef..b2d52214ebc 100644 --- a/studio/frontend/src/hooks/use-tauri-backend.ts +++ b/studio/frontend/src/hooks/use-tauri-backend.ts @@ -164,6 +164,9 @@ export function useTauriBackend() { const externalPollAbortedRef = useRef(false); const authFailureRef = useRef(getTauriAuthFailure()); const elevationResumeRef = useRef<"install" | "repair" | null>(null); + // Whether the repair in flight was asked to skip straight to the installer. Read back by + // approveElevation, which restarts the repair after the system packages land. + const forcedRepairRef = useRef(false); const [tauriEventsReady, setTauriEventsReady] = useState(!isTauri); // Read through rather than mirrored into state: the app-closing listener is registered // inside the long event effect below, which cannot reach a setState from this render. @@ -369,7 +372,14 @@ export function useTauriBackend() { startingRef.current = false; } - async function startRepair() { + // `forceInstaller` runs the bundled installer without trying `studio update` first. The + // automatic callers leave it off, because an out-of-date venv is the common case. Settings' + // manual repair turns it on: an update reuses the environment it finds, so a venv whose + // PyTorch was replaced by a CPU-only wheel comes back from one still CPU-only. + async function startRepair(options?: { forceInstaller?: boolean }) { + const forceInstaller = options?.forceInstaller ?? false; + // Survives the elevation round trip: approveElevation resumes by calling this again. + forcedRepairRef.current = forceInstaller; elevationResumeRef.current = null; setCurrentStepIndex(-1); setProgressDetail(null); @@ -384,7 +394,7 @@ export function useTauriBackend() { const { invoke } = await import("@tauri-apps/api/core"); try { - await invoke("start_managed_repair"); + await invoke("start_managed_repair", { forceInstaller }); setBackendStatus("starting"); elevationResumeRef.current = null; @@ -467,6 +477,19 @@ export function useTauriBackend() { } const retry = useCallback(() => { + // Retry on a FORCED repair has to re-run that repair, not the preflight. The installer + // is transactional, so a failed attempt over an existing install restores the desktop-ready + // environment it found: checkInstallAndStart() then sees a ready install and restarts the + // same CPU-only backend the user pressed Repair about, and the button does nothing. + // Elevation resumes already preserve this; the error path did not. + if (statusRef.current === "repair-error" && forcedRepairRef.current) { + clearAuthFailure(); + clearServerStopIntent(); + setError(null); + void startRepair({ forceInstaller: true }); + return; + } + forcedRepairRef.current = false; clearAuthFailure(); clearServerStopIntent(); setError(null); @@ -515,7 +538,7 @@ export function useTauriBackend() { setProgressDetail(null); elevationResumeRef.current = null; if (resume === "repair") { - await startRepair(); + await startRepair({ forceInstaller: forcedRepairRef.current }); } else { await startInstall(); } @@ -724,5 +747,8 @@ export function useTauriBackend() { currentStepIndex, progressDetail, startupMessage, elevationPackages, startServer, stopServer, startInstall, retry, retryInstall, approveElevation, copyDiagnostics, + // The same function startup uses, so a manual repair renders the same repairing screen + // and restarts the backend afterwards rather than leaving it stopped. + startRepair, }; } diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts index 65361275383..5f6fbbe0ffb 100644 --- a/studio/frontend/src/i18n/locales/ar.ts +++ b/studio/frontend/src/i18n/locales/ar.ts @@ -831,6 +831,16 @@ export const ar = { openError: "تعذّر فتح المجلد", copyError: "تعذّر نسخ المسار", }, + repairInstall: { + label: "إصلاح التثبيت", + description: + "يعيد تشغيل المثبّت على البيئة المُدارة. استخدمه إذا لم يتم اكتشاف وحدة معالجة الرسومات أو إذا تعذّر بدء التطبيق.", + action: "إصلاح التثبيت", + confirmTitle: "هل تريد إصلاح هذا التثبيت؟", + confirmDescription: + "يوقف الخادم ويعيد تشغيل المثبّت، الذي يعيد تثبيت PyTorch لوحدة معالجة الرسومات في هذا الجهاز. تُحفظ المحادثات والإعدادات. قد يستغرق ذلك عدة دقائق.", + confirmAction: "إصلاح الآن", + }, resetPreferences: { sectionTitle: "منطقة الخطر", label: "إعادة تعيين جميع التفضيلات المحلية", @@ -1086,6 +1096,8 @@ export const ar = { currentLoad: "الحمل الحالي", free: "{value} متاح", noGpu: "لا يوجد GPU مرئي", + gpuUnusable: "GPU غير قابل للاستخدام", + gpuUnusableDetail: "تم اكتشافه، لكن PyTorch لا يستطيع استخدامه", }, gpu: { title: "أجهزة GPU", @@ -1094,6 +1106,12 @@ export const ar = { detecting: "جارٍ البحث عن وحدات GPU...", unreadable: "تعذّرت قراءة عتاد هذا الخادم.", noGpu: "لم يُكتشف أي GPU مرئي. تُعرض موارد CPU فقط أعلاه.", + noUsableGpu: "لا يوجد على هذا الجهاز أي GPU يمكن لـ PyTorch استخدامه.", + mismatchCpuBuild: + "PyTorch إصدار للـ CPU فقط ({version})، لذا لا يمكن استخدام وحدات GPU أدناه. أصلح التثبيت لاستعادة دعم GPU.", + mismatchUnavailable: + "لا يستطيع PyTorch ({version}) تهيئة وحدات GPU أدناه، لذا لا يمكن استخدامها. تحقق من تعريف كرت الشاشة أو أصلح التثبيت.", + unusableDevice: "غير قابل للاستخدام", unknownDevice: "GPU غير معروف", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts index b3a17b8e380..cfb79302325 100644 --- a/studio/frontend/src/i18n/locales/de.ts +++ b/studio/frontend/src/i18n/locales/de.ts @@ -855,6 +855,16 @@ export const de = { openError: "Der Ordner konnte nicht geöffnet werden", copyError: "Der Pfad konnte nicht kopiert werden", }, + repairInstall: { + label: "Installation reparieren", + description: + "Führt das Installationsprogramm erneut über die verwaltete Umgebung aus. Nützlich, wenn die GPU nicht erkannt wird oder die App nicht startet.", + action: "Installation reparieren", + confirmTitle: "Diese Installation reparieren?", + confirmDescription: + "Stoppt den Server und führt das Installationsprogramm erneut aus, wodurch PyTorch für die GPU dieses Rechners neu installiert wird. Chats und Einstellungen bleiben erhalten. Dies kann einige Minuten dauern.", + confirmAction: "Jetzt reparieren", + }, resetPreferences: { sectionTitle: "Gefahrenzone", label: "Alle lokalen Einstellungen zurücksetzen", @@ -1112,6 +1122,8 @@ export const de = { currentLoad: "Aktuelle Auslastung", free: "{value} frei", noGpu: "Keine sichtbare GPU", + gpuUnusable: "GPU nicht nutzbar", + gpuUnusableDetail: "Erkannt, aber PyTorch kann sie nicht nutzen", }, gpu: { title: "GPU-Geräte", @@ -1121,6 +1133,12 @@ export const de = { unreadable: "Die Hardware dieses Servers konnte nicht gelesen werden.", noGpu: "Keine sichtbare GPU erkannt. Oben werden nur die CPU-Ressourcen angezeigt.", + noUsableGpu: "Auf diesem Rechner ist keine GPU für PyTorch nutzbar.", + mismatchCpuBuild: + "PyTorch ist ein reiner CPU-Build ({version}), daher können die GPUs unten nicht genutzt werden. Reparieren Sie die Installation, um die GPU-Unterstützung wiederherzustellen.", + mismatchUnavailable: + "PyTorch ({version}) kann die GPUs unten nicht initialisieren, daher sind sie nicht nutzbar. Prüfen Sie den Grafiktreiber oder reparieren Sie die Installation.", + unusableDevice: "nicht nutzbar", unknownDevice: "Unbekannte GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 1715543fe07..741b0b166b8 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -832,6 +832,16 @@ export const en = { openError: "Couldn't open the folder", copyError: "Couldn't copy the path", }, + repairInstall: { + label: "Repair installation", + description: + "Reruns the installer over the managed environment. Use this if the GPU is not detected or the app will not start.", + action: "Repair install", + confirmTitle: "Repair this installation?", + confirmDescription: + "Stops the server and reruns the installer, which reinstalls PyTorch for this machine's GPU. Chats and settings are kept. This can take several minutes.", + confirmAction: "Repair now", + }, resetPreferences: { sectionTitle: "Danger zone", label: "Reset all local preferences", @@ -1081,6 +1091,8 @@ export const en = { currentLoad: "Current load", free: "{value} free", noGpu: "No visible GPU", + gpuUnusable: "GPU unusable", + gpuUnusableDetail: "Detected, but PyTorch cannot use it", }, gpu: { title: "GPU devices", @@ -1089,6 +1101,12 @@ export const en = { detecting: "Checking for GPUs...", unreadable: "Could not read this server's hardware.", noGpu: "No visible GPU detected. CPU-only resources are shown above.", + noUsableGpu: "No GPU on this machine is usable by PyTorch.", + mismatchCpuBuild: + "PyTorch is a CPU-only build ({version}), so the GPUs below cannot be used. Repair the installation to restore GPU support.", + mismatchUnavailable: + "PyTorch ({version}) cannot initialise the GPUs below, so they cannot be used. Check the GPU driver, or repair the installation.", + unusableDevice: "unusable", unknownDevice: "Unknown GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts index f7320158a07..7325b72fc11 100644 --- a/studio/frontend/src/i18n/locales/es.ts +++ b/studio/frontend/src/i18n/locales/es.ts @@ -850,6 +850,16 @@ export const es = { openError: "No se pudo abrir la carpeta", copyError: "No se pudo copiar la ruta", }, + repairInstall: { + label: "Reparar la instalación", + description: + "Vuelve a ejecutar el instalador sobre el entorno gestionado. Útil si no se detecta la GPU o si la aplicación no arranca.", + action: "Reparar instalación", + confirmTitle: "¿Reparar esta instalación?", + confirmDescription: + "Detiene el servidor y vuelve a ejecutar el instalador, que reinstala PyTorch para la GPU de este equipo. Se conservan las conversaciones y los ajustes. Puede tardar varios minutos.", + confirmAction: "Reparar ahora", + }, resetPreferences: { sectionTitle: "Zona de peligro", label: "Restablecer todas las preferencias locales", @@ -1107,6 +1117,8 @@ export const es = { currentLoad: "Carga actual", free: "Libre: {value}", noGpu: "No hay GPU visible", + gpuUnusable: "GPU no utilizable", + gpuUnusableDetail: "Detectada, pero PyTorch no puede usarla", }, gpu: { title: "Dispositivos GPU", @@ -1116,6 +1128,12 @@ export const es = { unreadable: "No se pudo leer el hardware de este servidor.", noGpu: "No se detectó ninguna GPU visible. Arriba se muestran los recursos solo de CPU.", + noUsableGpu: "Ninguna GPU de esta máquina se puede usar con PyTorch.", + mismatchCpuBuild: + "PyTorch es una compilación solo para CPU ({version}), por lo que las GPU de abajo no se pueden usar. Repare la instalación para restaurar la compatibilidad con GPU.", + mismatchUnavailable: + "PyTorch ({version}) no puede inicializar las GPU de abajo, por lo que no se pueden usar. Revise el controlador de la GPU o repare la instalación.", + unusableDevice: "no utilizable", unknownDevice: "GPU desconocida", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts index a10d7ac6d6a..475fde15a8f 100644 --- a/studio/frontend/src/i18n/locales/fr.ts +++ b/studio/frontend/src/i18n/locales/fr.ts @@ -853,6 +853,16 @@ export const fr = { openError: "Impossible d'ouvrir le dossier", copyError: "Impossible de copier le chemin", }, + repairInstall: { + label: "Réparer l'installation", + description: + "Relance le programme d'installation sur l'environnement géré. Utile si le GPU n'est pas détecté ou si l'application ne démarre pas.", + action: "Réparer l'installation", + confirmTitle: "Réparer cette installation ?", + confirmDescription: + "Arrête le serveur et relance le programme d'installation, qui réinstalle PyTorch pour le GPU de cette machine. Les discussions et les paramètres sont conservés. Cela peut prendre plusieurs minutes.", + confirmAction: "Réparer maintenant", + }, resetPreferences: { sectionTitle: "Zone de danger", label: "Réinitialiser toutes les préférences locales", @@ -1111,6 +1121,8 @@ export const fr = { currentLoad: "Charge actuelle", free: "Disponible : {value}", noGpu: "Aucun GPU visible", + gpuUnusable: "GPU inutilisable", + gpuUnusableDetail: "Détecté, mais PyTorch ne peut pas l'utiliser", }, gpu: { title: "Périphériques GPU", @@ -1120,6 +1132,12 @@ export const fr = { unreadable: "Impossible de lire le matériel de ce serveur.", noGpu: "Aucun GPU visible n'a été détecté. Seules les ressources du CPU sont affichées ci-dessus.", + noUsableGpu: "Aucun GPU de cette machine n'est utilisable par PyTorch.", + mismatchCpuBuild: + "PyTorch est une version CPU uniquement ({version}), les GPU ci-dessous ne peuvent donc pas être utilisés. Réparez l'installation pour rétablir la prise en charge du GPU.", + mismatchUnavailable: + "PyTorch ({version}) ne parvient pas à initialiser les GPU ci-dessous, ils ne peuvent donc pas être utilisés. Vérifiez le pilote graphique ou réparez l'installation.", + unusableDevice: "inutilisable", unknownDevice: "GPU inconnu", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts index fc72c8106f5..b5f71690cc0 100644 --- a/studio/frontend/src/i18n/locales/hi.ts +++ b/studio/frontend/src/i18n/locales/hi.ts @@ -834,6 +834,16 @@ export const hi = { openError: "फ़ोल्डर नहीं खोला जा सका", copyError: "पथ कॉपी नहीं किया जा सका", }, + repairInstall: { + label: "इंस्टॉलेशन ठीक करें", + description: + "प्रबंधित परिवेश पर इंस्टॉलर को फिर से चलाता है। यदि GPU का पता नहीं चल रहा है या ऐप शुरू नहीं हो रहा है तो इसका उपयोग करें।", + action: "इंस्टॉल ठीक करें", + confirmTitle: "क्या यह इंस्टॉलेशन ठीक करना है?", + confirmDescription: + "सर्वर रोकता है और इंस्टॉलर फिर से चलाता है, जो इस मशीन के GPU के लिए PyTorch को दोबारा इंस्टॉल करता है। चैट और सेटिंग्स सुरक्षित रहती हैं। इसमें कई मिनट लग सकते हैं।", + confirmAction: "अभी ठीक करें", + }, resetPreferences: { sectionTitle: "खतरनाक क्षेत्र", label: "सभी स्थानीय प्राथमिकताएं रीसेट करें", @@ -1089,6 +1099,8 @@ export const hi = { currentLoad: "वर्तमान लोड", free: "{value} खाली", noGpu: "कोई दृश्यमान GPU नहीं", + gpuUnusable: "GPU अनुपयोगी", + gpuUnusableDetail: "पहचाना गया, लेकिन PyTorch इसका उपयोग नहीं कर सकता", }, gpu: { title: "GPU डिवाइस", @@ -1097,6 +1109,12 @@ export const hi = { detecting: "GPU खोजे जा रहे हैं...", unreadable: "इस सर्वर का हार्डवेयर नहीं पढ़ा जा सका।", noGpu: "कोई दृश्यमान GPU नहीं मिला। केवल-CPU संसाधन ऊपर दिखाए गए हैं।", + noUsableGpu: "इस मशीन का कोई भी GPU PyTorch के लिए उपयोग योग्य नहीं है।", + mismatchCpuBuild: + "PyTorch केवल-CPU बिल्ड ({version}) है, इसलिए नीचे दिए गए GPU उपयोग नहीं किए जा सकते। GPU समर्थन बहाल करने के लिए इंस्टॉलेशन की मरम्मत करें।", + mismatchUnavailable: + "PyTorch ({version}) नीचे दिए गए GPU को आरंभ नहीं कर पा रहा है, इसलिए उनका उपयोग नहीं किया जा सकता। GPU ड्राइवर जाँचें या इंस्टॉलेशन की मरम्मत करें।", + unusableDevice: "अनुपयोगी", unknownDevice: "अज्ञात GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/it.ts b/studio/frontend/src/i18n/locales/it.ts index c16df90eb32..17ef7f6e533 100644 --- a/studio/frontend/src/i18n/locales/it.ts +++ b/studio/frontend/src/i18n/locales/it.ts @@ -825,6 +825,16 @@ export const it = { openError: "Impossibile aprire la cartella", copyError: "Impossibile copiare il percorso", }, + repairInstall: { + label: "Ripara l'installazione", + description: + "Riesegue il programma di installazione sull'ambiente gestito. Utile se la GPU non viene rilevata o se l'app non si avvia.", + action: "Ripara installazione", + confirmTitle: "Riparare questa installazione?", + confirmDescription: + "Arresta il server e riesegue il programma di installazione, che reinstalla PyTorch per la GPU di questa macchina. Le chat e le impostazioni vengono mantenute. L'operazione può richiedere alcuni minuti.", + confirmAction: "Ripara ora", + }, resetPreferences: { sectionTitle: "Zona pericolosa", label: "Reimposta tutte le preferenze locali", @@ -1078,6 +1088,8 @@ export const it = { currentLoad: "Carico attuale", free: "Disponibili: {value}", noGpu: "Nessuna GPU visibile", + gpuUnusable: "GPU non utilizzabile", + gpuUnusableDetail: "Rilevata, ma PyTorch non può usarla", }, gpu: { title: "Dispositivi GPU", @@ -1087,6 +1099,12 @@ export const it = { unreadable: "Impossibile leggere l'hardware di questo server.", noGpu: "Nessuna GPU visibile rilevata. Sopra sono mostrate le risorse della sola CPU.", + noUsableGpu: "Nessuna GPU di questa macchina è utilizzabile da PyTorch.", + mismatchCpuBuild: + "PyTorch è una build solo CPU ({version}), quindi le GPU sottostanti non possono essere usate. Ripara l'installazione per ripristinare il supporto GPU.", + mismatchUnavailable: + "PyTorch ({version}) non riesce a inizializzare le GPU sottostanti, quindi non possono essere usate. Controlla il driver della GPU o ripara l'installazione.", + unusableDevice: "non utilizzabile", unknownDevice: "GPU sconosciuta", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index 623d793cc6b..26c866c8ee2 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -818,6 +818,16 @@ export const ja = { openError: "フォルダを開けませんでした", copyError: "パスをコピーできませんでした", }, + repairInstall: { + label: "インストールを修復", + description: + "管理環境に対してインストーラーを再実行します。GPU が検出されない場合やアプリが起動しない場合に使用します。", + action: "インストールを修復", + confirmTitle: "このインストールを修復しますか？", + confirmDescription: + "サーバーを停止してインストーラーを再実行し、このマシンの GPU に合わせて PyTorch を再インストールします。チャットと設定は保持されます。数分かかる場合があります。", + confirmAction: "今すぐ修復", + }, resetPreferences: { sectionTitle: "危険ゾーン", label: "すべてのローカル設定をリセット", @@ -1067,6 +1077,8 @@ export const ja = { currentLoad: "現在の負荷", free: "{value} 空き", noGpu: "認識できる GPU がありません", + gpuUnusable: "GPU を利用できません", + gpuUnusableDetail: "検出されましたが、PyTorch から利用できません", }, gpu: { title: "GPU デバイス", @@ -1076,6 +1088,12 @@ export const ja = { unreadable: "このサーバーのハードウェアを読み取れませんでした。", noGpu: "利用可能な GPU が検出されませんでした。CPU のみの環境向けのリソース情報は上に表示されています。", + noUsableGpu: "このマシンの GPU はいずれも PyTorch から利用できません。", + mismatchCpuBuild: + "PyTorch は CPU 専用ビルド ({version}) のため、下の GPU は利用できません。インストールを修復すると GPU サポートが復元されます。", + mismatchUnavailable: + "PyTorch ({version}) は下の GPU を初期化できないため、利用できません。GPU ドライバーを確認するか、インストールを修復してください。", + unusableDevice: "利用不可", unknownDevice: "不明な GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts index b2e81f482e2..f5a17b5bceb 100644 --- a/studio/frontend/src/i18n/locales/ko.ts +++ b/studio/frontend/src/i18n/locales/ko.ts @@ -827,6 +827,16 @@ export const ko = { openError: "폴더를 열지 못했습니다", copyError: "경로를 복사하지 못했습니다", }, + repairInstall: { + label: "설치 복구", + description: + "관리 환경에 대해 설치 프로그램을 다시 실행합니다. GPU가 감지되지 않거나 앱이 시작되지 않을 때 사용하세요.", + action: "설치 복구", + confirmTitle: "이 설치를 복구할까요?", + confirmDescription: + "서버를 중지하고 설치 프로그램을 다시 실행하여 이 컴퓨터의 GPU에 맞는 PyTorch를 다시 설치합니다. 채팅과 설정은 유지됩니다. 몇 분 정도 걸릴 수 있습니다.", + confirmAction: "지금 복구", + }, resetPreferences: { sectionTitle: "위험 구역", label: "모든 로컬 환경설정 재설정", @@ -1082,6 +1092,8 @@ export const ko = { currentLoad: "현재 부하", free: "{value} 여유", noGpu: "인식되는 GPU 없음", + gpuUnusable: "GPU 사용 불가", + gpuUnusableDetail: "감지되었지만 PyTorch에서 사용할 수 없습니다", }, gpu: { title: "GPU 장치", @@ -1090,6 +1102,12 @@ export const ko = { detecting: "GPU를 확인하는 중...", unreadable: "이 서버의 하드웨어를 읽을 수 없습니다.", noGpu: "인식되는 GPU가 없습니다. 위에는 CPU 관련 리소스만 표시됩니다.", + noUsableGpu: "이 컴퓨터의 GPU 중 PyTorch에서 사용할 수 있는 것이 없습니다.", + mismatchCpuBuild: + "PyTorch가 CPU 전용 빌드({version})이므로 아래 GPU를 사용할 수 없습니다. 설치를 복구하면 GPU 지원이 복원됩니다.", + mismatchUnavailable: + "PyTorch({version})가 아래 GPU를 초기화하지 못해 사용할 수 없습니다. GPU 드라이버를 확인하거나 설치를 복구하세요.", + unusableDevice: "사용 불가", unknownDevice: "알 수 없는 GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index c5d07a4a457..1885b182f8b 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -842,6 +842,16 @@ export const ptBR = { openError: "Não foi possível abrir a pasta", copyError: "Não foi possível copiar o caminho", }, + repairInstall: { + label: "Reparar a instalação", + description: + "Executa o instalador novamente sobre o ambiente gerenciado. Útil se a GPU não for detectada ou se o aplicativo não iniciar.", + action: "Reparar instalação", + confirmTitle: "Reparar esta instalação?", + confirmDescription: + "Para o servidor e executa o instalador novamente, que reinstala o PyTorch para a GPU desta máquina. As conversas e as configurações são mantidas. Isso pode levar alguns minutos.", + confirmAction: "Reparar agora", + }, resetPreferences: { sectionTitle: "Zona de perigo", label: "Redefinir todas as preferências locais", @@ -1096,6 +1106,8 @@ export const ptBR = { currentLoad: "Carga atual", free: "{value} livres", noGpu: "Nenhuma GPU visível", + gpuUnusable: "GPU inutilizável", + gpuUnusableDetail: "Detectada, mas o PyTorch não consegue usá-la", }, gpu: { title: "Dispositivos de GPU", @@ -1104,6 +1116,12 @@ export const ptBR = { detecting: "Procurando GPUs...", unreadable: "Não foi possível ler o hardware deste servidor.", noGpu: "Nenhuma GPU visível detectada. Os recursos somente CPU aparecem acima.", + noUsableGpu: "Nenhuma GPU desta máquina pode ser usada pelo PyTorch.", + mismatchCpuBuild: + "O PyTorch é uma compilação somente CPU ({version}), então as GPUs abaixo não podem ser usadas. Repare a instalação para restaurar o suporte a GPU.", + mismatchUnavailable: + "O PyTorch ({version}) não consegue inicializar as GPUs abaixo, então elas não podem ser usadas. Verifique o driver da GPU ou repare a instalação.", + unusableDevice: "inutilizável", unknownDevice: "GPU desconhecida", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts index e3e3be3f37b..36b73939726 100644 --- a/studio/frontend/src/i18n/locales/ru.ts +++ b/studio/frontend/src/i18n/locales/ru.ts @@ -835,6 +835,16 @@ export const ru = { openError: "Не удалось открыть папку", copyError: "Не удалось скопировать путь", }, + repairInstall: { + label: "Восстановить установку", + description: + "Повторно запускает установщик для управляемой среды. Пригодится, если графический процессор не определяется или приложение не запускается.", + action: "Восстановить установку", + confirmTitle: "Восстановить эту установку?", + confirmDescription: + "Останавливает сервер и повторно запускает установщик, который переустанавливает PyTorch для графического процессора этого компьютера. Чаты и настройки сохраняются. Это может занять несколько минут.", + confirmAction: "Восстановить", + }, resetPreferences: { sectionTitle: "Опасная зона", label: "Сбросить все локальные настройки", @@ -1091,6 +1101,8 @@ export const ru = { currentLoad: "Текущая нагрузка", free: "{value} свободно", noGpu: "GPU не обнаружен", + gpuUnusable: "GPU недоступен", + gpuUnusableDetail: "Обнаружен, но PyTorch не может его использовать", }, gpu: { title: "Устройства GPU", @@ -1100,6 +1112,12 @@ export const ru = { unreadable: "Не удалось прочитать оборудование этого сервера.", noGpu: "Доступные GPU не обнаружены. Выше показаны ресурсы только для CPU.", + noUsableGpu: "Ни один GPU на этой машине недоступен для PyTorch.", + mismatchCpuBuild: + "PyTorch собран только для CPU ({version}), поэтому GPU ниже использовать нельзя. Восстановите установку, чтобы вернуть поддержку GPU.", + mismatchUnavailable: + "PyTorch ({version}) не может инициализировать GPU ниже, поэтому использовать их нельзя. Проверьте драйвер видеокарты или восстановите установку.", + unusableDevice: "недоступен", unknownDevice: "Неизвестный GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index dcacd0fb1bb..4789c3b9bdf 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -811,6 +811,16 @@ export const zhCN = { openError: "无法打开文件夹", copyError: "无法复制路径", }, + repairInstall: { + label: "修复安装", + description: + "对托管环境重新运行安装程序。如果未检测到 GPU 或应用无法启动，可使用此功能。", + action: "修复安装", + confirmTitle: "要修复此安装吗？", + confirmDescription: + "停止服务器并重新运行安装程序，为本机 GPU 重新安装 PyTorch。聊天记录和设置会保留。这可能需要几分钟。", + confirmAction: "立即修复", + }, resetPreferences: { sectionTitle: "危险区域", label: "重置所有本地偏好设置", @@ -1061,6 +1071,8 @@ export const zhCN = { currentLoad: "当前负载", free: "{value} 可用", noGpu: "无可见 GPU", + gpuUnusable: "GPU 不可用", + gpuUnusableDetail: "已检测到，但 PyTorch 无法使用", }, gpu: { title: "GPU 设备", @@ -1069,6 +1081,10 @@ export const zhCN = { detecting: "正在检查 GPU...", unreadable: "无法读取此服务器的硬件信息。", noGpu: "未检测到可见 GPU。上方仅显示 CPU 资源。", + noUsableGpu: "本机没有 PyTorch 可用的 GPU。", + mismatchCpuBuild: "PyTorch 是仅 CPU 版本（{version}），因此无法使用下方的 GPU。修复安装即可恢复 GPU 支持。", + mismatchUnavailable: "PyTorch（{version}）无法初始化下方的 GPU，因此无法使用。请检查显卡驱动，或修复安装。", + unusableDevice: "不可用", unknownDevice: "未知 GPU", deviceWithIndex: "GPU {index}", vramUtilization: "VRAM", diff --git a/studio/frontend/tests/forced-repair-retry.test.ts b/studio/frontend/tests/forced-repair-retry.test.ts new file mode 100644 index 00000000000..f71e2d7cc8e --- /dev/null +++ b/studio/frontend/tests/forced-repair-retry.test.ts @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Retry, on a repair the user forced from Settings, has to re-run that repair. +// +// startRepair({ forceInstaller: true }) skips `studio update` and runs the bundled +// installer, which is the only thing that replaces a CPU-only wheel. The installer is +// transactional: a failed attempt over an existing install restores the desktop-ready +// environment it found. So the generic retry path -- clear state, run the preflight -- +// finds a ready install and restarts the very backend the user pressed Repair about. +// The elevation resume already carried the flag; the error path did not. +// +// The hook cannot be rendered here, so the callback is lifted by regex and evaluated, +// the way gpu-torch-mismatch.test.ts and system-status-verdict.test.ts do beside it. + +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import test from "node:test"; + +const src = await readFile( + new URL("../src/hooks/use-tauri-backend.ts", import.meta.url), + "utf8", +); + +function lift(pattern: RegExp, what: string): string { + const found = pattern.exec(src); + assert.ok(found, `could not find ${what} in use-tauri-backend.ts`); + return found[0]; +} + +const body = lift( + /const retry = useCallback\(\(\) => \{[\s\S]*?\n \}, \[\]\);/, + "the retry callback", +); + +type Run = { repairs: boolean[]; preflights: number; forcedAfter: boolean }; + +function runRetry(status: string, forced: boolean): Run { + const repairs: boolean[] = []; + const forcedRepairRef = { current: forced }; + let preflights = 0; + const noop = () => {}; + const scope = { + statusRef: { current: status }, + forcedRepairRef, + startRepair: (options?: { forceInstaller?: boolean }) => { + repairs.push(options?.forceInstaller === true); + return Promise.resolve(); + }, + checkInstallAndStart: () => { + preflights += 1; + }, + elevationResumeRef: { current: null as string | null }, + startingRef: { current: true }, + portRef: { current: 1 as number | null }, + startTimedOutRef: { current: true }, + seenStepsRef: { current: new Set() }, + useCallback: (fn: unknown) => fn, + clearAuthFailure: noop, + clearServerStopIntent: noop, + setError: noop, + setLogs: noop, + setCurrentStepIndex: noop, + setProgressDetail: noop, + setElevationPackages: noop, + setIsExternalServer: noop, + stopExternalServerPoll: noop, + }; + const keys = Object.keys(scope); + new Function( + ...keys, + `${body.replace(/^const retry = /, "return ")}`.replace(/;\s*$/, ";"), + )(...keys.map((key) => (scope as Record)[key]))(); + return { repairs, preflights, forcedAfter: forcedRepairRef.current }; +} + +test("retry after a forced repair re-runs the forced repair", () => { + const run = runRetry("repair-error", true); + assert.deepEqual(run.repairs, [true], "the retry must force the installer again"); + assert.equal(run.preflights, 0, "the preflight would restart the same broken backend"); +}); + +test("retry after an automatic repair still runs the preflight", () => { + // The automatic callers leave forceInstaller off, and an out-of-date venv really is + // the common case there: nothing about that path changed. + const run = runRetry("repair-error", false); + assert.deepEqual(run.repairs, []); + assert.equal(run.preflights, 1); +}); + +test("retry from any other failure is untouched", () => { + for (const status of ["error", "install-error", "not-installed", "stopped"]) { + const run = runRetry(status, true); + assert.deepEqual(run.repairs, [], `${status} must not start a repair`); + assert.equal(run.preflights, 1, `${status} must still run the preflight`); + assert.equal( + run.forcedAfter, + false, + `${status} leaves the generic path, so the forced flag must not survive it`, + ); + } +}); diff --git a/studio/frontend/tests/gpu-torch-mismatch.test.ts b/studio/frontend/tests/gpu-torch-mismatch.test.ts new file mode 100644 index 00000000000..bdb395097c5 --- /dev/null +++ b/studio/frontend/tests/gpu-torch-mismatch.test.ts @@ -0,0 +1,418 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// What the UI says about a host whose GPUs PyTorch cannot use. +// +// A Windows in-app update that resolved torch from PyPI left a 2.11.0+cpu wheel beside two +// working RTX A4000s (#8473). The backend now reports those cards in gpu.physical_devices +// with a gpu.mismatch reason, and three places have to stop saying the opposite: the System +// tab's VRAM tile and GPU section, the sidebar's Train hint, and videoNavHint. +// +// The derivations live in .tsx files that pull in the whole app, so they are lifted by +// regex and evaluated, as system-status-verdict.test.ts does beside them. + +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import test from "node:test"; + +import { registerBundlerResolver } from "./helpers/kit.ts"; + +registerBundlerResolver(); + +const { videoNavHint } = await import("../src/config/hardware-verdict.ts"); +const { en } = await import("../src/i18n/locales/en.ts"); + +const tabSrc = await readFile( + new URL("../src/features/settings/tabs/resources-tab.tsx", import.meta.url), + "utf8", +); +const sidebarSrc = await readFile( + new URL("../src/components/app-sidebar.tsx", import.meta.url), + "utf8", +); + +function lift( + src: string, + pattern: RegExp, + what: string, + where: string, +): string { + const found = pattern.exec(src); + assert.ok(found, `could not find ${what} in ${where}`); + return found[0]; +} + +const t = (key: string) => key; + +const CPU_BUILD = "settings.resources.gpu.mismatchCpuBuild"; +const UNAVAILABLE = "settings.resources.gpu.mismatchUnavailable"; +const NO_USABLE_GPU = "settings.resources.gpu.noUsableGpu"; +const NO_GPU = "settings.resources.gpu.noGpu"; +const UNKNOWN = "settings.resources.environment.unknown"; + +// gpuInventory carries a TS cast, so it is asserted on as source below and supplied here +// as an input instead. +const derivation = [ + lift( + tabSrc, + /const gpuMismatch = [\s\S]*?;/, + "gpuMismatch", + "resources-tab.tsx", + ), + lift( + tabSrc, + /const physicalDevices = [\s\S]*?;/, + "physicalDevices", + "resources-tab.tsx", + ), + lift( + tabSrc, + /const gpuMismatchMessage = [\s\S]*?;/, + "gpuMismatchMessage", + "resources-tab.tsx", + ), +].join("\n"); + +interface Inventory { + mismatch?: { reason?: string; torch_version?: string | null } | null; + physical_devices?: { name?: string }[]; +} + +function mismatchFor(gpuInventory: Inventory | null) { + const run = new Function( + "gpuInventory", + "t", + "unknownLabel", + `${derivation} + return { gpuMismatch, physicalDevices, gpuMismatchMessage };`, + ); + return run(gpuInventory, t, UNKNOWN) as { + gpuMismatch: { reason?: string } | null; + physicalDevices: { name?: string }[]; + gpuMismatchMessage: string | null; + }; +} + +test("a CPU-only wheel and a dead accelerator wheel get different sentences", () => { + const cpuBuild = mismatchFor({ + mismatch: { reason: "torch_cpu_build", torch_version: "2.11.0+cpu" }, + physical_devices: [ + { name: "NVIDIA RTX A4000" }, + { name: "NVIDIA RTX A4000" }, + ], + }); + assert.equal(cpuBuild.gpuMismatchMessage, CPU_BUILD); + assert.equal(cpuBuild.physicalDevices.length, 2); + + // Reinstalling torch is the wrong advice for a healthy wheel whose runtime will not + // start, so the two reasons must not collapse into one string. + const dead = mismatchFor({ + mismatch: { + reason: "torch_cuda_unavailable", + torch_version: "2.6.0+cu124", + }, + physical_devices: [{ name: "NVIDIA RTX A4000" }], + }); + assert.equal(dead.gpuMismatchMessage, UNAVAILABLE); +}); + +test("a healthy host, and one that really has no GPU, get no banner at all", () => { + for (const inventory of [ + null, + {}, + { mismatch: null }, + ] as (Inventory | null)[]) { + const out = mismatchFor(inventory); + assert.equal(out.gpuMismatch, null); + assert.equal(out.gpuMismatchMessage, null); + assert.deepEqual(out.physicalDevices, []); + } + const strayRows = mismatchFor({ + physical_devices: [{ name: "NVIDIA RTX A4000" }], + }); + assert.deepEqual(strayRows.physicalDevices, []); +}); + +test("the verdict is taken from a settled read only, and from the training view", () => { + // The placeholder useSystemInfo starts from is shaped like a CPU-only host, so a banner + // derived from it would accuse a host nobody has measured yet. + const inventory = lift( + tabSrc, + /const gpuInventory = [\s\S]*?;\n/, + "gpuInventory", + "resources-tab.tsx", + ); + assert.match( + inventory, + /hostUnread\s*\n?\s*\?\s*null/, + "gated on the read having settled", + ); + // systemInfo.gpu, NOT displayedGpu: a Vulkan llama.cpp makes displayedGpu fall back to the + // inference inventory, and that host is exactly the second report in #8473. + assert.match(inventory, /systemInfo\.gpu/); + assert.doesNotMatch(inventory, /displayedGpu/); +}); + +test("the GPU section stops telling this host there is no GPU", () => { + assert.match( + tabSrc, + new RegExp(`\\) : gpuMismatch \\? \\([\\s\\S]*?t\\("${NO_USABLE_GPU}"\\)`), + "a host with unusable cards gets its own line", + ); + assert.match( + tabSrc, + /gpuUnknown \? gpuUnknownLabel : t\("settings\.resources\.gpu\.noGpu"\)/, + "and a host that really has no GPU still gets the CPU-only one", + ); +}); + +test("the VRAM tile stops reading as a CPU-only host", () => { + const tiles = tabSrc.match(//g) ?? []; + const vram = tiles.find((tile) => tile.includes("liveMonitor.vram")); + assert.ok(vram, "the VRAM tile"); + const mismatchAt = vram.indexOf("liveMonitor.gpuUnusable"); + const noGpuAt = vram.indexOf("liveMonitor.noGpu"); + assert.ok(mismatchAt > -1, "the tile has a mismatch state"); + assert.ok(noGpuAt > -1, "and still has the CPU-only state"); + assert.ok(mismatchAt < noGpuAt, "the mismatch state is reached first"); +}); + +test("the physically detected cards are shown, and never offered as devices", () => { + // The banner renders physicalDevices and the selectable rows metrics.devices. If the + // banner ever read metrics.devices the two would merge, which the field split prevents. + const banner = lift( + tabSrc, + /\{gpuMismatch \? \(\n[\s\S]*?\n\s*\) : null\}/, + "the mismatch banner", + "resources-tab.tsx", + ); + assert.match(banner, /physicalDevices\.map/); + assert.doesNotMatch(banner, /metrics\.devices/); + assert.match(banner, /settings\.resources\.gpu\.unusableDevice/); +}); + + +test("videoNavHint stops telling a two-GPU host to get a GPU", () => { + for (const reason of ["torch_cpu_build", "torch_cuda_unavailable"]) { + const hint = videoNavHint(true, reason); + assert.ok(hint, `${reason} explains the disabled Video row`); + assert.doesNotMatch( + hint, + /needs an NVIDIA or AMD GPU/, + `${reason} is not a missing-GPU host`, + ); + assert.match(hint, /PyTorch/, `${reason} names what is actually wrong`); + assert.equal(videoNavHint(false, reason), undefined); + } + // And the genuine no-GPU host keeps the sentence that is true for it. + assert.equal( + videoNavHint(true, "no_gpu"), + "Video generation needs an NVIDIA or AMD GPU.", + ); +}); + +test("the sidebar's Train hint stops doing the same", () => { + const hint = lift( + sidebarSrc, + /const trainDisabledHint: string \| undefined = [\s\S]*?\n\s*: undefined;/, + "trainDisabledHint", + "app-sidebar.tsx", + ); + const forReason = (chatOnlyReason: string, chatOnlyDetail: string | null) => + new Function( + "chatOnlyMeasured", + "chatOnlyReason", + "chatOnlyDetail", + `${hint.replace(": string | undefined", "")}\nreturn trainDisabledHint;`, + )(true, chatOnlyReason, chatOnlyDetail) as string | undefined; + + for (const reason of ["torch_cpu_build", "torch_cuda_unavailable"]) { + const withDetail = forReason(reason, "2.11.0+cpu"); + assert.ok(withDetail); + assert.doesNotMatch(withDetail, /needs an NVIDIA or AMD GPU/); + // The installed build is what makes this actionable to someone whose update already ran. + assert.match(withDetail, /2\.11\.0\+cpu/); + const withoutDetail = forReason(reason, null); + assert.ok(withoutDetail); + assert.doesNotMatch(withoutDetail, /needs an NVIDIA or AMD GPU/); + } + assert.equal( + forReason("no_gpu", null), + "Training needs an NVIDIA or AMD GPU.", + ); + assert.equal(forReason("detection_failed", null), undefined); +}); + + +test("every string the banner reaches for exists", () => { + const gpu = en.settings.resources.gpu as Record; + const liveMonitor = en.settings.resources.liveMonitor as Record< + string, + string + >; + for (const key of [ + "noUsableGpu", + "mismatchCpuBuild", + "mismatchUnavailable", + "unusableDevice", + ]) { + assert.equal(typeof gpu[key], "string", `settings.resources.gpu.${key}`); + } + for (const key of ["gpuUnusable", "gpuUnusableDetail"]) { + assert.equal( + typeof liveMonitor[key], + "string", + `settings.resources.liveMonitor.${key}`, + ); + } + // The version is what a user can check against their own install, so both sentences + // have to carry it. + assert.match(gpu.mismatchCpuBuild, /\{version\}/); + assert.match(gpu.mismatchUnavailable, /\{version\}/); + // And the CPU-only host's line is still the one it always was. + assert.equal(t(NO_GPU), NO_GPU); + assert.match(gpu.noGpu, /No visible GPU detected/); +}); + +// The repair row must not be offered for a backend the desktop does not manage. +// +// start_managed_repair rejects that mutation, but only after startRepair has cleared +// isExternalServer, stopped the external-server poll and swapped the shell to the repairing +// screen, so a connected user lands on the repair-error screen instead of on their server. +test("the repair row hides itself for an externally started backend", async () => { + const source = await readFile( + new URL( + "../src/features/settings/components/desktop-repair-control.tsx", + import.meta.url, + ), + "utf8", + ); + assert.match( + source, + /if\s*\(!repair\s*\|\|\s*repair\.isExternalServer\)\s*return null;/, + "the control must bail out on an external server as well as outside Tauri", + ); + + const context = await readFile( + new URL("../src/hooks/tauri-repair-context.ts", import.meta.url), + "utf8", + ); + assert.match( + context, + /isExternalServer:\s*boolean;/, + "the controller has to carry the flag for the control to read it", + ); + + const provider = await readFile( + new URL("../src/app/provider.tsx", import.meta.url), + "utf8", + ); + const memo = provider.slice(provider.indexOf("const repairController")); + assert.match( + memo.slice(0, 300), + /isExternalServer,/, + "the provider has to publish the flag", + ); + assert.match( + memo.slice(0, 300), + /\[isExternalServer\]/, + "and list it as a dependency, or the context freezes on the first render's value", + ); +}); + +// The verdict can change without a restart, so the sidebar has to keep asking. +// +// The backend refreshes its physical inventory on a 60s TTL: attach an eGPU to a CPU-torch +// machine and no_gpu becomes torch_cpu_build. The polling effect stopped at the first +// settled verdict, so the new hint was unreachable for the rest of the session. +test("the sidebar keeps polling while the inventory can still change the verdict", async () => { + const source = await readFile( + new URL("../src/components/app-sidebar.tsx", import.meta.url), + "utf8", + ); + + assert.match( + source, + /INVENTORY_SENSITIVE_REASONS = new Set\(\[[^\]]*"no_gpu"[^\]]*"torch_cpu_build"[^\]]*"torch_cuda_unavailable"/s, + "the three verdicts the inventory can move must all keep the poll alive", + ); + + const set = source.slice( + source.indexOf("INVENTORY_SENSITIVE_REASONS = new Set(["), + ); + const listed = set.slice(0, set.indexOf("]")); + for (const settled of ["mlx_unavailable", "intel_mac", "detection_failed"]) { + assert.ok( + !listed.includes(settled), + `${settled} cannot change on a probe and must not keep polling`, + ); + } + + assert.match( + source, + /if \(selfHealSettled && !capabilitiesUnknown && !inventorySensitive\) return;/, + "the early return has to consider the inventory-sensitive case", + ); + assert.match(source, /const INVENTORY_POLL_MS = 60000;/); + assert.match( + source, + /selfHealSettled\s*\?\s*INVENTORY_POLL_MS\s*:\s*SELF_HEAL_POLL_MS/, + "a settled host polls at the inventory cadence, not the self-heal one", + ); +}); + +// The poll decision itself, evaluated rather than pattern-matched. +// +// The test above pins the shape of the early return; this one runs it. A regression that +// keeps the guard's text but inverts its sense would give every working install a forced +// /api/system read a minute for the life of the session, which is the opposite of what +// this change is for. +test("only a host the inventory can still reclassify keeps polling", () => { + const guard = lift( + sidebarSrc, + /const inventorySensitive =[\s\S]*?if \(selfHealSettled && !capabilitiesUnknown && !inventorySensitive\) return;/, + "the polling guard", + "app-sidebar.tsx", + ); + const reasons = lift( + sidebarSrc, + /const INVENTORY_SENSITIVE_REASONS = new Set\(\[[\s\S]*?\]\);/, + "INVENTORY_SENSITIVE_REASONS", + "app-sidebar.tsx", + ); + const polls = ( + chatOnly: boolean, + chatOnlyReason: string | null, + selfHealSettled = true, + capabilitiesUnknown = false, + ) => + new Function( + "chatOnly", + "chatOnlyReason", + "selfHealSettled", + "capabilitiesUnknown", + `${reasons} + ${guard} + return true;`, + )(chatOnly, chatOnlyReason, selfHealSettled, capabilitiesUnknown) === true; + + // The hosts this change exists for. Their verdict moves on the next inventory refresh. + assert.ok(polls(true, "torch_cpu_build")); + assert.ok(polls(true, "torch_cuda_unavailable")); + assert.ok(polls(true, "no_gpu"), "an eGPU can arrive on a CPU-only box"); + + // And the hosts that were working before this PR and must keep working the same way. + assert.ok( + !polls(false, null), + "a healthy GPU host must not gain a forced read a minute", + ); + assert.ok(!polls(true, "intel_mac"), "an Intel Mac stays an Intel Mac"); + assert.ok( + !polls(true, "detection_failed"), + "a probe cannot undo a detection that already failed", + ); + + // The two pre-existing polls are untouched. + assert.ok(polls(true, "mlx_unavailable", false), "the MLX self-heal poll"); + assert.ok(polls(false, null, true, true), "the unknown-verdict poll"); +}); diff --git a/studio/install_manifest.py b/studio/install_manifest.py index 269b67386cf..22272111026 100644 --- a/studio/install_manifest.py +++ b/studio/install_manifest.py @@ -274,6 +274,8 @@ def write_manifest( steps_total: int = 0, package_name: str = "unsloth", no_torch: Optional[bool] = None, + expected_torch_tag: Optional[str] = None, + expected_torch_tag_pinned: Optional[bool] = None, ) -> Optional[Path]: """Record a completed install. Never raises: no manifest reads as incomplete, which is the safe answer.""" @@ -301,6 +303,17 @@ def write_manifest( # exports nothing and would otherwise reinstall torch into a GGUF-only venv. if no_torch is not None: payload["no_torch"] = bool(no_torch) + # The FLAVOR, never the index URL it came from: a pinned index can carry a token in + # its userinfo, query or fragment, and this file lives in the venv and is read back + # by verify-install, desktop-capabilities and the setup fast path. + if expected_torch_tag: + payload["expected_torch_tag"] = str(expected_torch_tag).strip().lower() + # Whether that flavor was NAMED by whoever ran the install, or merely what the selection + # landed on: setup.ps1 picks /cpu automatically on a GPU-less host and publishes it exactly + # as it publishes a pinned one, and reading the automatic case as deliberate leaves a later + # eGPU with no repair offered. Absent means unknown, as with every other additive key. + if expected_torch_tag_pinned is not None: + payload["expected_torch_tag_pinned"] = bool(expected_torch_tag_pinned) path = manifest_path(root) try: tmp = path.with_suffix(".json.tmp") @@ -374,6 +387,45 @@ def recorded_no_torch(root: Optional[Path] = None) -> Optional[bool]: return None +def recorded_torch_flavor(root: Optional[Path] = None) -> Optional[str]: + """The torch flavor this venv was installed with, or None when unknown. + + None means nothing recorded it: no manifest, or one written before the key + existed. Callers must treat None as "unknown" and fall back to their own + detection, never as "cpu" -- claiming a flavor nobody selected would let a + repair reinstall over a deliberate build. + + There is no marker companion here (unlike no_torch): the manifest is dropped + before every dependency pass, so this answers only for the PREVIOUS install, + which is exactly the question a repair asks. A run whose own setup script + exported the flavor never reaches this. + """ + manifest = read_manifest(root) + if manifest is None: + return None + value = manifest.get("expected_torch_tag") + if not isinstance(value, str): + return None + value = value.strip().lower() + return value or None + + +def recorded_torch_flavor_was_pinned(root: Optional[Path] = None) -> bool: + """Whether the recorded flavor was NAMED rather than automatically selected. + + False when nothing recorded it, including a manifest written before the key + existed. That is the safe direction here and the opposite of the usual "unknown + falls back to the old behaviour": treating an unproven CPU record as deliberate is + what leaves a host that has since gained a GPU with no repair offered at all, which + is the failure this whole field exists to distinguish. A repair is something the + user can decline; a silently CPU-only GPU box is not. + """ + manifest = read_manifest(root) + if manifest is None: + return False + return bool(manifest.get("expected_torch_tag_pinned")) + + def _parse_requirement_line(line: str) -> Optional[Tuple[str, str, str]]: """(distribution name, marker, specifier) for a requirement, or None. diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 596d87502b1..727fed77cfd 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -266,6 +266,14 @@ def _torch_index_leaf(url: str) -> str: # mismatched. _CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC +# Byte-identical to the non-XPU arm of install.ps1's $_fixSpecs, NOT _CUDA_TORCH_PKG_SPEC: +# `studio update` must repair to the same wheels install.ps1 does. +_TORCH_FLAVOR_REPAIR_PKG_SPEC: tuple[str, str, str] = ( + "torch>=2.4,<2.11.0", + "torchvision>=0.19,<0.26.0", + "torchaudio>=2.4,<2.11.0", +) + # torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch # mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails # to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a @@ -329,6 +337,9 @@ def _select_torchao_spec(torch_version: str | None) -> str: # pip_install_try(), the only things here that change what is installed. None means # "not probed yet". _TORCH_RUNTIME_PROBE: "tuple[bool, bool, str | None, str, str] | None" = None +# Not in the tuple above: thirteen call sites unpack that, and only the GPU-build +# verdict needs this one. +_TORCH_RUNTIME_XPU: str = "" # Prefix on the probe's own stdout line. Import chatter can arrive before the answer, and # an atexit handler or a CUDA teardown notice can arrive after it, so "the last non-empty @@ -387,7 +398,9 @@ def _probe_torch_runtime() -> "tuple[bool, bool, str | None, str, str]": "_v = getattr(torch, 'version', None); " "h = getattr(_v, 'hip', '') or ''; " "c = getattr(_v, 'cuda', '') or ''; " - f"print('{_TORCH_PROBE_MARKER}' + '|'.join((v, h, c)))" + # XPU too: such a wheel carries its runtime here and nowhere else. + "x = getattr(_v, 'xpu', '') or ''; " + f"print('{_TORCH_PROBE_MARKER}' + '|'.join((v, h, c, x)))" ), ], stdout = subprocess.PIPE, @@ -413,9 +426,12 @@ def _probe_torch_runtime() -> "tuple[bool, bool, str | None, str, str]": ] version: "str | None" = None hip = cuda = "" + global _TORCH_RUNTIME_XPU if _marked: _fields = _marked[-1][len(_TORCH_PROBE_MARKER) :].split("|") - version, hip, cuda = (_fields + ["", ""])[:3] + version, hip, cuda = (_fields + ["", "", ""])[:3] + # Beside the tuple, not in it: thirteen call sites unpack five values. + _TORCH_RUNTIME_XPU = (_fields + ["", "", "", ""])[3] _TORCH_RUNTIME_PROBE = (True, probe.returncode == 0, version, hip, cuda) return _TORCH_RUNTIME_PROBE @@ -2622,9 +2638,13 @@ def _ensure_xpu_triton() -> None: Lives here, not in install.sh, because install.sh runs setup.sh which runs this file: one copy covers the fresh install AND `unsloth studio update`, which never touches install.sh. - Windows is excluded -- studio/setup.ps1 owns the same swap there. + + On Windows setup.ps1 performs the same swap after this script exits, so its handover + variable means "someone else will"; a direct run has no such postlude. """ - if NO_TORCH or IS_MACOS or IS_WINDOWS: + if NO_TORCH or IS_MACOS: + return + if IS_WINDOWS and os.environ.get("UNSLOTH_EXPECTED_TORCH_TAG", "").strip(): return pin = _explicit_xpu_torch_index_url() if pin is None: @@ -2856,6 +2876,532 @@ def _ensure_cpu_torch() -> None: ) +def _torch_flavor_tag(version: str) -> str: + """Classify a torch.__version__ into the installers' flavor vocabulary. + + +cuNNN -> "cuNNN", +rocm -> "rocm", +xpu -> "xpu", +cpu or untagged -> "cpu". + MUST match install.ps1's ConvertTo-TorchFlavorTag and setup.ps1's stale-venv probe, + because the tag this returns is compared against one those produced. "" only for an + empty version, which means the classification failed rather than "cpu". + + Untagged reads as "cpu" deliberately: PyPI forbids the local +cuNNN label, so an + untagged wheel is the PyPI build, which on Windows is CPU-only. A repair off that + verdict is self-resolving -- the replacement carries a +cuNNN tag and matches on the + next pass -- and _torch_build_is_gpu, not this function, decides whether to FAIL. + """ + value = str(version).strip().lower() + if not value: + return "" + match = re.search(r"\+(cu\d+)", value) + if match: + return match.group(1) + if "+rocm" in value: + return "rocm" + if "+xpu" in value: + return "xpu" + return "cpu" + + +def _gpu_family_from_runtime_markers(hip: str, cuda: str) -> str: + """Which GPU family an untagged wheel's runtime markers name. + + torch.version.xpu alongside .hip and .cuda, for the reason _torch_build_is_gpu already + reads all three: an untagged source, conda or private-index XPU build carries its + runtime only there. Omitting it let an explicit /cpu pin over such a wheel compare + equal, return success without replacing anything, and then record a PINNED cpu flavor + for an environment that still holds an XPU build. + """ + if hip: + return "rocm" + if cuda: + return "cuda" + return "xpu" + + +def _torch_build_is_gpu() -> bool: + """Whether the installed torch can use a GPU at all, on the evidence available. + + Weaker and more forgiving than _torch_flavor_tag, and used only for the FAIL verdict + in _ensure_expected_torch_flavor: a wrong family is worth a reinstall, but only a + build with no GPU support whatsoever is worth failing the update over. + + torch.version.cuda / .hip count alongside the local label, so an untagged wheel that + does carry a CUDA runtime is not called CPU-only. An answer that never arrived (a + wedged driver hanging `import torch` -- the host these repairs exist for) falls back + to the on-disk label and, failing that, reads as a GPU build: ambiguity must not fail + an update by itself. + """ + _ran, _importable, _version, _hip, _cuda = _probe_torch_runtime() + if _ran and _importable and _version: + return ( + _is_gpu_torch_label(_version.lower()) + or bool(_hip) + or bool(_cuda) + # torch.version.xpu, for the same reason .cuda and .hip are here: an untagged + # source, conda or private-index XPU build carries its runtime there and nowhere else. + or bool(_TORCH_RUNTIME_XPU) + ) + label = _installed_torch_label_on_disk() + return (not label) or _is_gpu_torch_label(label) + + +def _expected_torch_flavor_tag() -> str: + """The torch flavor this venv is SUPPOSED to hold, or "" when nothing can say. + + Resolution order, most authoritative first: + 1. UNSLOTH_EXPECTED_TORCH_TAG -- the setup script's own answer, exported by + setup.ps1 immediately before it hands over, so it describes the index the torch + install arm actually used (pin, preserved venv, GPU probe and all). + 2. An explicit index pin, when its family is one this vocabulary can name. The + manifest records what a PREVIOUS run installed; a pin is the instruction for + THIS one, so a freshly selected family has to outrank a stale record. Resolving + the manifest first let a cu128 pin lose to a cu124 manifest, and + _expected_torch_index_url then rejected the cu128 pin as a family mismatch and + repaired from the PUBLIC cu124 index -- undoing both the family and the source + the user had just chosen. + 3. The flavor the last completed install recorded in the manifest. Read at import + (_RECORDED_TORCH_TAG), because install_python_stack() drops the manifest before + the dependency pass. + 4. A live probe, for a run nothing set up: `python install_python_stack.py` by hand. + Only an NVIDIA host, or an explicit pin, can expect a GPU build -- otherwise + return "" rather than invent an expectation from an absent GPU. + """ + env = os.environ.get("UNSLOTH_EXPECTED_TORCH_TAG", "").strip().lower() + if env: + return env + pin = _explicit_torch_index_url() + if pin is not None: + leaf = _torch_index_leaf(pin) + # "rocm" names every AMD leaf (rocm6.4, gfx1151); an unreadable one falls through. + if _is_pip_rocm_family_leaf(leaf): + return "rocm" + if _is_cuda_family_leaf(leaf) or leaf in ("xpu", "cpu"): + return leaf + # A resolved backend is a stated choice, not a probe result: an AMD host taking setup.sh's + # documented UNSLOTH_TORCH_BACKEND=cpu recorded nothing here, because the NVIDIA probe + # answers "" for it, and the next launch called the deliberate install broken. Only when + # the family agrees with the wheel actually installed. Ahead of the manifest, which + # describes the PREVIOUS install, or a reinstall that changes flavor re-records the old tag. + if _TORCH_BACKEND in ("cpu", "rocm", "xpu"): + _installed = _torch_flavor_tag(_installed_torch_version_label()) + if _installed == _TORCH_BACKEND: + return _TORCH_BACKEND + if _RECORDED_TORCH_TAG: + return _RECORDED_TORCH_TAG + # An absent NVIDIA GPU with no pin means no CUDA expectation exists to enforce. + if _explicit_torch_index_url() is None and not _has_usable_nvidia_gpu(): + return "" + return _torch_index_leaf(_detect_cuda_torch_index_url()) + + +def _expected_torch_flavor_is_explicit() -> bool: + """Whether the expectation came from someone SAYING so, rather than from a probe. + + True for the setup script's handover, an explicit index pin, and the flavor the last + completed install recorded: the first three steps of _expected_torch_flavor_tag. + False when only the live hardware probe can answer, which is the one case a + visibility mask has any business overruling. + """ + if os.environ.get("UNSLOTH_EXPECTED_TORCH_TAG", "").strip(): + return True + if _explicit_torch_index_url() is not None: + return True + return bool(_RECORDED_TORCH_TAG) + + +def _recordable_torch_flavor_tag(resolved: str) -> str: + """The flavor worth writing to the manifest, or "" when nothing is. + + Normally the flavor this run resolved, falling back to the previous install's so an + update does not erase a record it simply had no occasion to recompute. An explicit + pin whose leaf names no family (a corporate /simple mirror, /current) breaks that + fallback: the wheel now in the venv came from that mirror, the old record describes + a venv that no longer exists, and carrying it forward would hand a later unpinned run + a flavor to "repair" the mirror's build back to. + """ + if resolved: + return resolved + if _explicit_unknown_family_torch_index_url() is not None: + return "" + return _RECORDED_TORCH_TAG or "" + + +def _index_leaf_flavor_family(leaf: str) -> str: + """The flavor family a pip index leaf names: cpu, xpu, rocm, cuda, or "" for none.""" + leaf = (leaf or "").strip().lower() + if leaf in ("cpu", "xpu"): + return leaf + if _is_pip_rocm_family_leaf(leaf): + return "rocm" + if _is_cuda_family_leaf(leaf): + return "cuda" + return "" + + +def _flavor_tag_family(tag: str) -> str: + """The family a flavor tag belongs to. cu124 and cu128 are both "cuda".""" + tag = (tag or "").strip().lower() + return "cuda" if tag.startswith("cu") else tag + + +def _expected_torch_flavor_was_pinned(flavor: str = "") -> bool: + """Whether ``flavor`` was NAMED by whoever ran this install. + + Distinct from _expected_torch_flavor_is_explicit(), which counts setup.ps1's + handover variable. setup.ps1 publishes that variable for an AUTOMATIC /cpu choice on + a GPU-less host exactly as it does for a pinned one, so the handover cannot answer + this question. An index pin, an index family, and UNSLOTH_TORCH_BACKEND all can: + each of them is someone saying which build they want. Carried forward from the + previous manifest, so an update that names nothing does not erase the fact. + + Each of them only answers for the family it NAMES. setup.ps1 falls back to the CPU + index when a pinned ROCm or XPU install fails, and publishes the resolved cpu tag + while the original GPU pin is still in the environment: counting that pin would + record a pinned CPU flavor, and _expected_cpu_flavor_was_chosen() would then read a + failed install as a deliberate one and suppress the repair guidance for good. + ``flavor`` empty means nobody asked about a specific one, and every pin counts. + """ + + def _names_it(family: str) -> bool: + return True if not flavor else family == _flavor_tag_family(flavor) + + pin = _explicit_torch_index_url() + if pin is not None and _names_it(_index_leaf_flavor_family(_torch_index_leaf(pin))): + return True + _family_var = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if _family_var and _names_it(_index_leaf_flavor_family(_torch_index_leaf(_family_var))): + return True + # install.sh derives UNSLOTH_TORCH_BACKEND from the index it RESOLVED -- "cpu" on any GPU-less + # machine, asked for or not -- and marks it derived. Only an unmarked value is a preference. + if ( + _TORCH_BACKEND in ("cpu", "cuda", "rocm", "xpu") + and os.environ.get("UNSLOTH_TORCH_BACKEND_SOURCE", "").strip().lower() != "resolved" + and _names_it(_TORCH_BACKEND) + ): + return True + return bool(_RECORDED_TORCH_TAG_PINNED) and _names_it( + _flavor_tag_family(_RECORDED_TORCH_TAG or "") + ) + + +def _expected_torch_index_url(tag: str) -> str: + """The wheel index to repair `tag` from. + + Prefers the exact URL the setup script installed from (UNSLOTH_TORCH_INSTALL_INDEX_URL) + and then the explicit pin, because an authenticated mirror can only be repaired from + the credentialed URL -- neither is reconstructible from a family leaf. Both are only + used when their leaf IS this tag: setup.ps1 sends the /cpu index alongside a "rocm" + tag on the AMD Windows path, and repairing a cu* mismatch from that URL would install + the very CPU wheel this exists to remove. Otherwise rebuild it the way setup.ps1 does + when nothing is pinned: /. + """ + url = os.environ.get("UNSLOTH_TORCH_INSTALL_INDEX_URL", "").strip() + if url: + url = _trim_index_path_slashes(url) + if _torch_index_leaf(url) == tag: + return url + pin = _explicit_torch_index_url() + if pin is not None and _torch_index_leaf(pin) == tag: + return pin + return f"{_PYTORCH_WHL_BASE}/{tag}" + + +def _explicit_cpu_torch_index_pin() -> bool: + """Whether this run was pinned to a CPU wheel index, by URL or by family. + + Only a pin counts. setup.ps1's published tag also reads "cpu" for a host whose + nvidia-smi probe simply returned nothing, and treating that as an instruction would + let a wedged driver downgrade a healthy CUDA venv. + """ + pin = _explicit_torch_index_url() + return pin is not None and _torch_index_leaf(pin) == "cpu" + + +def _installed_flavor_tag_now(expected: str = "") -> str: + """The venv's CURRENT flavor tag, re-probed, or "" when nothing could be read. + + ``expected`` only matters for "cpu": _torch_flavor_tag reads every untagged version + as cpu, so a private index serving an untagged CUDA or ROCm wheel would satisfy a + CPU expectation here exactly as it did in the pre-repair comparison. The same + runtime markers settle it, and the two comparisons have to agree or the repair is + accepted on a rule the check that triggered it rejected. + + _torch_build_is_gpu answers a weaker question ("can this torch use a GPU at all") + and is deliberately family-blind, so it cannot tell a repair that installed the + requested family from one that left a different GPU wheel in place. "" is returned + for an unreadable venv rather than "cpu", so ambiguity stays distinguishable from + a positive CPU reading and never fails an update on its own. + """ + _ran, _importable, _version, _hip, _cuda = _probe_torch_runtime() + if _ran and _importable and _version: + tag = _torch_flavor_tag(_version) + if expected == "cpu" and tag == "cpu" and (_hip or _cuda or _TORCH_RUNTIME_XPU): + return _gpu_family_from_runtime_markers(_hip, _cuda) + return tag + label = _installed_torch_label_on_disk() + return _torch_flavor_tag(label) if label else "" + + +def _warn_wrong_flavor(expected: str, installed: str) -> bool: + """Report a repair that installed the wrong family, and fail. Always returns False. + + Distinct from _warn_still_cpu because "PyTorch is CPU-only" would be false here and + would send the reader after the wrong problem: the venv holds a GPU build, just not + the one this host asked for. + """ + _safe_print("") + _safe_print( + f" [WARN] PyTorch is a {installed} build but {expected} was expected for this machine." + ) + _safe_print(" [WARN] The repair did not install the requested build.") + _safe_print(" [WARN] Re-run this installer, or reinstall the build for your GPU manually.") + _safe_print(" [WARN] irm https://unsloth.ai/install.ps1 | iex") + return False + + +def _warn_still_cpu(expected: str) -> bool: + """Report a repair that did not take, and fail the install. Always returns False. + + install.ps1 warns here and exits 0. Failing instead is the entire point: a CPU-only + torch on a host that expects a GPU build is precisely the state an update used to + report as "dependencies up to date" while the app ran everything on CPU. + """ + _safe_print("") + _safe_print( + f" [WARN] PyTorch is CPU-only but a {expected} GPU build was expected for this machine." + ) + _safe_print(" [WARN] Training and GPU inference will run on CPU until this is fixed.") + _safe_print( + " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." + ) + _safe_print(" [WARN] irm https://unsloth.ai/install.ps1 | iex") + return False + + +def _uninstall_distribution(name: str) -> bool: + """Remove one distribution from the venv this script targets. True iff it is gone. + + Same shape as the flash-attn removal below: --python sys.executable so a uv that + also needs --system cannot remove from the system Python instead, and a pip + fallback for the same interpreter. Output is swallowed; the caller reports. + """ + if USE_UV and shutil.which("uv"): + cmd = ["uv", "pip", "uninstall"] + if UV_NEEDS_SYSTEM: + cmd.append("--system") + cmd.extend(["--python", sys.executable, name]) + else: + cmd = [sys.executable, "-m", "pip", "uninstall", "-y", name] + removed = subprocess.run(cmd, stdout = subprocess.DEVNULL, stderr = subprocess.DEVNULL) + return removed.returncode == 0 + + +def _resident_xformers_build_torch() -> "str | None": + """The torch build the installed xFormers extension was compiled against. + + Read from ``xformers/cpp_lib.json``, the same file install.ps1's resident probe + reads. None when xFormers is absent or carries no build metadata. Never raises, and + never imports xformers -- a mismatched _C.pyd logs its own warning on import, which + would land in the middle of the installer's output. + """ + try: + spec = importlib.util.find_spec("xformers") + except Exception: + return None + locations = list(getattr(spec, "submodule_search_locations", None) or []) if spec else [] + if not locations: + return None + try: + with open(os.path.join(locations[0], "cpp_lib.json"), encoding = "utf-8") as fh: + recorded = json.load(fh).get("version", {}).get("torch") + except (OSError, ValueError, AttributeError): + return None + return recorded.strip() if isinstance(recorded, str) and recorded.strip() else None + + +def _resync_torch_coupled_packages(label_before: str) -> bool: + """Re-settle the packages whose compiled extensions are tied to the torch build. + + Returns False when this pass left the venv in a state the caller must re-verify. + + torchao's cpp extensions are tied to the torch release AND its CUDA major, both of + which _select_torchao_spec branches on; xFormers is tied to the exact (torch, CUDA) + pair, and beside a pair it was not built for its ops vanish behind a log line rather + than an error. --no-deps is the whole safety of the torchao call: torchao depends on + torch, so resolving dependencies would pull PyPI's CPU wheel back in. Never fatal -- + both are secondary to the flavor repair that has just succeeded. + """ + _label_after = str(_probe_installed_torch_version() or "") + if not _label_after or _label_after == label_before: + return True + _touched_torch = False + # Release OR CUDA major: cu124 to cu130 at one release still changes the build. + _release_moved = _label_after.split("+", 1)[0] != str(label_before).split("+", 1)[0] + _cuda_moved = _cuda_major_from_torch_version(_label_after) != ( + _cuda_major_from_torch_version(str(label_before)) + ) + if _release_moved or _cuda_moved: + try: + _spec = _select_torchao_spec(_label_after) + if not _exact_distribution_spec_is_installed(_spec): + _note(f"torch {_label_after} after repair -- reinstalling {_spec}") + _touched_torch = True + if not pip_install_try( + "Re-matching torchao to the repaired torch", + "--force-reinstall", + "--no-deps", + "--no-cache-dir", + _spec, + ): + # Returns False rather than raising: an unreachable PyPI would + # otherwise leave the incompatible build in place silently. + _safe_print( + f" [WARN] could not install {_spec} for the repaired torch; the " + f"torchao kernels will fall back to the slow path." + ) + except Exception as e: + _safe_print(f" [WARN] could not re-match torchao after the repair: {e}") + try: + _built_for = _resident_xformers_build_torch() + if _built_for and _built_for != _label_after: + _note( + f"xFormers was built for torch {_built_for}, which is no longer " + f"installed -- removing it so attention falls back to torch SDPA" + ) + if not _uninstall_distribution("xformers"): + _safe_print( + " [WARN] could not remove the mismatched xFormers; its compiled " + "operations will stay unavailable until it is uninstalled by hand." + ) + except Exception as e: + _safe_print(f" [WARN] could not re-check xFormers after the repair: {e}") + return not _touched_torch + + +def _ensure_expected_torch_flavor(expected: "str | None" = None) -> bool: + """Enforce that the venv still holds the torch flavor the install selected. + + `unsloth studio update` runs setup.ps1 and this script, never install.ps1, which held + the only flavor repair; the dependency steps above resolve torch from PyPI, whose + Windows wheel is 2.11.0+cpu. Returns False when the flavor is wrong and could not be + repaired, which fails the install: that state used to be reported as success. + + ROCm is delegated to _ensure_rocm_torch -- AMD's Windows wheels live on a + per-architecture repo.amd.com index a generic "rocm" tag cannot reconstruct. + """ + if NO_TORCH: + return True + # rocm/xpu/cpu fall THROUGH: an explicit GPU pin sets _TORCH_BACKEND, and rejecting it + # here would skip the invariant on the hosts that asked for that family. + if _TORCH_BACKEND not in ("", "cuda", "rocm", "xpu", "cpu"): + return True + if expected is None: + expected = _expected_torch_flavor_tag() + # The PIN, not the handover: setup.ps1 also publishes "cpu" when its nvidia-smi probe + # comes back empty, and that host must not be downgraded. + _cpu_pinned = expected == "cpu" and _explicit_cpu_torch_index_pin() + if not (_is_cuda_family_leaf(expected) or expected in ("xpu", "rocm") or _cpu_pinned): + return True + if _TORCH_BACKEND in ("rocm", "xpu", "cpu") and _TORCH_BACKEND != expected: + return True + # A pin whose leaf names no flavor was applied verbatim at install time, so acting on an + # older manifest overrides an administrator's mirror. Compared, not vetoed: the helper's + # known set predates XPU. + _unknown_pin = _explicit_unknown_family_torch_index_url() + if _unknown_pin is not None and _torch_index_leaf(_unknown_pin) != expected: + return True + # CUDA only, and only for an expectation INFERRED from hardware: an emptied mask is a + # reason not to conclude cu124 from a probe, not to ignore a stated one. + if _is_cuda_family_leaf(expected) and not _expected_torch_flavor_is_explicit(): + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if _cvd is not None and _cvd.strip() in ("", "-1"): + return True + + _ran, _importable, _version, _hip, _cuda = _probe_torch_runtime() + if _ran and _importable and _version: + installed_version = _version + elif not _ran: + # A wedged driver hangs `import torch`; version.py names the wheel without one. + installed_version = _installed_torch_label_on_disk() + else: + # Missing or unimportable is the base install's job, a louder failure than this. + return True + if not installed_version: + return True + + installed = _torch_flavor_tag(installed_version) + # _torch_flavor_tag reads untagged as "cpu", so a private index's untagged CUDA build + # compares equal under a /cpu pin. + if expected == "cpu" and installed == "cpu" and (_hip or _cuda or _TORCH_RUNTIME_XPU): + installed = _gpu_family_from_runtime_markers(_hip, _cuda) + if installed == expected: + return True + + # install.ps1's line, word for word, so a support log from either path reads the same. + _safe_print( + f" PyTorch flavor mismatch (installed {installed}, need {expected}) -- " + f"reinstalling correct build..." + ) + if expected == "rocm": + # AMD's Windows wheels live on a per-architecture repo.amd.com index no "rocm" tag + # can name, and the handed-over URL still points at /cpu there. + _ensure_rocm_torch() + # The FAMILY: a transient repo.amd.com failure is non-fatal in there, and the cu124 + # wheel it leaves passes _torch_build_is_gpu. + _now = _installed_flavor_tag_now(expected) + if _now == expected: + return True + if not _now: + # Ambiguity must not fail an update by itself. + return True if _torch_build_is_gpu() else _warn_still_cpu(expected) + if not _torch_build_is_gpu(): + return _warn_still_cpu(expected) + return _warn_wrong_flavor(expected, _now) + + index_url = _expected_torch_index_url(expected) + # XPU floor is 2.6, not 2.4: unsloth/models/_utils.py raises at import below it. + _torch_pkg, _vision_pkg, _audio_pkg = ( + _XPU_TORCH_PKG_SPEC if expected == "xpu" else _TORCH_FLAVOR_REPAIR_PKG_SPEC + ) + # No win_arm64 torchaudio wheel exists on any index ($WinArm64NoAudio in setup.ps1). + _trio = [_torch_pkg, _vision_pkg, _audio_pkg] + if _is_windows_arm64(): + _trio = [_torch_pkg, _vision_pkg] + _label_before = str(installed_version) + # --force-reinstall, not install.ps1's uv-only --reinstall-package: pip_install falls back + # to pip. constrain=False: constraints.txt resolves against the PyPI torch that did this. + pip_install( + "PyTorch flavor repair", + "--force-reinstall", + "--no-cache-dir", + *_trio, + "--index-url", + index_url, + constrain = False, + ) + + # The family: a mirror can answer /cu128 with a cached cu124 wheel, and + # _torch_build_is_gpu is family-blind. + _now = _installed_flavor_tag_now(expected) + if _now == expected: + if _resync_torch_coupled_packages(_label_before): + return True + # The resync installs --no-deps, but this function's verification is behind us. + _after = _installed_flavor_tag_now(expected) + if _after in (expected, ""): + return True + _safe_print(" [WARN] the post-repair package resync changed the torch build.") + return _warn_wrong_flavor(expected, _after) + if not _now: + # Ambiguity must not fail an update by itself. + if expected == "cpu": + return True + return True if _torch_build_is_gpu() else _warn_still_cpu(expected) + if expected != "cpu" and not _torch_build_is_gpu(): + return _warn_still_cpu(expected) + return _warn_wrong_flavor(expected, _now) + + def _amd_torch_needs_dependency_pass() -> bool: """Return True when setup must run the dependency pass to repair non-ROCm torch. @@ -2991,6 +3537,11 @@ def _ensure_rocm_torch() -> None: _torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get( gfx_arch, ("torch", "torchvision", "torchaudio") ) + # Same win_arm64 exception setup.ps1 applies: no torchaudio wheel exists + # there, so asking for one makes the trio unresolvable. + _rocm_trio = [_torch_pkg, _vision_pkg, _audio_pkg] + if _is_windows_arm64(): + _rocm_trio = [_torch_pkg, _vision_pkg] # Nonfatal: a transient AMD-index failure must not abort the install. # --force-reinstall resolves before uninstalling, so a failed index keeps the # existing build intact; let the user retry. @@ -2999,9 +3550,7 @@ def _ensure_rocm_torch() -> None: "--force-reinstall", "--index-url", index_url, - _torch_pkg, - _vision_pkg, - _audio_pkg, + *_rocm_trio, constrain = False, ): _safe_print( @@ -3434,6 +3983,12 @@ def _infer_no_torch() -> bool: NO_TORCH = _infer_no_torch() +# Read at import: install_python_stack() drops the manifest before its dependency pass. +_RECORDED_TORCH_TAG = install_manifest.recorded_torch_flavor() +# Whether that record came from someone NAMING a flavor: setup.ps1 publishes an automatic +# /cpu choice the same way it publishes a pinned one. +_RECORDED_TORCH_TAG_PINNED = install_manifest.recorded_torch_flavor_was_pinned() + # UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() ("cuda", "rocm", # "cpu"; empty = standalone `studio update`, where we re-detect). _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() @@ -5274,6 +5829,8 @@ def install_python_stack() -> int: base_total += 1 # ROCm torch check (step 2b), non-macOS if not IS_WINDOWS: base_total += 2 # flash-attn + torch final repair (step 13), Linux + else: + base_total += 1 # torch flavor invariant (step 13w), Windows if IS_MAC_ARM and not skip_base: base_total += 1 # MLX stack, same gate as the step itself base_requirements = _shared_base_requirements() if skip_base else None @@ -5685,6 +6242,7 @@ def _win_amd_smi_has_gpu(stdout: str) -> bool: ) # 13. Final torch repair. Steps above can pull CUDA torch from PyPI, so repair last. + torch_flavor_tag = "" if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("final")) _ensure_cuda_torch() @@ -5695,6 +6253,22 @@ def _win_amd_smi_has_gpu(stdout: str) -> bool: # CPU pin over an XPU venv would leave XPU triton under a CPU torch. _ensure_xpu_triton() + # 13w. Windows torch flavor invariant, separate from step 13's Linux-shaped repair set + # but in the same position: last, after the with-deps steps re-resolved torch. + if IS_WINDOWS and not NO_TORCH: + _progress(_torch_step_label("flavor")) + torch_flavor_tag = _expected_torch_flavor_tag() + if not _ensure_expected_torch_flavor(torch_flavor_tag): + return 1 + # A direct run has no setup.ps1 postlude to swap triton back. After the invariant, + # because the swap keys off the installed +xpu label. + _ensure_xpu_triton() + elif not NO_TORCH: + # Resolve it on the other platforms too, for the RECORD only. Without this a Linux GPU box + # installed with a transient explicit CPU pin looks, on the next launch, like a CPU wheel + # beside a physical GPU. + torch_flavor_tag = _expected_torch_flavor_tag() + # 14. Final check (silent; third-party conflicts are expected) subprocess.run( [sys.executable, "-m", "pip", "check"], @@ -5724,6 +6298,12 @@ def _win_amd_smi_has_gpu(stdout: str) -> bool: steps_total = _TOTAL, package_name = package_name, no_torch = NO_TORCH, + # A platform that never resolves a flavor carries the old record forward. An unknown-family + # pin is the exception: the previous record describes a venv that no longer exists, and + # writing it back would give a later unpinned run a flavor to "repair" the mirror's to. + expected_torch_tag = _recordable_torch_flavor_tag(torch_flavor_tag), + expected_torch_tag_pinned = bool(_recordable_torch_flavor_tag(torch_flavor_tag)) + and _expected_torch_flavor_was_pinned(_recordable_torch_flavor_tag(torch_flavor_tag)), ) is None ): diff --git a/studio/setup.ps1 b/studio/setup.ps1 index aac61eb140a..1064b244111 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1052,7 +1052,9 @@ function Get-RocmPinStaleTags { # that answer. Mirrors install.ps1's copy. function Invoke-BoundedPythonProbe { param([string]$PythonExe, [string]$Code, [int]$TimeoutSec = 30) - $result = [pscustomobject]@{ Ok = $false; Output = ""; Error = "" } + # TimedOut separates "never answered" from "answered with a failure": both leave Ok + # false, only the second says anything about the installation. + $result = [pscustomobject]@{ Ok = $false; Output = ""; Error = ""; TimedOut = $false } if (-not $PythonExe -or -not $Code) { return $result } try { $psi = New-Object System.Diagnostics.ProcessStartInfo @@ -1070,6 +1072,7 @@ function Invoke-BoundedPythonProbe { # Synthesised, not read back: waiting on the reader tasks of a wedged child would # reintroduce the hang this helper exists to bound. $result.Error = "python did not answer within $TimeoutSec seconds" + $result.TimedOut = $true return $result } $result.Output = $outTask.GetAwaiter().GetResult() @@ -1231,6 +1234,30 @@ function Test-VenvTorchIsRocm { } catch { return $false } } +# The NVIDIA counterpart of the XPU and ROCm on-disk rescues: a wedged display driver hangs or +# faults `import torch`, and the chain then fell through to a rebuild with a null tag, deleting a +# healthy CUDA venv with no rollback copy. Returns the FAMILY, since the stale check needs it. +function Get-VenvTorchCudaTag { + param([string]$VenvPath) + if (-not $VenvPath) { return $null } + try { + $verPy = Join-Path $VenvPath "Lib\site-packages\torch\version.py" + if (-not (Test-Path -LiteralPath $verPy)) { return $null } + $line = (Get-Content -LiteralPath $verPy -TotalCount 40 -ErrorAction Stop | + Where-Object { $_ -match "__version__\s*=\s*'[^']*\+(cu[0-9]+)" } | + Select-Object -First 1) + if (-not $line) { return $null } + $m = [regex]::Match($line, "__version__\s*=\s*'[^']*\+(cu[0-9]+)") + if ($m.Success) { return $m.Groups[1].Value.ToLowerInvariant() } + return $null + } catch { return $null } +} + +function Test-VenvTorchIsCuda { + param([string]$VenvPath) + return [bool](Get-VenvTorchCudaTag -VenvPath $VenvPath) +} + # Same free disk read, plus the supported range: unsloth/models/_utils.py raises at import for an # XPU device on torch < 2.6, so flavour alone would call a 2.5+xpu venv fine; 2.11 is the trio's # ceiling. Flavour and range ONLY -- whether the runtime reaches the GPU is a driver question. @@ -4048,6 +4075,7 @@ $installedTorchTag = $null # installer-managed repair below raises it, yet only the block a fresh install never enters # assigns it. $script:PinChangedForceReinstall = $false +$script:TorchImportDefinitivelyFailed = $false if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode) { $VenvPyExe = Join-Path $VenvDir "Scripts\python.exe" $installedTorchTag = $null @@ -4098,6 +4126,18 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode $installedTorchTag = "rocm" substep "PyTorch did not respond but this venv holds a ROCm build -- keeping it." "Yellow" substep "If training fails, reboot and update the AMD Adrenalin / HIP SDK driver." "Yellow" + } elseif (Test-VenvTorchIsCuda -VenvPath $VenvDir) { + # Without this the chain fell through with a NULL tag, so the no-wipe escape + # below saw no cu* wheel to preserve. Keep the FAMILY, not a generic "cuda". + $installedTorchTag = Get-VenvTorchCudaTag -VenvPath $VenvDir + substep "PyTorch did not respond but this venv holds a $installedTorchTag build -- keeping it." "Yellow" + substep "If training fails, reboot and update the NVIDIA driver." "Yellow" + # A half-written torch also leaves a +cu* version.py behind, and the matched + # install below would write a completion manifest over it. Force the reinstall. + if ($_verProbe -and -not $_verProbe.TimedOut) { + $script:TorchImportDefinitivelyFailed = $true + substep "PyTorch failed to import rather than timing out -- reinstalling the same family in place." "Yellow" + } } else { $shouldRebuild = $true } @@ -4288,6 +4328,23 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode $shouldRebuild = $false } + # A cu* venv is never wiped by a DIRECT update just because nvidia-smi did not answer: + # every way that bounded probe comes back empty on a working NVIDIA host collapses + # $expectedTorchTag to "cpu", no escape above catches it, and the wipe has no rollback + # copy. Narrow on purpose: unpinned only, cu* installed only, and only when the collapse + # was for want of an NVIDIA answer. The -and order also keeps the reads legal under + # Set-StrictMode, since both are assigned inside `if (-not $shouldRebuild)`. + if ($shouldRebuild -and -not $InstallerManagedSetup -and + $installedTorchTag -and (Test-CudaFamilyLeaf $installedTorchTag) -and + -not $_pinnedIdx -and -not $HasNvidiaSmi -and $expectedTorchTag -eq "cpu") { + substep "nvidia-smi did not answer, but this venv holds a $installedTorchTag build -- keeping it." "Yellow" + substep "If training runs on CPU, re-run install.ps1: irm https://unsloth.ai/install.ps1 | iex" "Yellow" + $shouldRebuild = $false + # Keeping the wheel is half the job: the index selection below rescans, sees no + # NVIDIA either, and would route the install to the /cpu arm. + $script:PreservedInstallerTorchTag = $installedTorchTag + } + if ($shouldRebuild) { substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" # why: mirror install.ps1 env-mode guard so an update against a custom @@ -4945,7 +5002,13 @@ sys.exit(0 if installed is not None and required is not None and installed >= re # A torch-index pin change repairs in place: force the dependency pass so the torch install # below force-reinstalls from the new pin (else the fast path keeps the old wheel). -if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false } +# A torch that will not import needs the same pass for the same reason, and it is the only +# thing that can act on it: its --force-reinstall lives INSIDE this block, so on a current +# core package with a valid manifest the fast path skipped the repair entirely and the +# update reported dependencies up to date over a broken wheel. +if ($script:PinChangedForceReinstall -or $script:TorchImportDefinitivelyFailed) { + $SkipPythonDeps = $false +} if (-not $SkipPythonDeps) { @@ -5269,7 +5332,9 @@ if (-not $ROCmIndexUrl -and -not $XpuIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCp # requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126 # -> cu128) never applies. $cudaForce = @() - if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") } + if ($script:PinChangedForceReinstall -or $script:TorchImportDefinitivelyFailed) { + $cudaForce = @("--force-reinstall") + } # An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound # the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion # against the capped torch. Known cu* leaves keep bare specs. @@ -5329,6 +5394,33 @@ if (-not $ROCmIndexUrl -and -not $XpuIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCp # Duplicate metadata repair is the one pass that DOES reinstall unsloth under # SKIP_STUDIO_BASE, so the CLI wraps this script in its launcher transaction. +# ── Publish the torch flavor this run settled on ── +# so install_python_stack.py can enforce it: its dependency steps resolve torch from PyPI, whose +# Windows wheel is 2.11.0+cpu, and only install.ps1 -- never on the updater's path -- repaired +# that. Vocabulary is Get-InstalledTorchTag's; an unknown leaf publishes nothing. +if (-not $NoTorchMode) { + $_expectedLeaf = Get-TorchIndexLeaf $TorchInstallIndexUrl + # $ROCmIndexUrl first: on the AMD path $TorchInstallIndexUrl still points at /cpu. + $_expectedTag = if ($ROCmIndexUrl) { "rocm" } + elseif (Test-CudaFamilyLeaf $_expectedLeaf) { $_expectedLeaf } + elseif ($_expectedLeaf -eq "cpu" -or $_expectedLeaf -eq "xpu") { $_expectedLeaf } + elseif (Test-PipRocmFamilyLeaf $_expectedLeaf) { "rocm" } + else { $null } + # Remove-Item, NOT `= ""`: PowerShell 7.5 took .NET 9's change and KEEPS a name + # assigned an empty string. Deleting also stops a value inherited from the caller's + # shell surviving a run that decided it cannot name this host's flavor. + if ($_expectedTag) { + $env:UNSLOTH_EXPECTED_TORCH_TAG = $_expectedTag + } else { + Remove-Item Env:\UNSLOTH_EXPECTED_TORCH_TAG -ErrorAction SilentlyContinue + } + if ($TorchInstallIndexUrl) { + $env:UNSLOTH_TORCH_INSTALL_INDEX_URL = $TorchInstallIndexUrl + } else { + Remove-Item Env:\UNSLOTH_TORCH_INSTALL_INDEX_URL -ErrorAction SilentlyContinue + } +} + # Ordered heavy dependency installation -- shared cross-platform script substep "running ordered dependency installation..." python "$PSScriptRoot\install_python_stack.py" diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index ee7206ecd2c..ba16a1493f1 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -933,6 +933,15 @@ pub async fn native_path_leases_usable( ) } +/// `force_installer` skips the update attempt and runs the bundled installer directly. +/// +/// The automatic repair tries `studio update` first because it is the cheap fix for a venv +/// one release behind. But an update reuses the environment it finds, so a managed venv +/// whose PyTorch has been replaced by a CPU-only wheel comes back from a successful update +/// still CPU-only. Only install.ps1 / install.sh re-select the torch index. +/// +/// Settings' manual "Repair installation" therefore passes true. Absent (the startup +/// auto-repair path) reads as false, so that path is unchanged. #[tauri::command] pub async fn start_managed_repair( app: AppHandle, @@ -941,8 +950,13 @@ pub async fn start_managed_repair( update_state: tauri::State<'_, update::UpdateState>, install_state: tauri::State<'_, install::InstallState>, diagnostics: tauri::State<'_, DiagnosticsState>, + force_installer: Option, ) -> Result<(), String> { - info!("start_managed_repair command called"); + let force_installer = force_installer.unwrap_or(false); + info!( + "start_managed_repair command called (force_installer={})", + force_installer + ); if install_state .lock() @@ -979,29 +993,38 @@ pub async fn start_managed_repair( let repair_group_id = install::take_pending_repair_group_for_resume(&install_state) .unwrap_or_else(|| diagnostics::begin_repair_group(&diagnostics_state)); - let _ = app.emit("repair-progress", "Updating existing Unsloth install..."); - let update_app = app.clone(); - let update_state = update_state.inner().clone(); - let update_diagnostics = diagnostics_state.clone(); - let update_repair_group_id = repair_group_id.clone(); - let update_result = tokio::task::spawn_blocking(move || { - update::run_backend_update_for_repair( - update_app, - update_state, - update_diagnostics, - update_repair_group_id, - ) - }) - .await - .map_err(|e| format!("Repair update task panicked: {e}"))?; + // Ok(()) rather than skipping the match: a second copy of the fallback under a + // `if force_installer` would be the one place a future change could miss. + let update_result = if force_installer { + let _ = app.emit("repair-progress", "Running bundled installer..."); + Ok(()) + } else { + let _ = app.emit("repair-progress", "Updating existing Unsloth install..."); + let update_app = app.clone(); + let update_state = update_state.inner().clone(); + let update_diagnostics = diagnostics_state.clone(); + let update_repair_group_id = repair_group_id.clone(); + tokio::task::spawn_blocking(move || { + update::run_backend_update_for_repair( + update_app, + update_state, + update_diagnostics, + update_repair_group_id, + ) + }) + .await + .map_err(|e| format!("Repair update task panicked: {e}"))? + }; match update_result { - Ok(()) if managed_install_ready_after_repair().await => { + Ok(()) if !force_installer && managed_install_ready_after_repair().await => { info!("Managed repair complete after update"); diagnostics::finish_repair_group(&diagnostics_state, &repair_group_id, "success", None); let _ = app.emit("repair-complete", ()); return Ok(()); } + // The forced path already emitted its own progress line and ran no update. + Ok(()) if force_installer => {} Ok(()) => { warn!("Managed repair update finished, but preflight is still not ready; falling back to installer"); let _ = app.emit( diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index 4b04f27a2f6..8b7098e103d 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -2371,3 +2371,98 @@ def fake_run( assert excinfo.value.code == 1 assert len(attempts) == 1, "a failure with nothing to clear must not be retried" + + +class TestExpectedTorchFlavorResolution: + """_expected_torch_flavor_tag / _expected_torch_index_url: the two pure inputs to the + Windows flavor invariant. The invariant itself is covered in + tests/studio/install/test_cuda_repair.py; these pin the resolution ORDER, which is what + decides whether a repair fires against the right index or not at all.""" + + _KEYS = ( + "UNSLOTH_EXPECTED_TORCH_TAG", + "UNSLOTH_TORCH_INSTALL_INDEX_URL", + "UNSLOTH_TORCH_INDEX_URL", + "UNSLOTH_TORCH_INDEX_FAMILY", + ) + + @contextlib.contextmanager + def _env(self, **values): + """Set the named vars and REMOVE every other one this resolution reads, so an + ambient pin on the developer's box cannot change the answer.""" + with mock.patch.dict(os.environ, {k: v for k, v in values.items() if v is not None}): + for key in self._KEYS: + if values.get(key) is None: + os.environ.pop(key, None) + yield + + def test_the_handover_tag_wins(self): + with self._env(UNSLOTH_EXPECTED_TORCH_TAG = "cu124"): + with mock.patch.object(ips, "_RECORDED_TORCH_TAG", "cu128"): + assert ips._expected_torch_flavor_tag() == "cu124" + + def test_the_handover_tag_is_normalised(self): + with self._env(UNSLOTH_EXPECTED_TORCH_TAG = " CU128 "): + assert ips._expected_torch_flavor_tag() == "cu128" + + def test_the_manifest_answers_next(self): + with self._env(): + with mock.patch.object(ips, "_RECORDED_TORCH_TAG", "cu128"): + assert ips._expected_torch_flavor_tag() == "cu128" + + def test_a_gpuless_host_with_nothing_recorded_says_nothing(self): + # Inventing a CUDA expectation from an absent GPU would reinstall CUDA torch + # onto a CPU box on every update. + with self._env(): + with ( + mock.patch.object(ips, "_RECORDED_TORCH_TAG", None), + mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False), + ): + assert ips._expected_torch_flavor_tag() == "" + + def test_a_pin_answers_without_probing_the_gpu(self): + with self._env(UNSLOTH_TORCH_INDEX_FAMILY = "cu126"): + with ( + mock.patch.object(ips, "_RECORDED_TORCH_TAG", None), + mock.patch.object(ips, "_has_usable_nvidia_gpu") as probe, + ): + assert ips._expected_torch_flavor_tag() == "cu126" + probe.assert_not_called() + + def test_a_cpu_pin_resolves_to_cpu_not_to_the_host_gpu(self): + with self._env(UNSLOTH_TORCH_INDEX_FAMILY = "cpu"): + with mock.patch.object(ips, "_RECORDED_TORCH_TAG", None): + assert ips._expected_torch_flavor_tag() == "cpu" + + def test_the_index_url_is_reused_only_for_its_own_family(self): + # setup.ps1 hands over the /cpu index alongside a "rocm" tag on AMD Windows, so + # repairing from it would install the very CPU wheel the repair exists to remove. + with self._env(UNSLOTH_TORCH_INSTALL_INDEX_URL = "https://mirror.local/whl/cu124/"): + assert ips._expected_torch_index_url("cu124") == "https://mirror.local/whl/cu124" + with self._env(UNSLOTH_TORCH_INSTALL_INDEX_URL = "https://download.pytorch.org/whl/cpu"): + assert ips._expected_torch_index_url("cu124") == f"{ips._PYTORCH_WHL_BASE}/cu124" + + def test_a_credentialed_index_survives_intact(self): + # Why the URL is forwarded rather than rebuilt: userinfo and a token query are + # not reconstructible from a family leaf. + url = "https://user:tok@mirror.local/whl/cu128?token=abc" + with self._env(UNSLOTH_TORCH_INSTALL_INDEX_URL = url): + assert ips._expected_torch_index_url("cu128") == url + + def test_the_pin_supplies_the_index_when_the_setup_script_did_not(self): + with self._env(UNSLOTH_TORCH_INDEX_URL = "https://mirror.local/whl/cu126"): + assert ips._expected_torch_index_url("cu126") == "https://mirror.local/whl/cu126" + + def test_the_default_index_is_the_pytorch_mirror(self): + with self._env(): + assert ips._expected_torch_index_url("cu124") == f"{ips._PYTORCH_WHL_BASE}/cu124" + + def test_no_index_url_is_ever_persisted_by_the_manifest_write(self): + # The manifest lives in the venv, so a token in a pinned URL must not reach it. + source = inspect.getsource(ips.install_python_stack) + assert "expected_torch_tag = _recordable_torch_flavor_tag(torch_flavor_tag)," in source + assert "torch_index_url" not in source + # And the helper that answers it records a FLAVOR, never a URL, for the same + # reason: it is reached with the pin still in the environment. + helper = inspect.getsource(ips._recordable_torch_flavor_tag) + assert "return" in helper and "_explicit_torch_index_url()" not in helper diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index da6b46a4a59..91b1ae5921e 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -1,10 +1,18 @@ """_ensure_cuda_torch reinstalls CUDA torch when an NVIDIA-host venv carries a ROCm build (the pre-fix KFD gpu_id false positive), but leaves healthy CUDA / CPU / ROCm / -macOS / Windows untouched. Fully mocked -- no GPU required.""" +macOS / Windows untouched. Fully mocked -- no GPU required. + +Also covers _ensure_expected_torch_flavor, the Windows counterpart: _ensure_cuda_torch +returns early on Windows because setup.ps1 owns torch there, which left the update path +with no flavor invariant at all. See the bottom of this file.""" import importlib.util +import inspect +import re +import subprocess import sys from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -21,6 +29,8 @@ sys.modules[_STACK_SPEC.name] = stack_mod _STACK_SPEC.loader.exec_module(stack_mod) +_SETUP_SRC = (PACKAGE_ROOT / "studio" / "setup.ps1").read_text(encoding = "utf-8") + # The probe prints its answer behind this marker, so chatter on either side of it cannot # be mistaken for the answer. Mocked stdout has to carry it too. _MARK = stack_mod._TORCH_PROBE_MARKER @@ -700,5 +710,1295 @@ def _sms(stdout, returncode = 0): assert _sms("7.0\n", returncode = 1) is None +# The updater runs setup.ps1 -> install_python_stack.py, never install.ps1, which held the +# only flavor repair. + +_ensure_expected_torch_flavor = stack_mod._ensure_expected_torch_flavor +_UNSET = object() + + +def _flavor_probe_stdout(version): + """The shared probe's marked line for a literal torch.__version__. + + torch.version.cuda / .hip follow the local label, because a real wheel sets them + together and _torch_build_is_gpu reads all three. + """ + match = re.search(r"\+(cu\d+)", version) + cuda = _dotted_cuda(match.group(1)) if match else "" + hip = "6.4" if "+rocm" in version else "" + return _MARK + f"{version}|{hip}|{cuda}\n" + + +def _run_flavor_invariant( + *, + installed = "2.11.0+cpu", + repaired = None, + expected_env = "cu124", + recorded = None, + install_index_url = None, + backend = "", + no_torch = False, + nvidia = True, + cuda_version = "12.4", + torch_rc = 0, + probe_timeout = False, + disk_label = _UNSET, + cvd = None, + index_family = None, + index_url = None, + win_arm64 = False, + probe_cuda = None, + probe_hip = None, +): + """Invoke _ensure_expected_torch_flavor against a fully mocked venv. + + `installed` is torch.__version__ before the pass. `repaired` is what the mocked + pip_install leaves behind: None means the reinstall changed nothing, which is the + state that must FAIL the update rather than report success. + + `expected_env` sets UNSLOTH_EXPECTED_TORCH_TAG (setup.ps1's handover), `recorded` the + flavor read out of the previous manifest, and leaving both None forces the live probe. + `disk_label` overrides the on-disk torch/version.py label the wedged-probe path reads. + + Returns (ok, pip_mock). + """ + state = {"version": installed} + + env = {} + if expected_env is not None: + env["UNSLOTH_EXPECTED_TORCH_TAG"] = expected_env + if install_index_url is not None: + env["UNSLOTH_TORCH_INSTALL_INDEX_URL"] = install_index_url + if cvd is not None: + env["CUDA_VISIBLE_DEVICES"] = cvd + if index_family is not None: + env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family + if index_url is not None: + env["UNSLOTH_TORCH_INDEX_URL"] = index_url + + def _run(cmd, *args, **kwargs): + result = MagicMock() + exe = str(cmd[0]) if cmd else "" + if exe == sys.executable: + if probe_timeout: + raise subprocess.TimeoutExpired(cmd, 90) + result.returncode = torch_rc + out = _flavor_probe_stdout(state["version"]) + if (probe_cuda is not None or probe_hip is not None) and state["version"] == installed: + out = _MARK + (f"{state['version']}|{probe_hip or ''}|{probe_cuda or ''}\n") + else: + result.returncode = 0 + if len(cmd) > 1 and str(cmd[1]) == "--query-gpu=compute_cap": + out = "8.6\n" + else: + out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices\n" + result.stdout = out if kwargs.get("text") else out.encode() + return result + + def _pip(*args, **kwargs): + stack_mod._invalidate_torch_runtime_probe() + if repaired is not None: + state["version"] = repaired + + def _rocm(*args, **kwargs): + # The ROCm arm delegates, so the stand-in must move the venv too. + _pip() + + def _which(name, *a, **k): + return "/usr/bin/nvidia-smi" if name == "nvidia-smi" else None + + stack_mod._invalidate_torch_runtime_probe() + + with ( + patch.object(stack_mod, "_TORCH_BACKEND", backend), + patch.object(stack_mod, "NO_TORCH", no_torch), + patch.object(stack_mod, "_RECORDED_TORCH_TAG", recorded), + patch.object(stack_mod.platform, "machine", return_value = "AMD64"), + patch.object(stack_mod, "_is_windows_arm64", return_value = win_arm64), + patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia), + patch.object(stack_mod.shutil, "which", side_effect = _which), + patch.object(stack_mod.os.path, "isfile", return_value = True), + patch.object( + stack_mod, + "_installed_torch_label_on_disk", + side_effect = ( + (lambda: state["version"].lower()) if disk_label is _UNSET else (lambda: disk_label) + ), + ), + patch.object(stack_mod, "pip_install", side_effect = _pip) as mock_pip, + patch.object(stack_mod, "_ensure_rocm_torch", side_effect = _rocm) as mock_rocm, + patch.object(stack_mod.subprocess, "run", side_effect = _run), + patch.dict(stack_mod.os.environ, env, clear = False), + ): + for name, value in ( + ("UNSLOTH_EXPECTED_TORCH_TAG", expected_env), + ("UNSLOTH_TORCH_INSTALL_INDEX_URL", install_index_url), + ("CUDA_VISIBLE_DEVICES", cvd), + ("UNSLOTH_TORCH_INDEX_FAMILY", index_family), + ("UNSLOTH_TORCH_INDEX_URL", index_url), + ): + if value is None: + stack_mod.os.environ.pop(name, None) + ok = _ensure_expected_torch_flavor() + mock_pip.rocm_repair = mock_rocm + return ok, mock_pip + + +class TestExpectedTorchFlavorRepairs: + def test_cpu_torch_under_a_cu124_expectation_is_repaired(self): + ok, mock_pip = _run_flavor_invariant(repaired = "2.10.0+cu124") + assert ok is True + assert mock_pip.call_count == 1 + call_args = [str(a) for a in mock_pip.call_args.args] + assert "--force-reinstall" in call_args + assert "--no-cache-dir" in call_args + assert _index_url(mock_pip).endswith("/cu124") + assert mock_pip.call_args.kwargs["constrain"] is False + + def test_the_repair_uses_install_ps1s_bounded_trio(self): + _ok, mock_pip = _run_flavor_invariant(repaired = "2.10.0+cu124") + call_args = [str(a) for a in mock_pip.call_args.args] + # NOT the <2.12 Linux range: an unbounded trio resolves back to the 2.11 wheel. + for spec in ("torch>=2.4,<2.11.0", "torchvision>=0.19,<0.26.0", "torchaudio>=2.4,<2.11.0"): + assert spec in call_args + + def test_untagged_pypi_wheel_is_repaired(self): + ok, mock_pip = _run_flavor_invariant(installed = "2.11.0", repaired = "2.10.0+cu124") + assert ok is True + assert mock_pip.call_count == 1 + + def test_wrong_cuda_family_is_repaired_to_the_expected_one(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.9.1+cu118", + expected_env = "cu128", + repaired = "2.10.0+cu128", + ) + assert ok is True + assert _index_url(mock_pip).endswith("/cu128") + + def test_a_wedged_probe_classifies_from_disk_and_still_repairs(self): + # A stalled driver hangs `import torch`; version.py on disk still names the wheel. + ok, mock_pip = _run_flavor_invariant( + probe_timeout = True, + repaired = "2.10.0+cu124", + ) + assert ok is True + assert mock_pip.call_count == 1 + + +class TestExpectedTorchFlavorFailsTheUpdate: + def test_a_repair_that_leaves_cpu_torch_fails(self): + ok, mock_pip = _run_flavor_invariant(repaired = None) + assert ok is False + assert mock_pip.call_count == 1 + + def test_the_failure_verdict_reads_torch_version_cuda_not_just_the_tag(self): + # Untagged but carrying a CUDA runtime is a GPU build: reinstall, do not fail. + ok, _mock_pip = _run_flavor_invariant(installed = "2.11.0", repaired = "2.11.0+cu124") + assert ok is True + + def test_a_rocm_build_left_behind_is_not_called_cpu_only(self, capsys): + ok, _mock_pip = _run_flavor_invariant(repaired = "2.9.1+rocm6.4") + assert ok is False + out = capsys.readouterr().out + assert "rocm build but cu124 was expected" in out + assert "CPU-only" not in out + + +class TestExpectedTorchFlavorSkips: + def test_expected_cpu_is_a_no_op(self): + ok, mock_pip = _run_flavor_invariant(expected_env = "cpu") + assert ok is True + mock_pip.assert_not_called() + + @pytest.mark.parametrize("backend", ["cpu", "auto"]) + def test_a_non_gpu_backend_is_a_no_op(self, backend): + """Decided without resolving the expectation: cpu and unrecognised values are + deliberate whatever the handover said, and this exit costs no GPU probe.""" + ok, mock_pip = _run_flavor_invariant(backend = backend) + assert ok is True + mock_pip.assert_not_called() + + @pytest.mark.parametrize("backend", ["rocm", "xpu"]) + def test_a_gpu_backend_that_disagrees_with_the_handover_is_a_no_op(self, backend): + """An explicit ROCm/XPU pin against a cu124 handover. The pin is the newer and + more specific instruction; the handover only describes what the install arm + above it happened to do, so the pin wins and this pass stays out of it.""" + ok, mock_pip = _run_flavor_invariant(backend = backend, expected_env = "cu124") + assert ok is True + mock_pip.assert_not_called() + + def test_backend_cuda_still_runs(self): + ok, mock_pip = _run_flavor_invariant(backend = "cuda", repaired = "2.10.0+cu124") + assert ok is True + assert mock_pip.call_count == 1 + + def test_no_torch_mode_is_a_no_op(self): + ok, mock_pip = _run_flavor_invariant(no_torch = True) + assert ok is True + mock_pip.assert_not_called() + + def test_a_matching_flavor_is_a_no_op(self): + ok, mock_pip = _run_flavor_invariant(installed = "2.9.1+cu124") + assert ok is True + mock_pip.assert_not_called() + + @pytest.mark.parametrize("cvd", ["", "-1", " ", " -1 "]) + def test_hidden_cuda_devices_is_a_no_op_for_an_inferred_expectation(self, cvd): + ok, mock_pip = _run_flavor_invariant(cvd = cvd, expected_env = None) + assert ok is True + mock_pip.assert_not_called() + + @pytest.mark.parametrize("cvd", ["", "-1", " ", " -1 "]) + def test_a_mask_does_not_veto_an_explicit_expectation(self, cvd): + # A reason not to CONCLUDE cu124 from a probe, not to ignore a stated one. + for kwargs in ( + {"expected_env": "cu124"}, + {"expected_env": None, "index_family": "cu124"}, + {"expected_env": None, "recorded": "cu124"}, + ): + ok, mock_pip = _run_flavor_invariant(cvd = cvd, repaired = "2.10.0+cu124", **kwargs) + assert ok is True, kwargs + assert mock_pip.call_count == 1, kwargs + + def test_an_explicit_visible_device_still_repairs(self): + ok, mock_pip = _run_flavor_invariant(cvd = "0", repaired = "2.10.0+cu124") + assert ok is True + assert mock_pip.call_count == 1 + + @pytest.mark.parametrize("tag", ["current", "custom", "simple", "cu"]) + def test_an_unenforceable_expectation_is_a_no_op(self, tag): + # "xpu" and "rocm" are NOT here: both are published and both are enforced. + ok, mock_pip = _run_flavor_invariant(expected_env = tag) + assert ok is True + mock_pip.assert_not_called() + + def test_an_empty_handover_tag_falls_through_rather_than_deciding(self): + # PowerShell 7.5+ keeps an entry assigned "", which must read as "nobody said". + ok, mock_pip = _run_flavor_invariant( + expected_env = "", + recorded = "cu128", + repaired = "2.10.0+cu128", + ) + assert ok is True + assert _index_url(mock_pip).endswith("/cu128") + + def test_missing_or_unimportable_torch_is_a_no_op(self): + # Reinstalling over an unimportable torch turns a driver fault into a wheel one. + ok, mock_pip = _run_flavor_invariant(torch_rc = 1) + assert ok is True + mock_pip.assert_not_called() + + def test_an_unreadable_venv_is_a_no_op(self): + ok, mock_pip = _run_flavor_invariant(probe_timeout = True, disk_label = "") + assert ok is True + mock_pip.assert_not_called() + + +class TestExpectedTorchFlavorResolution: + def test_the_manifest_answers_when_the_environment_is_silent(self): + ok, mock_pip = _run_flavor_invariant( + expected_env = None, + recorded = "cu128", + repaired = "2.10.0+cu128", + ) + assert ok is True + assert _index_url(mock_pip).endswith("/cu128") + + def test_the_environment_wins_over_the_manifest(self): + _ok, mock_pip = _run_flavor_invariant( + expected_env = "cu124", + recorded = "cu128", + repaired = "2.10.0+cu124", + ) + assert _index_url(mock_pip).endswith("/cu124") + + def test_the_live_probe_answers_when_nothing_recorded_it(self): + _ok, mock_pip = _run_flavor_invariant( + expected_env = None, + recorded = None, + cuda_version = "12.8", + repaired = "2.10.0+cu128", + ) + assert _index_url(mock_pip).endswith("/cu128") + + def test_the_live_probe_declines_without_an_nvidia_gpu(self): + # Inventing an expectation here reinstalls CUDA torch onto a CPU-only box. + ok, mock_pip = _run_flavor_invariant(expected_env = None, recorded = None, nvidia = False) + assert ok is True + mock_pip.assert_not_called() + + def test_a_cpu_pin_beats_the_live_probe(self): + ok, mock_pip = _run_flavor_invariant(expected_env = None, recorded = None, index_family = "cpu") + assert ok is True + mock_pip.assert_not_called() + + def test_the_setup_scripts_index_url_is_used_when_its_leaf_matches(self): + # Credentials are not reconstructible from a family leaf. + _ok, mock_pip = _run_flavor_invariant( + install_index_url = "https://mirror.local/whl/cu124?token=secret", + repaired = "2.10.0+cu124", + ) + assert _index_url(mock_pip) == "https://mirror.local/whl/cu124?token=secret" + + def test_an_index_url_naming_another_family_is_ignored(self): + # setup.ps1 hands over /cpu alongside a "rocm" tag on the AMD Windows path. + _ok, mock_pip = _run_flavor_invariant( + install_index_url = "https://download.pytorch.org/whl/cpu", + repaired = "2.10.0+cu124", + ) + assert _index_url(mock_pip) == f"{stack_mod._PYTORCH_WHL_BASE}/cu124" + + def test_a_matching_family_pin_supplies_the_index(self): + _ok, mock_pip = _run_flavor_invariant( + index_url = "https://mirror.local/whl/cu124", + repaired = "2.10.0+cu124", + ) + assert _index_url(mock_pip) == "https://mirror.local/whl/cu124" + + +class TestTorchFlavorTagVocabulary: + """_torch_flavor_tag is compared against tags install.ps1 and setup.ps1 produced, so + it has to classify identically.""" + + @pytest.mark.parametrize( + "version,tag", + [ + ("2.9.1+cu124", "cu124"), + ("2.11.0+cu130", "cu130"), + ("2.9.1+rocm6.4", "rocm"), + ("2.11.0+rocm7.2.1", "rocm"), + ("2.9.1+xpu", "xpu"), + ("2.11.0+cpu", "cpu"), + ("2.11.0", "cpu"), + ("2.9.1+CU124", "cu124"), + ("", ""), + ], + ) + def test_tags(self, version, tag): + assert stack_mod._torch_flavor_tag(version) == tag + + +class TestUnknownFamilyPinIsNotOverridden: + """An explicit index pin whose leaf names no flavor is applied verbatim at install + time, so this pass has no standing to second-guess it. The only expectation it could + act on comes from the manifest, i.e. from whatever was installed BEFORE the pin was + set, and repairing off that stale tag reinstalls from the public pytorch index -- + overriding a deliberate package source, and failing outright on an air-gapped host. + _ensure_cuda_torch already declines on the same test.""" + + def test_a_simple_mirror_pin_suppresses_the_manifest_fallback(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = None, + recorded = "cu124", + index_url = "https://mirror.corp.example/simple", + ) + assert ok is True + mock_pip.assert_not_called() + + def test_a_readable_cuda_pin_still_repairs(self): + """Narrowness: the escape must not disable a normal cu* mirror.""" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.6.0+cu124", + expected_env = None, + recorded = "cu124", + index_url = "https://mirror.corp.example/whl/cu124", + ) + assert ok is True + args = mock_pip.call_args[0] + assert "https://mirror.corp.example/whl/cu124" in args + + +class TestExpectedXpuFlavorIsEnforced: + """setup.ps1 publishes "xpu" for an Arc host and installs the XPU trio before handing + over, so declining to act on that expectation would leave the invariant carrying an + answer it refuses to use. The exposure is identical to the CUDA one: the dependency + steps re-resolve torch from PyPI, and _ensure_xpu_torch cannot clean up afterwards + because step 13's whole repair set is gated off Windows.""" + + def test_an_xpu_venv_that_lost_torch_to_pypi_is_repaired(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", repaired = "2.9.1+xpu", expected_env = "xpu" + ) + assert ok is True + args = mock_pip.call_args[0] + # XPU floor is 2.6, not the CUDA trio's 2.4: unsloth raises at import below it. + assert "torch>=2.6,<2.11.0" in args + assert "torchvision>=0.21,<0.26.0" in args + assert any("xpu" in str(a) for a in args) + + def test_a_healthy_xpu_venv_is_untouched(self): + ok, mock_pip = _run_flavor_invariant(installed = "2.9.1+xpu", expected_env = "xpu") + assert ok is True + mock_pip.assert_not_called() + + def test_an_xpu_repair_that_does_not_take_fails_the_update(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", repaired = None, expected_env = "xpu" + ) + assert ok is False + mock_pip.assert_called_once() + + def test_a_cuda_mask_does_not_cancel_an_xpu_repair(self): + """CUDA_VISIBLE_DEVICES hides an NVIDIA GPU, not an Arc one.""" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.9.1+xpu", + expected_env = "xpu", + cvd = "", + ) + assert ok is True + mock_pip.assert_called_once() + + def test_rocm_is_delegated_rather_than_rebuilt_here(self): + """AMD's Windows wheels live on a per-architecture repo.amd.com index that a + generic "rocm" tag cannot name, and setup.ps1 hands over an index URL that still + points at /cpu on that path. _ensure_rocm_torch already detects the arch, maps + it, and honours an explicit pin, so the repair is delegated to it rather than + rebuilt from a guessed URL -- but it IS repaired, because it runs at step 2b, + before the dependency steps that can put PyPI's CPU wheel here.""" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.11.0+rocm7.2", + expected_env = "rocm", + ) + mock_pip.rocm_repair.assert_called_once() + mock_pip.assert_not_called() + assert ok is True + + def test_a_healthy_rocm_venv_does_not_call_the_repair(self): + ok, mock_pip = _run_flavor_invariant(installed = "2.11.0+rocm7.2", expected_env = "rocm") + mock_pip.rocm_repair.assert_not_called() + mock_pip.assert_not_called() + assert ok is True + + def test_a_rocm_repair_that_does_not_take_fails_the_update(self): + """The state that used to be written as a successful CPU-only manifest.""" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", repaired = None, expected_env = "rocm" + ) + mock_pip.rocm_repair.assert_called_once() + mock_pip.assert_not_called() + assert ok is False + + +class TestTheDelegatedRocmRepairIsVerifiedByFamily: + """_torch_build_is_gpu is family-blind, so it cannot judge a ROCm repair. + + A transient repo.amd.com failure is non-fatal inside _ensure_rocm_torch, and the + cu124 wheel it leaves behind passes that check, so the update exited 0 and the + manifest recorded "rocm" over an environment that never received it. + """ + + def test_a_repair_that_left_a_cuda_wheel_now_fails(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0+cu124", + repaired = "2.6.0+cu124", + expected_env = "rocm", + backend = "rocm", + ) + assert ok is False + mock_pip.rocm_repair.assert_called_once() + + def test_a_repair_that_took_still_passes(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0+cu124", + repaired = "2.11.0+rocm7.2", + expected_env = "rocm", + backend = "rocm", + ) + assert ok is True + mock_pip.rocm_repair.assert_called_once() + + def test_a_repair_that_left_cpu_torch_still_reads_as_cpu(self): + ok, _mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.11.0+cpu", + expected_env = "rocm", + backend = "rocm", + ) + assert ok is False + + def test_an_unreadable_venv_does_not_fail_the_update_on_its_own(self): + ok, _mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.11.0+cpu", + expected_env = "rocm", + backend = "rocm", + probe_timeout = True, + disk_label = "", + ) + assert ok is True + + +class TestAPinnedCpuIndexIsEnforcedToo: + """UV_TORCH_BACKEND is honoured by the unpinned dependency steps. + + A venv deliberately built against /cpu can therefore come out of them holding a GPU + wheel, and the update would record expected_torch_tag: cpu over it. Only a PIN + counts: setup.ps1 also publishes "cpu" for a host whose nvidia-smi probe returned + nothing, and acting on that would push a healthy cu124 venv down to CPU. + """ + + def test_a_gpu_wheel_under_a_pinned_cpu_index_is_repaired(self): + pin = "https://download.pytorch.org/whl/cpu" + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0+cu124", + repaired = "2.11.0+cpu", + expected_env = "cpu", + index_url = pin, + backend = "cpu", + ) + assert ok is True + assert _index_url(mock_pip) == pin + assert "--force-reinstall" in [str(a) for a in mock_pip.call_args.args] + + def test_a_published_cpu_tag_with_no_pin_is_still_left_alone(self): + # setup.ps1 publishes "cpu" when nvidia-smi answers nothing; the healthy cu124 + # venv underneath must not be downgraded. + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0+cu124", + expected_env = "cpu", + nvidia = False, + ) + assert ok is True + mock_pip.assert_not_called() + + def test_a_cpu_venv_under_a_cpu_pin_is_a_no_op(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = "cpu", + index_url = "https://download.pytorch.org/whl/cpu", + backend = "cpu", + ) + assert ok is True + mock_pip.assert_not_called() + + def test_a_cpu_repair_that_did_not_take_fails(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0+cu124", + repaired = "2.6.0+cu124", + expected_env = "cpu", + backend = "cpu", + index_url = "https://download.pytorch.org/whl/cpu", + ) + assert ok is False + assert mock_pip.call_count == 1 + + def test_a_cpu_backend_under_a_gpu_expectation_is_still_left_alone(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = "cu124", + backend = "cpu", + ) + assert ok is True + mock_pip.assert_not_called() + + +class TestWindowsOnArmKeepsTheNoTorchaudioException: + """No win_arm64 torchaudio wheel is published on any index. + + setup.ps1 drops it from all four of its install trios ($WinArm64NoAudio), so asking + for it here would turn a repairable venv into a failed install, and the venv the + repair rebuilds could not have been installed in the first place. + """ + + def test_torchaudio_is_dropped_on_windows_arm64(self): + ok, mock_pip = _run_flavor_invariant(repaired = "2.10.0+cu124", win_arm64 = True) + assert ok is True + args = [str(a) for a in mock_pip.call_args.args] + assert not any(a.startswith("torchaudio") for a in args) + assert any(a.startswith("torch>=") for a in args) + assert any(a.startswith("torchvision") for a in args) + assert _index_url(mock_pip).endswith("/cu124") + + def test_the_xpu_trio_drops_it_too(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.9.1+xpu", + expected_env = "xpu", + win_arm64 = True, + ) + assert ok is True + args = [str(a) for a in mock_pip.call_args.args] + assert not any(a.startswith("torchaudio") for a in args) + assert "torch>=2.6,<2.11.0" in args + + def test_x64_windows_still_installs_all_three(self): + _ok, mock_pip = _run_flavor_invariant(repaired = "2.10.0+cu124") + args = [str(a) for a in mock_pip.call_args.args] + assert any(a.startswith("torchaudio") for a in args) + + +class TestTheRequestedFamilyIsVerifiedAfterEveryRepair: + """A GPU build is not the same answer as THE GPU build that was asked for. + + A misconfigured mirror can answer a /cu128 request with a cached cu124, rocm or xpu + wheel. _torch_build_is_gpu is deliberately family-blind, so the update exited 0 and + the manifest recorded the requested tag over a build that never arrived. + """ + + @pytest.mark.parametrize("landed", ["2.6.0+cu124", "2.11.0+rocm7.2", "2.9.1+xpu"]) + def test_a_wheel_from_the_wrong_family_fails(self, landed): + ok, _mock_pip = _run_flavor_invariant( + expected_env = "cu128", + repaired = landed, + ) + assert ok is False + + def test_the_requested_family_passes(self): + ok, _mock_pip = _run_flavor_invariant( + expected_env = "cu128", + repaired = "2.6.0+cu128", + ) + assert ok is True + + def test_an_xpu_repair_is_held_to_the_same_rule(self): + ok, _mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = "xpu", + repaired = "2.6.0+cu124", + ) + assert ok is False + + def test_an_unreadable_venv_still_passes(self): + ok, _mock_pip = _run_flavor_invariant( + expected_env = "cu128", + repaired = "2.6.0+cu128", + probe_timeout = True, + disk_label = "", + ) + assert ok is True + + def test_a_still_cpu_venv_keeps_its_own_warning(self, capsys): + ok, _mock_pip = _run_flavor_invariant(expected_env = "cu124", repaired = None) + assert ok is False + assert "CPU-only" in capsys.readouterr().out + + +class TestAnUntaggedGpuWheelIsNotACpuMatch: + """_torch_flavor_tag reads every untagged version as "cpu". + + Right for PyPI, wrong for a private index serving an untagged CUDA or ROCm build: + under a /cpu pin that wheel compared equal to the expectation, skipped the repair, + and was recorded as cpu. The runtime probe already carries the markers that tell + them apart. + """ + + def test_an_untagged_cuda_wheel_under_a_cpu_pin_is_repaired(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0", + repaired = "2.11.0+cpu", + expected_env = "cpu", + backend = "cpu", + index_url = "https://download.pytorch.org/whl/cpu", + probe_cuda = "12.4", + ) + assert ok is True + assert mock_pip.call_count == 1 + + def test_a_genuinely_untagged_cpu_wheel_is_still_a_match(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.6.0", + expected_env = "cpu", + backend = "cpu", + index_url = "https://download.pytorch.org/whl/cpu", + ) + assert ok is True + mock_pip.assert_not_called() + + +class TestAnExplicitPinOutranksTheManifest: + """Direct `python install_python_stack.py` on Windows, which the invariant supports. + + The manifest records what a PREVIOUS run installed; a pin is the instruction for + THIS one. Resolving the manifest first let a freshly set cu128 pin lose to a stale + cu124 record, and _expected_torch_index_url then rejected the cu128 pin as a family + mismatch and repaired from the PUBLIC cu124 index -- undoing both the family and the + source the user had just chosen. + """ + + def test_a_new_family_pin_beats_a_stale_manifest(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.10.0+cu124", + repaired = "2.10.0+cu128", + expected_env = None, + recorded = "cu124", + index_family = "cu128", + ) + assert ok is True + assert "cu128" in _index_url(mock_pip) + assert "cu124" not in _index_url(mock_pip) + assert "--force-reinstall" in [str(a) for a in mock_pip.call_args.args] + + def test_a_pinned_url_beats_a_stale_manifest_and_is_the_repair_source(self): + pin = "https://mirror.corp.example/whl/cu128" + ok, mock_pip = _run_flavor_invariant( + installed = "2.10.0+cu124", + repaired = "2.10.0+cu128", + expected_env = None, + recorded = "cu124", + index_url = pin, + ) + assert ok is True + assert _index_url(mock_pip) == pin + + def test_a_rocm_pin_collapses_to_the_flavor_vocabulary(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+rocm7.2", + expected_env = None, + recorded = "cu124", + backend = "rocm", + index_family = "gfx1151", + ) + assert ok is True + mock_pip.assert_not_called() + + def test_a_cpu_pin_beats_a_gpu_manifest(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = None, + recorded = "cu124", + index_family = "cpu", + ) + assert ok is True + mock_pip.assert_not_called() + + def test_an_unrecognised_pin_still_falls_through_to_the_manifest(self): + # A /simple mirror names no family; the caller's unknown-pin gate is the guard. + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = None, + recorded = "cu124", + index_url = "https://mirror.corp.example/simple", + ) + assert ok is True + mock_pip.assert_not_called() + + def test_the_setup_handover_still_wins_over_a_pin(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.10.0+cu124", + repaired = "2.10.0+cu126", + expected_env = "cu126", + index_family = "cu128", + ) + assert ok is True + assert "cu126" in _index_url(mock_pip) + + +class TestExplicitlyPinnedGpuFlavorsAreStillEnforced: + """An explicit GPU pin sets _TORCH_BACKEND at import, and an XPU pin additionally + reads as an "unknown family" to the shared helper, whose known set predates XPU. + Between them those two gates skipped the invariant on exactly the hosts that asked + for that GPU family on purpose -- so a later dependency install could put PyPI's CPU + wheel there and the update would still report success.""" + + def test_a_pinned_xpu_backend_is_enforced_when_it_agrees(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.9.1+xpu", + backend = "xpu", + expected_env = "xpu", + ) + assert ok is True + mock_pip.assert_called_once() + + def test_a_pinned_rocm_backend_is_enforced_when_it_agrees(self): + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.11.0+rocm7.2", + backend = "rocm", + expected_env = "rocm", + ) + assert ok is True + mock_pip.rocm_repair.assert_called_once() + + def test_an_xpu_index_pin_repairs_from_that_pin(self): + """The pin's leaf IS the expected family, so it is not an unknown pin at all -- + and it is the right index to repair from.""" + pin = "https://mirror.corp.example/whl/xpu" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + repaired = "2.9.1+xpu", + backend = "xpu", + expected_env = "xpu", + index_url = pin, + ) + assert ok is True + assert pin in mock_pip.call_args[0] + + def test_an_unknown_pin_is_still_left_alone(self): + """Regression guard: narrowing the veto must not reopen the case it closed.""" + ok, mock_pip = _run_flavor_invariant( + installed = "2.11.0+cpu", + expected_env = None, + recorded = "cu124", + index_url = "https://mirror.corp.example/simple", + ) + assert ok is True + mock_pip.assert_not_called() + + +class TestSetupPs1CudaOnDiskFallback: + """A wedged NVIDIA driver hangs or faults `import torch` exactly as a faulted HIP + runtime does. XPU and ROCm both rescue that venv from torch/version.py on disk; CUDA + had no such arm, so the probe-failure chain fell through with a NULL tag, the no-wipe + escape could not see a cu* wheel to preserve, and a direct update deleted a healthy + CUDA environment before aborting.""" + + def test_the_classifier_exists_and_keeps_the_family(self): + assert "function Get-VenvTorchCudaTag" in _SETUP_SRC + assert "function Test-VenvTorchIsCuda" in _SETUP_SRC + block = _SETUP_SRC[_SETUP_SRC.index("function Get-VenvTorchCudaTag") :][:1400] + # The family, not a flat "cuda": the stale comparison below it is cu126-vs-cu128. + assert "cu[0-9]+" in block + assert "site-packages\\torch\\version.py" in block + + def test_it_joins_the_probe_failure_chain(self): + chain_start = _SETUP_SRC.index("elseif (Test-VenvTorchIsXpu -VenvPath $VenvDir)") + chain = _SETUP_SRC[chain_start : _SETUP_SRC.index("$shouldRebuild = $true", chain_start)] + assert "Test-VenvTorchIsCuda -VenvPath $VenvDir" in chain + assert "$installedTorchTag = Get-VenvTorchCudaTag" in chain + + def test_it_does_not_disturb_the_callers_match_state(self): + """A bare -match would clobber $Matches in the caller's scope; the stale-venv + block reads $Matches[1] a few lines above.""" + block = _SETUP_SRC[_SETUP_SRC.index("function Get-VenvTorchCudaTag") :][:1400] + assert "[regex]::Match(" in block + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-q"])) + + +class TestThePackagesTiedToTheTorchReleaseAreResettled: + """A repair that MOVES the torch release invalidates two compiled extensions. + + torchao's cpp extensions are torch-release-specific, and step 4 chose its pin from + the torch this repair then replaced: 0.17.0 selected for a 2.11 that the <2.11.0 + repair spec takes down to 2.10, whose matched build is 0.16.0. Leaving the wrong one + installed silently drops to the slow fallback. + + xFormers is stricter: its _C.pyd is linked against one exact (torch, CUDA) pair, and + beside any other pair torch.ops.load_library raises, which xformers/_cpp_lib.py + downgrades to a log line, so the import "succeeds" with memory-efficient attention, + SwiGLU and the sparse ops silently gone. This script never installs xFormers, so + removal is the only correct action, and it is what the backend's own resolver + already concludes: torch SDPA beats an extension that cannot load. + """ + + def _resync( + self, + before, + after, + *, + resident_xformers = None, + installed_spec = False, + ): + calls = {"torchao": [], "removed": []} + + with ( + patch.object(stack_mod, "_probe_installed_torch_version", return_value = after), + patch.object( + stack_mod, + "_exact_distribution_spec_is_installed", + return_value = installed_spec, + ), + patch.object( + stack_mod, + "_resident_xformers_build_torch", + return_value = resident_xformers, + ), + patch.object( + stack_mod, + "pip_install_try", + side_effect = lambda *a, **k: calls["torchao"].append([str(x) for x in a]), + ), + patch.object( + stack_mod, + "_uninstall_distribution", + side_effect = lambda name: calls["removed"].append(name) or True, + ), + patch.object(stack_mod, "_note", lambda *a, **k: None), + ): + calls["ok"] = stack_mod._resync_torch_coupled_packages(before) + return calls + + def test_a_release_that_moved_re_pins_torchao(self): + calls = self._resync("2.11.0+cu124", "2.10.0+cu124") + assert calls["torchao"], "torchao must be re-selected for the new release" + assert any("torchao==0.16.0" in " ".join(c) for c in calls["torchao"]) + + def test_a_build_that_did_not_move_at_all_does_nothing(self): + calls = self._resync("2.10.0+cu124", "2.10.0+cu124", resident_xformers = "2.11.0+cu128") + assert calls["torchao"] == [] + assert calls["removed"] == [] + assert calls["ok"] is True + + def test_a_cuda_major_change_alone_re_pins_torchao(self): + # _select_torchao_spec branches on cuda>=13, so cu124 -> cu130 moves the build. + calls = self._resync("2.10.0+cu124", "2.10.0+cu130") + assert calls["torchao"], "the CUDA major moved, so the pin has to be re-selected" + assert any("torchao==0.17.0" in " ".join(c) for c in calls["torchao"]) + + def test_a_flavor_change_within_one_cuda_major_leaves_torchao_alone(self): + calls = self._resync("2.10.0+cu124", "2.10.0+cu128") + assert calls["torchao"] == [] + + def test_a_flavor_change_alone_still_rechecks_xformers(self): + # xFormers links against the exact (torch, CUDA) pair. + calls = self._resync("2.10.0+cu124", "2.10.0+cu128", resident_xformers = "2.10.0+cu124") + assert calls["torchao"] == [], "the release did not move; torchao is fine" + assert calls["removed"] == ["xformers"] + assert calls["ok"] is True, "nothing here touched torch" + + def test_the_torchao_reinstall_cannot_drag_torch_back(self): + # torchao depends on torch: resolving deps re-pulls the wheel just removed. + calls = self._resync("2.11.0+cu124", "2.10.0+cu124") + assert calls["torchao"], "the release moved, so torchao is re-pinned" + assert all("--no-deps" in c for c in calls["torchao"]) + + def test_a_pass_that_ran_a_torch_touching_install_asks_to_be_re_verified(self): + calls = self._resync("2.11.0+cu124", "2.10.0+cu124") + assert calls["ok"] is False + + def test_a_pass_that_installed_nothing_does_not(self): + calls = self._resync("2.11.0+cu124", "2.10.0+cu124", installed_spec = True) + assert calls["torchao"] == [] + assert calls["ok"] is True + + def test_a_torchao_already_matching_is_left_alone(self): + calls = self._resync("2.11.0+cu124", "2.10.0+cu124", installed_spec = True) + assert calls["torchao"] == [] + + def test_a_mismatched_xformers_is_removed(self): + calls = self._resync("2.11.0", "2.10.0+cu124", resident_xformers = "2.11.0+cu128") + assert calls["removed"] == ["xformers"] + + def test_an_xformers_built_for_the_resident_torch_is_kept(self): + calls = self._resync("2.11.0", "2.10.0+cu124", resident_xformers = "2.10.0+cu124") + assert calls["removed"] == [] + + def test_an_absent_xformers_is_not_touched(self): + calls = self._resync("2.11.0+cu124", "2.10.0+cu124", resident_xformers = None) + assert calls["removed"] == [] + + def test_an_unreadable_torch_after_the_repair_does_nothing(self): + calls = self._resync("2.11.0+cu124", None, resident_xformers = "2.11.0+cu128") + assert calls["torchao"] == [] + assert calls["removed"] == [] + assert calls["ok"] is True + + def test_neither_half_can_fail_the_update(self, capsys): + with ( + patch.object(stack_mod, "_probe_installed_torch_version", return_value = "2.10.0+cu124"), + patch.object( + stack_mod, + "_select_torchao_spec", + side_effect = RuntimeError("index down"), + ), + patch.object( + stack_mod, + "_resident_xformers_build_torch", + side_effect = RuntimeError("unreadable"), + ), + ): + assert stack_mod._resync_torch_coupled_packages("2.11.0+cu124") is True + out = capsys.readouterr().out + assert "could not re-match torchao" in out + assert "could not re-check xFormers" in out + + +class TestTheResidentXformersBuildIsReadFromDisk: + def test_the_recorded_torch_is_returned(self, tmp_path): + pkg = tmp_path / "xformers" + pkg.mkdir() + (pkg / "cpp_lib.json").write_text( + '{"version": {"torch": "2.10.0+cu128"}}', encoding = "utf-8" + ) + with patch.object( + stack_mod.importlib.util, + "find_spec", + return_value = SimpleNamespace(submodule_search_locations = [str(pkg)]), + ): + assert stack_mod._resident_xformers_build_torch() == "2.10.0+cu128" + + @pytest.mark.parametrize( + "body", ['{"version": {}}', "{not json", '{"version": {"torch": " "}}'] + ) + def test_an_unusable_record_reads_as_unknown(self, tmp_path, body): + pkg = tmp_path / "xformers" + pkg.mkdir() + (pkg / "cpp_lib.json").write_text(body, encoding = "utf-8") + with patch.object( + stack_mod.importlib.util, + "find_spec", + return_value = SimpleNamespace(submodule_search_locations = [str(pkg)]), + ): + assert stack_mod._resident_xformers_build_torch() is None + + def test_an_absent_xformers_reads_as_unknown(self): + with patch.object(stack_mod.importlib.util, "find_spec", return_value = None): + assert stack_mod._resident_xformers_build_torch() is None + + def test_a_find_spec_that_raises_reads_as_unknown(self): + with patch.object(stack_mod.importlib.util, "find_spec", side_effect = ValueError("boom")): + assert stack_mod._resident_xformers_build_torch() is None + + +class TestTheResyncNoticesItsOwnFailures: + """Both halves report failure by return value, not by raising. + + Ignoring that let the update write a completion manifest over an incompatible + torchao, or over an xFormers whose removal was blocked, with nothing said. + """ + + def _resync_with( + self, + *, + torchao_ok = True, + uninstall_ok = True, + ): + with ( + patch.object(stack_mod, "_probe_installed_torch_version", return_value = "2.10.0+cu124"), + patch.object(stack_mod, "_exact_distribution_spec_is_installed", return_value = False), + patch.object(stack_mod, "_resident_xformers_build_torch", return_value = "2.11.0+cu128"), + patch.object(stack_mod, "pip_install_try", return_value = torchao_ok), + patch.object(stack_mod, "_uninstall_distribution", return_value = uninstall_ok), + patch.object(stack_mod, "_note", lambda *a, **k: None), + ): + return stack_mod._resync_torch_coupled_packages("2.11.0+cu124") + + def test_a_torchao_install_that_did_not_take_is_reported(self, capsys): + self._resync_with(torchao_ok = False) + assert "could not install" in capsys.readouterr().out + + def test_an_xformers_removal_that_was_blocked_is_reported(self, capsys): + self._resync_with(uninstall_ok = False) + out = capsys.readouterr().out + assert "could not remove the mismatched xFormers" in out + + def test_a_clean_pass_says_nothing(self, capsys): + self._resync_with() + out = capsys.readouterr().out + assert "[WARN]" not in out + + def test_a_failure_still_does_not_fail_the_update(self): + # The return value reports whether torch may have MOVED, not that all was well. + assert self._resync_with(torchao_ok = False, uninstall_ok = False) is False + + +class TestThePostRepairCheckUsesTheSameRuleAsThePreRepairOne: + def test_an_untagged_gpu_wheel_does_not_satisfy_a_cpu_expectation(self): + # Untagged reads as cpu, so the post-repair check would accept the very wheel + # the pre-repair check rejected. + with ( + patch.object( + stack_mod, + "_probe_torch_runtime", + return_value = (True, True, "2.6.0", "", "12.4"), + ), + ): + assert stack_mod._installed_flavor_tag_now("cpu") == "cuda" + with patch.object( + stack_mod, + "_probe_torch_runtime", + return_value = (True, True, "2.6.0", "6.4", ""), + ): + assert stack_mod._installed_flavor_tag_now("cpu") == "rocm" + + def test_a_genuine_untagged_cpu_wheel_still_reads_as_cpu(self): + with patch.object( + stack_mod, + "_probe_torch_runtime", + return_value = (True, True, "2.6.0", "", ""), + ): + assert stack_mod._installed_flavor_tag_now("cpu") == "cpu" + + def test_the_adjustment_is_scoped_to_a_cpu_expectation(self): + with patch.object( + stack_mod, + "_probe_torch_runtime", + return_value = (True, True, "2.6.0", "", "12.4"), + ): + assert stack_mod._installed_flavor_tag_now("cu124") == "cpu" + assert stack_mod._installed_flavor_tag_now() == "cpu" + + def test_an_untagged_xpu_wheel_does_not_satisfy_a_cpu_expectation(self): + """torch.version.xpu is where an untagged source, conda or private-index XPU + build carries its runtime -- .hip and .cuda are both empty there. Reading only + those two accepted the XPU wheel under a /cpu pin, returned success without + replacing it, and then recorded a PINNED cpu flavor for a venv still holding it. + """ + with ( + patch.object( + stack_mod, + "_probe_torch_runtime", + return_value = (True, True, "2.9.0", "", ""), + ), + patch.object(stack_mod, "_TORCH_RUNTIME_XPU", "20250101"), + ): + assert stack_mod._installed_flavor_tag_now("cpu") == "xpu" + + def test_the_gpu_family_reading_prefers_the_explicit_runtimes(self): + # An XPU marker beside a CUDA or HIP one names the accelerator that wheel was + # BUILT for; xpu is the answer only when it is the sole marker. + assert stack_mod._gpu_family_from_runtime_markers("6.4", "12.4") == "rocm" + assert stack_mod._gpu_family_from_runtime_markers("", "12.4") == "cuda" + assert stack_mod._gpu_family_from_runtime_markers("", "") == "xpu" + + +class TestAFailedGpuPinIsNotADeliberateCpuChoice: + """setup.ps1 falls back to the CPU index when a pinned ROCm or XPU install fails and + publishes the resolved cpu tag, while the original GPU pin is still in the + environment. Recording that as pinned makes _expected_cpu_flavor_was_chosen() read a + failed install as an intentional one and suppress the repair guidance for good.""" + + @staticmethod + def _pinned( + monkeypatch, + flavor, + *, + url = "", + family = "", + backend = "", + recorded = None, + ): + for var in ( + "UNSLOTH_TORCH_INDEX_URL", + "UNSLOTH_TORCH_INDEX_FAMILY", + "UNSLOTH_TORCH_INSTALL_INDEX_URL", + ): + monkeypatch.delenv(var, raising = False) + if url: + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", url) + if family: + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_FAMILY", family) + monkeypatch.delenv("UNSLOTH_TORCH_BACKEND_SOURCE", raising = False) + monkeypatch.setattr(stack_mod, "_TORCH_BACKEND", backend) + monkeypatch.setattr(stack_mod, "_RECORDED_TORCH_TAG", (recorded or ("", False))[0]) + monkeypatch.setattr(stack_mod, "_RECORDED_TORCH_TAG_PINNED", (recorded or ("", False))[1]) + return stack_mod._expected_torch_flavor_was_pinned(flavor) + + def test_a_rocm_pin_that_settled_on_cpu_is_not_a_cpu_choice(self, monkeypatch): + assert ( + self._pinned(monkeypatch, "cpu", url = "https://download.pytorch.org/whl/rocm6.4") + is False + ) + + def test_an_xpu_family_that_settled_on_cpu_is_not_a_cpu_choice(self, monkeypatch): + assert self._pinned(monkeypatch, "cpu", family = "xpu") is False + + def test_a_gpu_backend_that_settled_on_cpu_is_not_a_cpu_choice(self, monkeypatch): + assert self._pinned(monkeypatch, "cpu", backend = "rocm") is False + + def test_a_carried_forward_gpu_record_does_not_pin_a_cpu_fallback(self, monkeypatch): + assert self._pinned(monkeypatch, "cpu", recorded = ("cu124", True)) is False + + def test_a_cpu_pin_that_settled_on_cpu_still_counts(self, monkeypatch): + assert self._pinned(monkeypatch, "cpu", url = "https://download.pytorch.org/whl/cpu") is True + + def test_a_rocm_pin_that_settled_on_rocm_still_counts(self, monkeypatch): + assert ( + self._pinned(monkeypatch, "rocm", url = "https://download.pytorch.org/whl/rocm6.4") + is True + ) + + def test_a_cuda_pin_counts_for_any_cuda_flavor(self, monkeypatch): + # cu124 and cu128 are both the cuda family: the pin names the family, and the + # recorded tag names the exact index within it. + assert ( + self._pinned(monkeypatch, "cu128", url = "https://download.pytorch.org/whl/cu124") is True + ) + + def test_a_carried_forward_cpu_record_still_counts_for_cpu(self, monkeypatch): + assert self._pinned(monkeypatch, "cpu", recorded = ("cpu", True)) is True + + def test_asking_without_a_flavor_answers_as_it_always_did(self, monkeypatch): + assert self._pinned(monkeypatch, "", url = "https://download.pytorch.org/whl/rocm6.4") is True + + +class TestTheWindowsXpuTritonSwapReachesADirectRun: + """setup.ps1 performs the swap after this script exits; a direct run has no such + postlude, and the core install pulls triton-windows over torch's XPU triton.""" + + def test_the_handover_variable_gates_it(self): + source = inspect.getsource(stack_mod._ensure_xpu_triton) + assert "if NO_TORCH or IS_MACOS:" in source, "Windows must no longer be excluded outright" + assert ( + 'IS_WINDOWS and os.environ.get("UNSLOTH_EXPECTED_TORCH_TAG"' in source + ), "under setup.ps1 the swap is still that script's job" + + def test_the_windows_branch_runs_it_after_the_invariant(self): + source = inspect.getsource(stack_mod.install_python_stack) + marker = "if IS_WINDOWS and not NO_TORCH:" + assert marker in source + block = source[source.index(marker) :][:1200] + assert "_ensure_expected_torch_flavor" in block + assert "_ensure_xpu_triton()" in block + # After, for the reason step 13 puts it last: the swap keys off the +xpu label. + assert block.index("_ensure_expected_torch_flavor") < block.index("_ensure_xpu_triton()") + + +class TestTheDelegatedRocmRepairKeepsTheArm64Exception: + """_ensure_rocm_torch's Windows branch asked for the full trio unconditionally. + + No win_arm64 torchaudio wheel exists, so the whole trio is unresolvable there. On + the delegated path that failure is nonfatal, so the CPU build the repair was meant + to replace stays put and the family verification then fails the update -- a worse + outcome than the plain trio case, where the failure is at least immediate. + """ + + def test_the_windows_rocm_install_drops_torchaudio_on_arm64(self): + source = inspect.getsource(stack_mod._ensure_rocm_torch) + block = source[source.index("_WINDOWS_ROCM_TORCH_PKG_SPECS.get") :][:1200] + assert ( + "_is_windows_arm64()" in block + ), "the delegated ROCm repair needs the same exception as the flavor repair" + assert "*_rocm_trio" in block, "the trio has to be built, not passed positionally" + + def test_x64_windows_still_asks_for_all_three(self): + source = inspect.getsource(stack_mod._ensure_rocm_torch) + block = source[source.index("_WINDOWS_ROCM_TORCH_PKG_SPECS.get") :][:1200] + assert "_rocm_trio = [_torch_pkg, _vision_pkg, _audio_pkg]" in block + + +class TestADefinitiveImportFailureIsNotADriverHang: + """setup.ps1's disk-label rescue treated both the same. + + A wedged driver and a truncated torch both leave a +cu* version.py behind. Keeping + the venv is right in both cases -- deleting it does not fix a driver, which is the + whole point of the rescue (#8335, #7275) -- but only the first means the + installation is sound, and the family-matched install below runs with bare + requirements and no reinstall flag, so the second could write a completion manifest + over a torch that still cannot import. + """ + + _SOURCE = (PACKAGE_ROOT / "studio" / "setup.ps1").read_text(encoding = "utf-8") + + def test_the_probe_distinguishes_a_timeout_from_an_answer(self): + assert "TimedOut = $false" in self._SOURCE + assert "$result.TimedOut = $true" in self._SOURCE + + def test_the_cuda_rescue_forces_a_reinstall_on_a_definitive_failure(self): + block = self._SOURCE[self._SOURCE.index("Test-VenvTorchIsCuda -VenvPath $VenvDir") :][:2000] + assert "$_verProbe -and -not $_verProbe.TimedOut" in block + assert "$script:TorchImportDefinitivelyFailed = $true" in block + + def test_the_venv_is_still_kept_either_way(self): + # A faulted driver raises at DLL load rather than timing out. + start = self._SOURCE.index("Test-VenvTorchIsCuda -VenvPath $VenvDir") + # The CUDA arm only: the trailing "no family matched" else SHOULD rebuild. + arm = self._SOURCE[start : self._SOURCE.index("} else {", start)] + code = "\n".join(line for line in arm.splitlines() if not line.strip().startswith("#")) + assert "$shouldRebuild" not in code + assert "$script:TorchImportDefinitivelyFailed = $true" in arm + + def test_the_flag_reaches_the_cuda_install(self): + assert ( + "if ($script:PinChangedForceReinstall -or $script:TorchImportDefinitivelyFailed) {" + in self._SOURCE + ) + assert "$script:TorchImportDefinitivelyFailed = $false" in self._SOURCE diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py index d2fd7ae8db4..c98b8c4e74d 100644 --- a/tests/studio/install/test_gpu_detection_followups.py +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -284,9 +284,9 @@ def install_src(self) -> str: def test_export_block_uses_leaf(self, install_src): anchor = install_src.find("_torch_index_leaf=") assert anchor >= 0, "backend export must classify on the final path segment" - # Window spans the leaf-normalization prelude (query/frag drop + all-slash trim loop) - # through the export case arms. - window = install_src[anchor : anchor + 900] + # Wide enough to clear the provenance check that now sits between the leaf normalisation + # and the case arms. + window = install_src[anchor : anchor + 1600] assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py index 53d3b49ecb0..91b5e480c40 100644 --- a/tests/studio/install/test_pr5940_followups.py +++ b/tests/studio/install/test_pr5940_followups.py @@ -793,9 +793,15 @@ def test_install_python_stack_windows_rocm_repair_pins_and_is_nonfatal(): k == -1 or j > k ), "Windows ROCm repair must use the nonfatal pip_install_try wrapping the trio" window = text[i : i + 700] + # The trio is built just above the call now (win_arm64 drops torchaudio), but the + # requirement is unchanged: the PINNED companions, never bare names. + trio = text[max(0, i - 900) : i + 700] assert ( - "_torch_pkg" in window and "_vision_pkg" in window and "_audio_pkg" in window + "_torch_pkg" in trio and "_vision_pkg" in trio and "_audio_pkg" in trio ), "Windows ROCm repair must pass the pinned companion trio, not bare names" + assert "*_rocm_trio" in window or ( + "_torch_pkg" in window and "_vision_pkg" in window and "_audio_pkg" in window + ), "the trio the call receives must be the pinned one" assert ( "keeping the existing torch build" in window ), "Windows ROCm repair must keep the existing build (nonfatal) when the index fails" diff --git a/tests/studio/install/test_rocm_arch_table_parity.py b/tests/studio/install/test_rocm_arch_table_parity.py index 993419f2c61..8919754abe9 100644 --- a/tests/studio/install/test_rocm_arch_table_parity.py +++ b/tests/studio/install/test_rocm_arch_table_parity.py @@ -383,9 +383,36 @@ def _name_tables() -> dict[str, object]: "install.ps1": _name_table_ps(_INSTALL_PS1), "studio/setup.ps1": _name_table_ps(_SETUP_PS1), "studio/install_python_stack.py": list(stack_mod._WIN_GPU_NAME_ARCH_TABLE), + # The backend carries the seventh copy: it decides whether a Windows adapter the + # DirectX registry did not give an AdapterFamily is one a repair could help, and + # answering that from a stale table would offer the repair to a card no wheel + # index covers (or withhold it from one that is covered). + "studio/backend/utils/hardware/hardware.py": _name_table_py_literal( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py", + "_GPU_NAME_GFX_TABLE", + ), } +def _name_table_py_literal(path: Path, name: str) -> list: + """A module-level list-of-pairs literal, read without importing the module.""" + import ast + + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in tree.body: + targets = ( + [node.target] + if isinstance(node, ast.AnnAssign) + else node.targets + if isinstance(node, ast.Assign) + else [] + ) + for target in targets: + if isinstance(target, ast.Name) and target.id == name: + return [tuple(pair) for pair in ast.literal_eval(node.value)] + raise AssertionError(f"{name} not found in {path}") + + def _spoof_profiles() -> dict[str, str]: """gfx -> marketing name out of tests/_zoo_rocm_spoof.py::_PROFILES. @@ -577,6 +604,7 @@ def test_divergences_are_real_and_still_diverging(self): "studio/setup.sh", "studio/setup.ps1", "studio/install_python_stack.py", + "studio/backend/utils/hardware/hardware.py", "tests/_zoo_rocm_spoof.py", } diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index db067b5604e..852e2243b07 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2443,7 +2443,8 @@ def test_unsloth_torch_backend_exported_after_get_torch_index_url(self): sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)") - backend_pos = source.find("UNSLOTH_TORCH_BACKEND") + # The export itself, not the first mention: a comment above the block mentions it too. + backend_pos = source.find("export UNSLOTH_TORCH_BACKEND=") assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh" assert ( backend_pos > torch_url_pos diff --git a/tests/studio/install/test_windows_torch_flavor_invariant.py b/tests/studio/install/test_windows_torch_flavor_invariant.py new file mode 100644 index 00000000000..f72c5774c8c --- /dev/null +++ b/tests/studio/install/test_windows_torch_flavor_invariant.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Structural cover for the Windows torch-flavor invariant on the update path. + +The behavioural tests for _ensure_expected_torch_flavor live in test_cuda_repair.py. +This file asserts the parts that are not a Python call: that setup.ps1 hands the flavor +over before it invokes the stack, that setup.ps1 no longer wipes a healthy cu* venv when +nvidia-smi fails to answer, that the repair specs and the mismatch line stay identical to +install.ps1's, that the manifest round-trips the flavor, and that none of it reaches the +Linux/macOS branch of install_python_stack(). Source/AST only -- no Windows required.""" + +import ast +import importlib.util +import json +import re +import sys +import textwrap +from pathlib import Path + +import pytest + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" +_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1" +_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" + +_SETUP_SRC = _SETUP_PS1.read_text(encoding = "utf-8") +_INSTALL_SRC = _INSTALL_PS1.read_text(encoding = "utf-8") +_STACK_SRC = _STACK_PATH.read_text(encoding = "utf-8") + +_MANIFEST_SPEC = importlib.util.spec_from_file_location( + "studio_install_manifest_flavor", PACKAGE_ROOT / "studio" / "install_manifest.py" +) +assert _MANIFEST_SPEC is not None and _MANIFEST_SPEC.loader is not None +install_manifest = importlib.util.module_from_spec(_MANIFEST_SPEC) +sys.modules[_MANIFEST_SPEC.name] = install_manifest +_MANIFEST_SPEC.loader.exec_module(install_manifest) + + +def _line_of(source: str, needle: str) -> int: + """1-based line number of the first line containing `needle`.""" + for number, line in enumerate(source.splitlines(), start = 1): + if needle in line: + return number + raise AssertionError(f"not found in source: {needle!r}") + + +def _publication_block() -> str: + """The whole flavor-publication block, delimited by the section that follows it. + + Sliced by its end marker rather than a character count: a count silently + truncates the moment a comment inside the block grows, and a test that reads + half the block passes for the wrong reason. + """ + start = _SETUP_SRC.index("# ── Publish the torch flavor this run settled on ──") + end = _SETUP_SRC.index("# Ordered heavy dependency installation", start) + return _SETUP_SRC[start:end] + + +class TestSetupPs1NoWipeEscape: + """A direct `studio update` has no rollback copy -- only install.ps1 makes one -- so a + wipe there is unrecoverable. Every way the bounded nvidia-smi probe can come back + empty on a working NVIDIA box collapses the expected tag to "cpu", and a healthy cu124 + venv then reads as stale.""" + + def test_the_escape_sits_ahead_of_the_wipe(self): + escape = _line_of(_SETUP_SRC, "nvidia-smi did not answer, but this venv holds a") + wipe = _line_of(_SETUP_SRC, "Remove-Item -LiteralPath $VenvDir -Recurse -Force") + stale = _line_of(_SETUP_SRC, "Stale venv detected ($reason) -- rebuilding...") + assert escape < stale < wipe, ( + "the no-wipe escape must be evaluated before the stale-venv branch that deletes " + f"the venv (escape={escape}, stale={stale}, wipe={wipe})" + ) + + def test_the_escape_cancels_the_rebuild(self): + body = _SETUP_SRC[_SETUP_SRC.index("nvidia-smi did not answer, but this venv holds a") :][ + :1200 + ] + assert "$shouldRebuild = $false" in body + # Without this the index selection re-runs the same rescan and routes to /cpu. + assert "$script:PreservedInstallerTorchTag = $installedTorchTag" in body + + def test_the_escape_is_narrow(self): + start = _SETUP_SRC.index("if ($shouldRebuild -and -not $InstallerManagedSetup -and\n") + condition = _SETUP_SRC[start : _SETUP_SRC.index("{", start)] + for clause in ( + "-not $InstallerManagedSetup", # install.ps1 repairs in place instead + "-not $_pinnedIdx", # a cpu index PIN is deliberate and still rebuilds + "Test-CudaFamilyLeaf $installedTorchTag", # only a cu* wheel is preserved + "-not $HasNvidiaSmi", # only when the NVIDIA probe gave no answer + '$expectedTorchTag -eq "cpu"', # ... and that is why the expectation collapsed + ): + assert clause in condition, f"the escape must be gated on {clause!r}" + + def test_the_installed_tag_is_tested_before_the_variables_it_implies(self): + # $_pinnedIdx and $expectedTorchTag are assigned only inside `if (-not + # $shouldRebuild)`, so under Set-StrictMode the other -and order is a fatal read. + start = _SETUP_SRC.index("if ($shouldRebuild -and -not $InstallerManagedSetup -and\n") + condition = _SETUP_SRC[start : _SETUP_SRC.index("{", start)] + assert condition.index("$installedTorchTag -and") < condition.index("$_pinnedIdx") + assert condition.index("$installedTorchTag -and") < condition.index("$expectedTorchTag") + + def test_an_xpu_venv_keeps_its_own_escape(self): + # Regression guard: the pre-existing XPU escape must not have been folded in. + assert "Keeping the installed Intel XPU environment" in _SETUP_SRC + + +class TestSetupPs1PublishesTheFlavor: + def test_the_tag_is_exported_before_the_stack_runs(self): + export = _line_of(_SETUP_SRC, "$env:UNSLOTH_EXPECTED_TORCH_TAG =") + index = _line_of(_SETUP_SRC, "$env:UNSLOTH_TORCH_INSTALL_INDEX_URL =") + handoff = _line_of(_SETUP_SRC, 'python "$PSScriptRoot\\install_python_stack.py"') + assert export < handoff and index < handoff + + def test_the_rocm_index_decides_before_the_leaf(self): + # The AMD Windows path installs from repo.amd.com while $TorchInstallIndexUrl still + # points at /cpu. + block = _SETUP_SRC[_SETUP_SRC.index("$_expectedTag = if ($ROCmIndexUrl)") :][:600] + assert block.startswith('$_expectedTag = if ($ROCmIndexUrl) { "rocm" }') + assert "Test-CudaFamilyLeaf $_expectedLeaf" in block + assert "Test-PipRocmFamilyLeaf $_expectedLeaf" in block + # An unknown leaf publishes nothing rather than a tag nothing can verify. + assert "else { $null }" in block + + def test_no_torch_mode_publishes_nothing(self): + assert "if (-not $NoTorchMode) {" in _publication_block() + + def test_unsetting_does_not_depend_on_the_powershell_version(self): + """`$env:X = ""` is not a portable unset, so it must not be used here. + + Windows PowerShell 5.1 and PowerShell 7.0-7.4 delete the entry when it is + assigned an empty string. PowerShell 7.5 took .NET 9's change and KEEPS the + name with an empty value; only $null removes it there. Both spellings still + read as unset through the two readers on the Python side, which test + truthiness rather than presence, so this is about not shipping a line whose + meaning depends on which PowerShell the user happens to have installed. + Remove-Item behaves identically on every version. + + Deleting rather than blanking also matters on its own: a value inherited + from the caller's shell must not survive a run that decided it cannot name + this host's flavor, or the stack enforces a stale expectation. + """ + block = _publication_block() + for name in ( + "UNSLOTH_EXPECTED_TORCH_TAG", + "UNSLOTH_TORCH_INSTALL_INDEX_URL", + ): + assert ( + f"Remove-Item Env:\\{name} -ErrorAction SilentlyContinue" in block + ), f"{name} must be removed, not blanked" + assert f'$env:{name} = ""' not in block + assert f"$env:{name} = if (" not in block + + +class TestInstallPs1Parity: + """A venv repaired by `studio update` and one repaired by install.ps1 must land on the + same wheels, and a support log from either must read the same.""" + + def test_the_repair_trio_matches_install_ps1(self): + match = re.search(r'else\s*\{\s*@\((\s*"torch[^)]*?)\)\s*\}', _INSTALL_SRC, re.S) + assert match is not None, "install.ps1's flavor-repair spec array moved" + ps_specs = tuple(re.findall(r'"([^"]+)"', match.group(1))) + py_specs = tuple( + re.findall( + r'"([^"]+)"', + re.search( + r"_TORCH_FLAVOR_REPAIR_PKG_SPEC: tuple\[str, str, str\] = \((.*?)\)", + _STACK_SRC, + re.S, + ).group(1), + ) + ) + assert py_specs == ps_specs, ( + "_TORCH_FLAVOR_REPAIR_PKG_SPEC must mirror install.ps1's flavor-repair trio " + f"(python={py_specs}, install.ps1={ps_specs})" + ) + + def test_the_mismatch_line_matches_install_ps1(self): + assert ( + "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) " + "-- reinstalling correct build..." + ) in _INSTALL_SRC + assert ( + "PyTorch flavor mismatch (installed {installed}, need {expected}) -- " + ) in _STACK_SRC + assert "reinstalling correct build..." in _STACK_SRC + + def test_the_loud_warning_matches_install_ps1(self): + for line in ( + "PyTorch is CPU-only but a", + "GPU build was expected for this machine.", + "Training and GPU inference will run on CPU until this is fixed.", + "Re-run this installer, or reinstall the GPU build manually for your GPU.", + ): + assert line in _INSTALL_SRC, f"install.ps1 no longer prints {line!r}" + assert line in _STACK_SRC, f"install_python_stack.py no longer prints {line!r}" + + def test_the_flavor_vocabulary_matches_convertto_torchflavortag(self): + arms = _INSTALL_SRC[_INSTALL_SRC.index("function ConvertTo-TorchFlavorTag") :][:900] + assert r"'\+(cu\d+)'" in arms + assert r"'\+rocm'" in arms + assert r"'\+xpu'" in arms + assert r"'\+cpu'" in arms + py = _STACK_SRC[_STACK_SRC.index("def _torch_flavor_tag(") :][:1600] + assert r'r"\+(cu\d+)"' in py + assert '"+rocm" in value' in py + assert '"+xpu" in value' in py + + +def _install_stack_ast(): + tree = ast.parse(_STACK_SRC) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == "install_python_stack": + return node + raise AssertionError("install_python_stack() not found") + + +def _calls_in(node) -> list: + """Every plain function name called under `node`, in source order. + + Depth first, not ast.walk: walk is breadth first, so a nested call reads as if it came + after its own siblings and the assertions below would encode the wrong order. + """ + names = [] + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + names.append(node.func.id) + for child in ast.iter_child_nodes(node): + names.extend(_calls_in(child)) + return names + + +def _guards_containing(call_name: str) -> list: + """Every TOP-LEVEL `if` in install_python_stack() whose body calls `call_name`. + + Plural on purpose: the four existing repair helpers run at two points (the step-2b + check and the step-13 final pass), so a test that took the first match would silently + assert about the wrong one. Top level only, or a guard's own nested `if` counts twice. + """ + found = [ + node + for node in _install_stack_ast().body + if isinstance(node, ast.If) and call_name in _calls_in(node) + ] + assert found, f"no guard around {call_name}()" + return found + + +class TestStepThirteenWiring: + def test_the_flavor_invariant_is_windows_only(self): + guards = _guards_containing("_ensure_expected_torch_flavor") + assert [ast.unparse(guard.test) for guard in guards] == ["IS_WINDOWS and (not NO_TORCH)"] + + def test_a_failed_invariant_returns_non_zero(self): + (guard,) = _guards_containing("_ensure_expected_torch_flavor") + returns = [ + node.value.value + for node in ast.walk(guard) + if isinstance(node, ast.Return) and isinstance(node.value, ast.Constant) + ] + assert returns == [1], ( + "a torch flavor that could not be repaired must fail the install; today that " + "state exits 0 and the app silently runs on CPU" + ) + + def test_the_existing_repair_set_is_untouched(self): + # Step 2b (which Windows enters; the four helpers return early there) and the + # Linux-only step 13. + guards = _guards_containing("_ensure_cuda_torch") + assert [ast.unparse(guard.test) for guard in guards] == [ + "not IS_MACOS and (not NO_TORCH)", + "not IS_WINDOWS and (not IS_MACOS) and (not NO_TORCH)", + ] + for guard in guards: + assert _calls_in(guard) == [ + "_progress", + "_torch_step_label", + "_ensure_cuda_torch", + "_ensure_rocm_torch", + "_ensure_xpu_torch", + "_ensure_cpu_torch", + "_ensure_xpu_triton", + ] + + def test_the_invariant_is_wired_in_exactly_once(self): + body = ast.unparse(_install_stack_ast()) + assert body.count("_ensure_expected_torch_flavor(") == 1 + + +def _base_total(**flags) -> int: + """Re-execute install_python_stack()'s step-total arithmetic under given flags. + + Read out of the function rather than duplicated, so a step added without a matching + total fails here instead of drawing a progress bar past 100%. + """ + lines = _STACK_SRC.splitlines() + start = next(i for i, line in enumerate(lines) if line.strip().startswith("base_total = 12")) + end = next(i for i, line in enumerate(lines) if line.strip().startswith("base_requirements =")) + block = textwrap.dedent("\n".join(lines[start:end])) + namespace = { + "IS_WINDOWS": False, + "IS_MACOS": False, + "NO_TORCH": False, + "IS_MAC_ARM": False, + "skip_base": False, + } + namespace.update(flags) + exec(block, namespace) # noqa: S102 -- the source under test, not user input + return namespace["base_total"] + + +class TestStepTotals: + def test_windows_gained_one_step(self): + assert _base_total(IS_WINDOWS = True) == 14 + assert _base_total(IS_WINDOWS = True, NO_TORCH = True) == 12 + + @pytest.mark.parametrize( + "flags,total", + [ + ({}, 16), # Linux, torch + ({"NO_TORCH": True}, 13), # Linux, GGUF-only + ({"IS_MACOS": True, "IS_MAC_ARM": True}, 13), # Apple Silicon + ({"IS_MACOS": True}, 12), # Intel Mac + ], + ) + def test_the_other_platforms_are_unchanged(self, flags, total): + assert _base_total(**flags) == total + + +class TestManifestRecordsTheFlavor: + def test_round_trip(self, tmp_path): + assert ( + install_manifest.write_manifest( + root = tmp_path, req_root = tmp_path, expected_torch_tag = "cu124" + ) + is not None + ) + assert install_manifest.recorded_torch_flavor(tmp_path) == "cu124" + + def test_the_tag_is_normalised(self, tmp_path): + install_manifest.write_manifest( + root = tmp_path, req_root = tmp_path, expected_torch_tag = " CU128 " + ) + assert install_manifest.recorded_torch_flavor(tmp_path) == "cu128" + + def test_absent_reads_as_unknown_not_cpu(self, tmp_path): + # Claiming a flavor nobody selected would let a repair reinstall over a + # deliberate build. + install_manifest.write_manifest(root = tmp_path, req_root = tmp_path) + assert install_manifest.recorded_torch_flavor(tmp_path) is None + + def test_no_manifest_reads_as_unknown(self, tmp_path): + assert install_manifest.recorded_torch_flavor(tmp_path) is None + + def test_a_hand_edited_non_string_reads_as_unknown(self, tmp_path): + path = install_manifest.manifest_path(tmp_path) + path.write_text(json.dumps({"schema": 1, "expected_torch_tag": 124}), encoding = "utf-8") + assert install_manifest.recorded_torch_flavor(tmp_path) is None + + def test_the_key_is_additive(self, tmp_path): + # MANIFEST_SCHEMA must not move: verify_install rejects a schema it does not know. + assert install_manifest.MANIFEST_SCHEMA == 1 + install_manifest.write_manifest( + root = tmp_path, req_root = tmp_path, expected_torch_tag = "cu124", no_torch = False + ) + payload = json.loads(install_manifest.manifest_path(tmp_path).read_text(encoding = "utf-8")) + assert payload["schema"] == 1 + assert payload["no_torch"] is False + assert payload["expected_torch_tag"] == "cu124" + + def test_no_index_url_is_ever_written(self, tmp_path): + # A pinned index can carry a token; this file sits in the venv and is read back. + install_manifest.write_manifest( + root = tmp_path, req_root = tmp_path, expected_torch_tag = "cu124" + ) + raw = install_manifest.manifest_path(tmp_path).read_text(encoding = "utf-8") + assert "http" not in raw + + def test_the_stack_carries_a_previous_record_forward(self): + # A platform that never resolves a flavor must not erase the one already recorded. + assert "expected_torch_tag = _recordable_torch_flavor_tag(torch_flavor_tag)," in _STACK_SRC + helper = _STACK_SRC[_STACK_SRC.index("def _recordable_torch_flavor_tag(") :] + helper = helper[: helper.index("\ndef ", 1)] + assert 'return _RECORDED_TORCH_TAG or ""' in helper + + def test_a_mirror_pin_does_not_carry_a_stale_flavor_forward(self): + # The wheel came from a mirror whose leaf names no family, so the previous record + # describes a venv that no longer exists and would hand a later unpinned run a flavor + # to "repair" the mirror's build back to. + helper = _STACK_SRC[_STACK_SRC.index("def _recordable_torch_flavor_tag(") :] + helper = helper[: helper.index("\ndef ", 1)] + assert "_explicit_unknown_family_torch_index_url() is not None" in helper + assert helper.index("_explicit_unknown_family_torch_index_url") < helper.index( + "_RECORDED_TORCH_TAG" + ), "the mirror check has to come before the carry-forward" + + def test_the_record_is_read_before_the_manifest_is_dropped(self): + # install_python_stack() removes the manifest before its dependency pass. + read = _line_of( + _STACK_SRC, "_RECORDED_TORCH_TAG = install_manifest.recorded_torch_flavor()" + ) + drop = _line_of(_STACK_SRC, "if not install_manifest.remove_manifest():") + assert read < drop + assert ( + "def install_python_stack" + not in _STACK_SRC[: _STACK_SRC.index("_RECORDED_TORCH_TAG =")] + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q"])) + + +class TestTheFlavorProvenance: + """A recorded flavor is only a CHOICE when someone named it.""" + + def test_install_sh_marks_a_derived_backend_as_derived(self): + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + block = source[source.index('case "$_torch_index_leaf" in') :] + block = block[: block.index("_is_pip_rocm_family_leaf")] + assert 'UNSLOTH_TORCH_BACKEND_SOURCE="resolved"' in block, ( + "install.sh derives the backend from the index it resolved -- cpu on any " + "GPU-less host -- so the manifest has to be told which it was" + ) + + def test_a_backend_the_caller_stated_is_not_marked_derived(self): + # On a GPU-less host the resolved value is cpu too, so a stated choice and the + # automatic one are indistinguishable unless install.sh checks BEFORE overwriting. + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + assert "_torch_backend_was_stated" in source + check = source.index("_torch_backend_was_stated=true") + overwrite = source.index('case "$_torch_index_leaf" in') + assert check < overwrite, "the check has to run before the assignment" + mark = source.index('UNSLOTH_TORCH_BACKEND_SOURCE="resolved"') + line_start = source.rindex("if ", 0, mark) + assert ( + "_torch_backend_was_stated" in source[line_start:mark] + ), "a stated backend must not be marked derived" + + def test_the_stack_reads_that_marker(self): + pinned = _STACK_SRC[_STACK_SRC.index("def _expected_torch_flavor_was_pinned(") :] + pinned = pinned[: pinned.index("\ndef ", 1)] + assert "UNSLOTH_TORCH_BACKEND_SOURCE" in pinned + assert "resolved" in pinned + + def test_an_untagged_xpu_runtime_counts_as_a_gpu_build(self): + # An untagged source or conda XPU build carries its runtime in torch.version.xpu and + # nowhere else, and calling it CPU-only fails the update outright. + assert "getattr(_v, 'xpu', '')" in _STACK_SRC + verdict = _STACK_SRC[_STACK_SRC.index("def _torch_build_is_gpu(") :] + verdict = verdict[: verdict.index("\ndef ", 1)] + assert "_TORCH_RUNTIME_XPU" in verdict + + +class TestABrokenTorchForcesTheDependencyPass: + """$script:TorchImportDefinitivelyFailed is set when the probe reports an import + error rather than a timeout. Its ONLY consumer -- the CUDA --force-reinstall -- lives + inside the `if (-not $SkipPythonDeps)` block, so on a current core package with a + valid manifest the fast path skipped it and the update reported dependencies up to + date over a wheel that will not load.""" + + def test_the_flag_clears_the_skip(self): + clear = _SETUP_SRC.index( + "if ($script:PinChangedForceReinstall -or $script:TorchImportDefinitivelyFailed) {" + ) + block = _SETUP_SRC[clear : _SETUP_SRC.index("if (-not $SkipPythonDeps) {", clear)] + assert "$SkipPythonDeps = $false" in block + + def test_the_clear_runs_before_the_block_it_unlocks(self): + clear = _line_of(_SETUP_SRC, "$script:TorchImportDefinitivelyFailed) {") + guard = _line_of(_SETUP_SRC, "if (-not $SkipPythonDeps) {") + # The LAST such line: the clear itself now reads the same flag, and _line_of + # takes the first match. + consumer = max( + number + for number, line in enumerate(_SETUP_SRC.splitlines(), start = 1) + if "$script:TorchImportDefinitivelyFailed" in line + ) + assert clear < guard < consumer, ( + f"the flag clears the skip at {clear}, the block opens at {guard} and its " + f"consumer is at {consumer}; any other order leaves the repair unreachable" + ) + + def test_the_flag_is_still_raised_where_the_import_definitively_failed(self): + assert "$script:TorchImportDefinitivelyFailed = $true" in _SETUP_SRC diff --git a/tests/studio/test_xpu_triton_swap.py b/tests/studio/test_xpu_triton_swap.py index f184597568f..0a40d22c6e8 100644 --- a/tests/studio/test_xpu_triton_swap.py +++ b/tests/studio/test_xpu_triton_swap.py @@ -461,20 +461,40 @@ def test_a_timeout_on_a_supported_wheel_reinstalls_nothing(self): class TestPlatformGuards: - @pytest.mark.parametrize("flag", ["NO_TORCH", "IS_MACOS", "IS_WINDOWS"]) + @pytest.mark.parametrize("flag", ["NO_TORCH", "IS_MACOS"]) def test_skipped_where_it_does_not_apply(self, monkeypatch, tmp_path, flag): - # Windows is setup.ps1's job; macOS has no XPU; --no-torch touches no wheels. + monkeypatch.delenv("UNSLOTH_EXPECTED_TORCH_TAG", raising = False) mod, log = _load(monkeypatch, tmp_path, spec = "pytorch-triton-xpu==3.5.0", generic = "3.7.1") mod.__dict__[flag] = True mod.__dict__["_ensure_xpu_triton"].__globals__[flag] = True mod.__dict__["_ensure_xpu_triton"]() assert log == [] + def test_windows_defers_to_setup_ps1_when_setup_ps1_ran(self, monkeypatch, tmp_path): + # setup.ps1 performs the same swap after this file exits, and publishes the + # handover variable immediately before invoking it. + monkeypatch.setenv("UNSLOTH_EXPECTED_TORCH_TAG", "xpu") + mod, log = _load(monkeypatch, tmp_path, spec = "pytorch-triton-xpu==3.5.0", generic = "3.7.1") + mod.__dict__["_ensure_xpu_triton"].__globals__["IS_WINDOWS"] = True + mod.__dict__["_ensure_xpu_triton"]() + assert log == [] + + def test_a_direct_windows_run_does_the_swap_itself(self, monkeypatch, tmp_path): + # Bare `python install_python_stack.py` on Windows has no setup.ps1 postlude, so + # the core install leaves triton-windows over torch's XPU triton. The absent + # handover variable is the signal that nobody else will fix it. + monkeypatch.delenv("UNSLOTH_EXPECTED_TORCH_TAG", raising = False) + mod, log = _load(monkeypatch, tmp_path, spec = "pytorch-triton-xpu==3.5.0", generic = "3.7.1") + mod.__dict__["_ensure_xpu_triton"].__globals__["IS_WINDOWS"] = True + mod.__dict__["_ensure_xpu_triton"]() + assert "INSTALL" in log -def test_the_swap_is_wired_in_at_both_repair_points(): - # The final repair pass would otherwise silently undo the first. + +def test_the_swap_is_wired_in_at_every_repair_point(): + # The final repair pass would otherwise silently undo the first. The third point is + # step 13w, the Windows flavor invariant. src = STACK.read_text(encoding = "utf-8") - assert src.count(" _ensure_xpu_triton()") == 2 + assert src.count(" _ensure_xpu_triton()") == 3 def test_the_swap_runs_after_every_torch_migration(): @@ -502,9 +522,14 @@ def test_the_swap_runs_after_every_torch_migration(): ] if "_ensure_xpu_triton" in calls: blocks.append(calls) - assert len(blocks) == 2, f"expected 2 repair blocks, found {len(blocks)}: {blocks}" + assert len(blocks) == 3, f"expected 3 repair blocks, found {len(blocks)}: {blocks}" for calls in blocks: assert calls[-1] == "_ensure_xpu_triton", calls + # Step 13w's migration is _ensure_expected_torch_flavor, which the walk above does + # not collect (its result is branched on, not discarded), so it is asserted on source. + migrating = [c for c in blocks if "_ensure_cuda_torch" in c] + assert len(migrating) == 2, blocks + for calls in migrating: for migration in ( "_ensure_cuda_torch", "_ensure_rocm_torch", @@ -512,6 +537,11 @@ def test_the_swap_runs_after_every_torch_migration(): "_ensure_cpu_torch", ): assert calls.index(migration) < calls.index("_ensure_xpu_triton"), (migration, calls) + src = STACK.read_text(encoding = "utf-8") + windows = src[src.index("# 13w.") : src.index("# 14.")] + assert windows.index("_ensure_expected_torch_flavor") < windows.index( + "_ensure_xpu_triton" + ), "the Windows swap must follow that platform's torch migration too" def test_install_sh_does_not_carry_a_second_copy():