diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index a037ea2579..c0b624e27b 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -307,6 +307,51 @@ def test_degrades_when_shared_helper_import_raises_importerror(): sys.modules["utils.hf_xet_fallback"] = saved_shim +def test_no_light_gpu_init_retry_on_an_accelerator_host(monkeypatch): + """The UNSLOTH_ZOO_DISABLE_GPU_INIT retry makes unsloth_zoo take its MLX/CPU path, which injects + triton and bitsandbytes STUBS into sys.modules for the whole process. On a GPU host whose + unsloth_zoo import failed for an unrelated reason (a bitsandbytes/CUDA mismatch, say), those + stubs raise "called on Apple Silicon / MLX" from the first CUDA-only kernel a later GGUF or + compiled diffusion generation touches, so a healthy GPU starts 500ing. The shim must degrade + instead of retrying there.""" + import importlib + import os + + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + attempts = [] + + class _Blocker: + def find_spec(self, name, path = None, target = None): + if name == "unsloth_zoo": + attempts.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise RuntimeError("CUDA Setup failed despite GPU being available") + return None + + finder = _Blocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + shim = importlib.import_module("utils.hf_xet_fallback") + monkeypatch.setattr(shim, "_gpu_present", lambda: True) + assert shim._load_shared() is False + # Exactly ONE attempt, made without the light-init flag. + assert attempts == [None], attempts + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. @@ -345,6 +390,9 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.meta_path.insert(0, finder) try: degraded = importlib.import_module("utils.hf_xet_fallback") + # The retry only applies to a host with no accelerator (see _gpu_present); pin that on the + # freshly imported module so this keeps exercising the light-init path. + monkeypatch.setattr(degraded, "_gpu_present", lambda: False) # Import is light (lazy backend); unsloth_zoo not loaded yet. assert seen_env == [], seen_env # First use of a heavy helper triggers the load (attempt without the light env, then a retry diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 49872f371e..f2a96d4490 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -38,6 +38,29 @@ _shared_import_error: Optional[BaseException] = None _load_lock = threading.Lock() +def _gpu_present() -> bool: + """Whether this host has a usable accelerator, decided WITHOUT importing unsloth_zoo. + + Only torch is consulted (already imported by the time any download helper runs), and any + failure answers False so a genuinely torch-less host keeps the light-init retry below. + """ + try: + import torch + except Exception: # noqa: BLE001 -- no torch at all: the light path is the right one + return False + for probe in ( + lambda: torch.cuda.is_available(), + lambda: torch.backends.mps.is_available(), + lambda: torch.xpu.is_available(), + ): + try: + if probe(): + return True + except Exception: # noqa: BLE001 -- a missing backend is just "not this one" + continue + return False + + def _load_shared() -> bool: """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so importing this module at worker startup does not pull transformers in before the sidecar is @@ -61,6 +84,26 @@ def _load_shared() -> bool: _shared_import_error = exc import os as _os + # ...but ONLY on a host that really has no accelerator. That flag makes unsloth_zoo take + # its MLX/CPU path, which injects triton and bitsandbytes STUBS into sys.modules for the + # rest of the process. On a working GPU box (where the import failed for an unrelated + # reason, e.g. a bitsandbytes/CUDA mismatch) those stubs then raise + # "called on Apple Silicon / MLX" from the first CUDA-only kernel a later GGUF or + # compiled diffusion generation touches, turning a healthy GPU into 500s. The Xet + # watchdog is optional; a poisoned process is not, so degrade instead. + if _gpu_present(): + _shared_available = False + import logging as _logging + + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Not retrying under UNSLOTH_ZOO_DISABLE_GPU_INIT because this host " + "has an accelerator and that path would stub out triton/bitsandbytes for the " + "whole process.", + exc, + ) + return False + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" try: