fix(registry): don't register deepseek models at import time (#7227)

* fix(registry): don't register deepseek models at import time

`_deepseek.py` called `register_deepseek_models(include_original_model=True)`
at module scope, so merely importing `unsloth.registry` registered models
(and reached the hub via `list_models`) as a side effect. None of the other
five families (`_gemma`/`_llama`/`_mistral`/`_phi`/`_qwen`) do this; they only
register when `register_models()` asks them to.

Two consequences:
- Importing the registry populated MODEL_REGISTRY on its own (32 entries,
  including 10 `deepseek-ai` original models that no other family leaks) and
  did network I/O at import time.
- Because the import-time call set the `_IS_DEEPSEEK_*_REGISTERED` guards with
  `include_original_model=True`, the later `register_models()` call (which uses
  the default `include_original_model=False`) early-returned, so the
  original-model set won permanently.

Remove the stray module-level call. The `if __name__ == "__main__"` block below
still registers with `include_original_model=True` for standalone use, so the
generator script is unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test(registry): make import-side-effect test pass on CPU-only runners

The new test spawned a fresh `python -c "import unsloth.registry"` that did
not inherit tests/conftest.py's GPU-free harness, so on no-accelerator CI
runners the child raised NotImplementedError from unsloth_zoo.device_type
before printing REGISTRY_SIZE. With check=True this surfaced only as an
opaque CalledProcessError, turning the "Repo tests (CPU)" job red even
though the registry fix is correct.

Import this directory's conftest inside the child first so it applies the
same device_type stubs and torch.cuda probe patches. Also use check=False
and include the child stdout/stderr in the assertion message so a future
import regression is legible instead of an opaque non-zero exit.

* test(registry): assert register_models() leaks no upstream originals

Adds a fresh-interpreter test that register_models() registers only
unsloth-org models (deepseek still present via the normal path) and never
leaks the upstream deepseek-ai originals that the import-time guard poisoning
used to leak (129 -> 139). Factors the conftest-harness subprocess runner
into a shared helper reused by both registry import tests.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
This commit is contained in:
Andrew Chen 2026-07-19 18:37:23 +08:00 committed by GitHub
commit aef36cfdd4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 85 additions and 2 deletions

View file

@ -1,5 +1,8 @@
"""Register each model set and check the registered ids exist on the HF Hub.""" """Register each model set and check the registered ids exist on the HF Hub."""
import os
import subprocess
import sys
from dataclasses import dataclass from dataclasses import dataclass
import pytest import pytest
@ -77,3 +80,85 @@ def test_quant_type():
assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models) assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH] quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
assert all(quant_tag in m.model_path for m in dynamic_quant_models) assert all(quant_tag in m.model_path for m in dynamic_quant_models)
def _run_registry_child(body: str) -> subprocess.CompletedProcess:
"""Run ``body`` in a fresh interpreter that first imports this directory's
``conftest`` so it inherits the same GPU-free harness the pytest session
uses (device_type stubs plus torch.cuda probe patches). Without it,
``import unsloth.registry`` raises ``NotImplementedError`` from
``unsloth_zoo.device_type`` on no-accelerator CI runners, so the child
would exit non-zero and the test would fail even though the registry code
is correct. A fresh process also keeps each check independent of any
``register_models()`` calls other tests make on the shared registry.
"""
tests_dir = os.path.dirname(os.path.abspath(__file__))
prelude = (
f"import sys; sys.path.insert(0, {tests_dir!r})\n"
"try:\n"
" import conftest # noqa: F401 GPU-free harness on no-accelerator runners\n"
"except Exception:\n"
" pass\n"
)
return subprocess.run(
[sys.executable, "-c", prelude + body],
capture_output = True,
text = True,
check = False,
)
def test_importing_registry_does_not_register_models():
"""Importing the registry must not populate MODEL_REGISTRY on its own.
``_deepseek`` used to call ``register_deepseek_models(...)`` at module
scope, so merely importing ``unsloth.registry`` registered models as an
import side effect, unlike every other family which only registers on
demand.
"""
result = _run_registry_child(
"import unsloth.registry\n"
"from unsloth.registry.registry import MODEL_REGISTRY\n"
"print('REGISTRY_SIZE', len(MODEL_REGISTRY))"
)
assert result.returncode == 0, (
f"registry import subprocess exited {result.returncode}\n"
f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
size_lines = [line for line in result.stdout.splitlines() if line.startswith("REGISTRY_SIZE")]
assert size_lines == ["REGISTRY_SIZE 0"], result.stdout + result.stderr
def test_register_models_registers_no_upstream_originals():
"""``register_models()`` must register each family's ``unsloth``-org models
and must NOT leak upstream vendor "original" models.
Before the fix, ``_deepseek``'s import-time
``register_deepseek_models(include_original_model = True)`` set the
``_IS_DEEPSEEK_*_REGISTERED`` guards, so the later default
``register_models()`` early-returned for deepseek and its 10 ``deepseek-ai``
originals leaked permanently (129 -> 139). This asserts the whole registry
is ``unsloth``-org after ``register_models()`` while deepseek is still
registered via the normal path. Runs in a fresh interpreter so it is
independent of other tests' registry mutations.
"""
result = _run_registry_child(
"import unsloth.registry\n"
"from unsloth.registry import register_models\n"
"from unsloth.registry.registry import MODEL_REGISTRY\n"
"register_models()\n"
"orgs = sorted({m.org for m in MODEL_REGISTRY.values()})\n"
"deepseek = [k for k in MODEL_REGISTRY if 'deepseek' in k.lower()]\n"
"print('ORGS', orgs)\n"
"print('NUM_DEEPSEEK', len(deepseek))"
)
assert result.returncode == 0, (
f"register_models subprocess exited {result.returncode}\n"
f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
out = result.stdout
# Every registered model is unsloth-org: no upstream "original" leaked.
assert "ORGS ['unsloth']" in out, out + result.stderr
# Deepseek is still registered via the normal path, just without originals.
deepseek_lines = [line for line in out.splitlines() if line.startswith("NUM_DEEPSEEK")]
assert deepseek_lines and int(deepseek_lines[0].split()[1]) > 0, out + result.stderr

View file

@ -171,8 +171,6 @@ def _list_deepseek_r1_distill_models():
return distill_models return distill_models
register_deepseek_models(include_original_model = True)
if __name__ == "__main__": if __name__ == "__main__":
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info