a6dc10dad2* Reduce and tighten comments and docstrings in tests Shorten verbose comments and docstrings across the test suite without changing any test logic. Remove narration that restates the next line, collapse long module and test docstrings to a single line, and drop banner separators. Keep regression context (issue and PR references, run ids), skip reasons, mocking and timing rationale, license headers, lint and type directives, and commented-out code. Comments and docstrings only: an AST signature check confirms no code, assertions, or string literals changed, and the suite byte-compiles cleanly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
148 lines
5 KiB
Python
148 lines
5 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""GPU-free test harness.
|
|
|
|
unsloth_zoo.device_type calls get_device_type() at import time and raises
|
|
NotImplementedError on CI runners with no CUDA/XPU/HIP. Pre-load it under a
|
|
mocked torch.cuda.is_available()==True so its @cache permanently captures
|
|
"cuda"; on a real accelerator the pre-load is skipped.
|
|
|
|
Mirrors the conftest harness in unslothai/unsloth-zoo PR #624.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
import types
|
|
|
|
|
|
def _has_real_accelerator() -> bool:
|
|
try:
|
|
import torch
|
|
except Exception:
|
|
return False
|
|
for probe in (
|
|
lambda: hasattr(torch, "cuda") and torch.cuda.is_available(),
|
|
lambda: hasattr(torch, "xpu") and torch.xpu.is_available(),
|
|
lambda: hasattr(torch, "accelerator") and torch.accelerator.is_available(),
|
|
):
|
|
try:
|
|
if probe():
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool:
|
|
"""Pre-load <package>.device_type under a mocked is_available()==True so its
|
|
@cache captures "cuda"; prereqs are submodules to load first (e.g. 'utils').
|
|
Returns False if anything is unimportable, so the caller falls back to a stub."""
|
|
target = f"{package}.device_type"
|
|
if target in sys.modules:
|
|
return True
|
|
pkg_spec = importlib.util.find_spec(package)
|
|
if pkg_spec is None or not pkg_spec.submodule_search_locations:
|
|
return False
|
|
pkg_path = pkg_spec.submodule_search_locations[0]
|
|
|
|
skeleton_already = package in sys.modules
|
|
if not skeleton_already:
|
|
skel = types.ModuleType(package)
|
|
skel.__path__ = [pkg_path]
|
|
skel.__spec__ = pkg_spec
|
|
skel.__package__ = package
|
|
sys.modules[package] = skel
|
|
|
|
try:
|
|
for prereq in prereqs:
|
|
full = f"{package}.{prereq}"
|
|
if full in sys.modules:
|
|
continue
|
|
prereq_path = os.path.join(pkg_path, f"{prereq}.py")
|
|
prereq_spec = importlib.util.spec_from_file_location(full, prereq_path)
|
|
prereq_mod = importlib.util.module_from_spec(prereq_spec)
|
|
sys.modules[full] = prereq_mod
|
|
prereq_spec.loader.exec_module(prereq_mod)
|
|
|
|
device_type_path = os.path.join(pkg_path, "device_type.py")
|
|
dt_spec = importlib.util.spec_from_file_location(target, device_type_path)
|
|
dt_mod = importlib.util.module_from_spec(dt_spec)
|
|
sys.modules[target] = dt_mod
|
|
|
|
import torch
|
|
|
|
_orig_is_avail = torch.cuda.is_available
|
|
torch.cuda.is_available = lambda: True # type: ignore[assignment]
|
|
try:
|
|
dt_spec.loader.exec_module(dt_mod)
|
|
finally:
|
|
torch.cuda.is_available = _orig_is_avail
|
|
except Exception:
|
|
sys.modules.pop(target, None)
|
|
return False
|
|
finally:
|
|
if not skeleton_already:
|
|
sys.modules.pop(package, None)
|
|
|
|
return True
|
|
|
|
|
|
def _patch_torch_cuda_for_import() -> None:
|
|
"""Stub the torch.cuda.* probes fired at import time once DEVICE_TYPE is
|
|
forced to "cuda"; returning plausible Ampere values lets the import finish
|
|
(real-tensor tests still run on CPU)."""
|
|
try:
|
|
import torch.cuda.memory as _cuda_memory # type: ignore
|
|
_cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
import torch
|
|
torch.cuda.get_device_capability = lambda *a, **k: (8, 0)
|
|
torch.cuda.is_bf16_supported = lambda *a, **k: True
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _install_device_type_stub(name: str) -> None:
|
|
stub = types.ModuleType(name)
|
|
stub.DEVICE_TYPE = "cuda"
|
|
stub.DEVICE_TYPE_TORCH = "cuda"
|
|
stub.DEVICE_COUNT = 1
|
|
stub.ALLOW_PREQUANTIZED_MODELS = False
|
|
stub.is_hip = lambda: False
|
|
stub.get_device_type = lambda: "cuda"
|
|
stub.get_device_count = lambda: 1
|
|
stub.device_synchronize = lambda *a, **k: None
|
|
stub.device_empty_cache = lambda *a, **k: None
|
|
stub.device_is_bf16_supported = lambda *a, **k: False
|
|
sys.modules[name] = stub
|
|
|
|
|
|
if not _has_real_accelerator():
|
|
if not _preload_device_type("unsloth_zoo", prereqs = ("utils",)):
|
|
_install_device_type_stub("unsloth_zoo.device_type")
|
|
if not _preload_device_type("unsloth"):
|
|
_install_device_type_stub("unsloth.device_type")
|
|
_patch_torch_cuda_for_import()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Apply upstream-drift fixes (vllm/triton/peft) by triggering ``import unsloth``
|
|
# (they run at import time in unsloth/import_fixes.py). The harness above lets
|
|
# the import survive CPU-only runners; the ImportError is swallowed otherwise.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _apply_upstream_import_fixes_for_tests() -> None:
|
|
try:
|
|
import unsloth # noqa: F401 # runs unsloth/import_fixes.py
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
_apply_upstream_import_fixes_for_tests()
|