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.
This commit is contained in:
Daniel Han 2026-07-27 07:15:52 +00:00
commit a8cfba2eb2

View file

@ -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