From a8cfba2eb2e58a555648969a9b94e4e22fafc779 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 07:15:52 +0000 Subject: [PATCH] Gate the unsloth retry in the diffusion patch backend The retry added for the clean-environment patch failures is not free: importing unsloth pulls torch in behind it, which costs ~940 MB of RSS measured in a process that had neither, and on a host with no accelerator it fails anyway. A cross-platform CI job that had generated fine at ~900 s later died 19 s in with SIGTERM and every 'if: always()' step skipped, which is the runner being torn down rather than a step failing. Retry only when torch is already imported (true of the server and of anything patching a real module, and the condition that stops the retry from being what loads torch), unsloth is installed but not yet imported, and the first failure was the ImportError the sentinel guard raises. The clean-environment case it was added for still passes 29/29. --- .../core/inference/diffusion_patch_backend.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion_patch_backend.py b/studio/backend/core/inference/diffusion_patch_backend.py index c7e3aaf8bf..29cf96e5d0 100644 --- a/studio/backend/core/inference/diffusion_patch_backend.py +++ b/studio/backend/core/inference/diffusion_patch_backend.py @@ -35,7 +35,17 @@ def _helpers() -> Optional[dict]: always resolved there, but ANY process that reaches the patch backend first -- the test suite, a worker subprocess -- got an ImportError and silently ran unpatched (every install returning False). So on failure, import ``unsloth`` and retry once, which is also the import order Unsloth - documents. A host with no accelerator still fails both attempts and stays a no-op.""" + documents. + + The retry is gated, because it is not free: importing ``unsloth`` pulls torch in behind it and + costs ~940 MB of RSS in a process that had neither, only to fail anyway on a host with no + accelerator, which is enough to matter on a small CI runner mid-generation. So it runs only when + + * ``torch`` is already imported -- true of the server and of anything that patches a real + module, and the condition that keeps the retry from being the thing that loads torch, + * ``unsloth`` is installed but not yet imported (if it were, the sentinel would be set and the + first attempt would have worked), + * and the first failure was the ImportError that guard raises.""" global _HELPERS if _HELPERS is not None: return _HELPERS or None @@ -44,12 +54,25 @@ def _helpers() -> Optional[dict]: from unsloth_zoo.temporary_patches.utils import patch_function, restore_original return {"patch": patch_function, "restore": restore_original} + def _retry_could_help(exc: BaseException) -> bool: + import importlib.util + import sys + + if not isinstance(exc, ImportError) or "unsloth" in sys.modules: + return False + if "torch" not in sys.modules: + return False + try: + return importlib.util.find_spec("unsloth") is not None + except Exception: # noqa: BLE001 — an unimportable package cannot set the sentinel either + return False + for attempt in (0, 1): try: _HELPERS = _load() return _HELPERS - except Exception: # noqa: BLE001 — no unsloth_zoo / no-GPU host -> optimisation skipped - if attempt: + except Exception as exc: # noqa: BLE001 — no unsloth_zoo / no-GPU host -> skip the patch + if attempt or not _retry_could_help(exc): break try: import unsloth # noqa: F401 — sets UNSLOTH_IS_PRESENT for the retry