mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
617 lines
26 KiB
Python
617 lines
26 KiB
Python
"""Offline model-loader contracts at provider, device, and quantization boundaries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import builtins
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from obliteratus.models import loader
|
|
|
|
|
|
def _config(**overrides):
|
|
values = {
|
|
"model_type": "gpt2",
|
|
"architectures": ["GPT2LMHeadModel"],
|
|
"num_hidden_layers": 2,
|
|
"num_attention_heads": 4,
|
|
"hidden_size": 8,
|
|
"intermediate_size": 16,
|
|
"vocab_size": 32,
|
|
"quantization_config": None,
|
|
}
|
|
values.update(overrides)
|
|
return SimpleNamespace(**values)
|
|
|
|
|
|
def _model():
|
|
model = MagicMock()
|
|
model.to.return_value = model
|
|
model.state_dict.return_value = {"weight": torch.ones(2)}
|
|
model.parameters.return_value = iter([torch.ones(2), torch.ones(3)])
|
|
return model
|
|
|
|
|
|
@pytest.fixture
|
|
def loader_boundary(monkeypatch):
|
|
config = _config()
|
|
model = _model()
|
|
tokenizer = SimpleNamespace(pad_token=None, eos_token="<eos>")
|
|
model_class = SimpleNamespace(from_pretrained=Mock(return_value=model))
|
|
monkeypatch.setattr(loader.AutoConfig, "from_pretrained", Mock(return_value=config))
|
|
monkeypatch.setattr(loader.AutoTokenizer, "from_pretrained", Mock(return_value=tokenizer))
|
|
monkeypatch.setitem(loader.TASK_MODEL_MAP, "causal_lm", model_class)
|
|
monkeypatch.setitem(loader.TASK_MODEL_MAP, "classification", model_class)
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda preference="auto": "cpu" if preference == "auto" else preference)
|
|
monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: False)
|
|
monkeypatch.setattr(loader.dev, "supports_bitsandbytes", lambda _device=None: False)
|
|
monkeypatch.setattr(loader.dev, "supports_bfloat16", lambda _device=None: True)
|
|
monkeypatch.setattr(loader.dev, "get_total_free_gb", lambda: 0.0)
|
|
monkeypatch.setattr(loader.dev, "empty_cache", Mock())
|
|
monkeypatch.setattr(loader.dev, "is_cuda", lambda: False)
|
|
return SimpleNamespace(config=config, model=model, tokenizer=tokenizer, model_class=model_class)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("kwargs", "message"),
|
|
[
|
|
({"model_name": ""}, "model_name must be"),
|
|
({"model_name": "x", "task": "embedding"}, "Unknown task"),
|
|
({"model_name": "x", "dtype": "int9"}, "Unknown dtype"),
|
|
({"model_name": "x", "quantization": "3bit"}, "Unknown quantization"),
|
|
({"model_name": "x", "gpu_memory_utilization": 0}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": -0.1}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": 1.1}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": float("nan")}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": float("inf")}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": True}, "gpu_memory_utilization"),
|
|
({"model_name": "x", "gpu_memory_utilization": "0.8"}, "gpu_memory_utilization"),
|
|
],
|
|
)
|
|
def test_invalid_requests_fail_before_provider_access(loader_boundary, monkeypatch, kwargs, message):
|
|
config_load = Mock(side_effect=AssertionError("provider should not be called"))
|
|
monkeypatch.setattr(loader.AutoConfig, "from_pretrained", config_load)
|
|
with pytest.raises(ValueError, match=message):
|
|
loader.load_model(**kwargs)
|
|
config_load.assert_not_called()
|
|
|
|
|
|
def test_revision_trust_and_offline_flags_reach_every_provider(loader_boundary):
|
|
handle = loader.load_model(
|
|
"local/model",
|
|
revision="deadbeef",
|
|
trust_remote_code=True,
|
|
local_files_only=True,
|
|
skip_snapshot=True,
|
|
)
|
|
expected = {
|
|
"trust_remote_code": True,
|
|
"token": None,
|
|
"revision": "deadbeef",
|
|
"local_files_only": True,
|
|
}
|
|
loader.AutoConfig.from_pretrained.assert_called_once_with("local/model", **expected)
|
|
loader_boundary.model_class.from_pretrained.assert_called_once()
|
|
for key, value in expected.items():
|
|
assert loader_boundary.model_class.from_pretrained.call_args.kwargs[key] == value
|
|
loader.AutoTokenizer.from_pretrained.assert_called_once_with("local/model", **expected)
|
|
assert handle.tokenizer.pad_token == "<eos>"
|
|
loader_boundary.model.eval.assert_called_once_with()
|
|
loader.dev.empty_cache.assert_called_once_with()
|
|
|
|
|
|
def test_hf_token_is_forwarded_without_logging_value(loader_boundary, monkeypatch):
|
|
monkeypatch.setenv("HF_TOKEN", "secret-token")
|
|
loader.load_model("x", skip_snapshot=True)
|
|
assert loader.AutoConfig.from_pretrained.call_args.kwargs["token"] == "secret-token"
|
|
assert loader_boundary.model_class.from_pretrained.call_args.kwargs["token"] == "secret-token"
|
|
assert loader.AutoTokenizer.from_pretrained.call_args.kwargs["token"] == "secret-token"
|
|
|
|
|
|
def test_device_and_dtype_failures_precede_provider_access(loader_boundary, monkeypatch):
|
|
provider = Mock(side_effect=AssertionError("provider should not be called"))
|
|
monkeypatch.setattr(loader.AutoConfig, "from_pretrained", provider)
|
|
monkeypatch.setattr(loader.dev, "get_device", Mock(side_effect=RuntimeError("CUDA unavailable")))
|
|
with pytest.raises(RuntimeError, match="CUDA unavailable"):
|
|
loader.load_model("x", device="cuda")
|
|
provider.assert_not_called()
|
|
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "mps")
|
|
monkeypatch.setattr(loader.dev, "supports_bfloat16", lambda _device=None: False)
|
|
with pytest.raises(RuntimeError, match="bfloat16 is not supported on 'mps'"):
|
|
loader.load_model("x", device="mps", dtype="bfloat16")
|
|
provider.assert_not_called()
|
|
|
|
|
|
def test_config_permission_error_retries_in_temp_cache(loader_boundary, monkeypatch, tmp_path):
|
|
monkeypatch.setattr(loader.tempfile, "gettempdir", lambda: str(tmp_path))
|
|
loader.AutoConfig.from_pretrained.side_effect = [PermissionError("cache"), loader_boundary.config]
|
|
loader.load_model("x", skip_snapshot=True)
|
|
assert loader.AutoConfig.from_pretrained.call_count == 2
|
|
assert loader.AutoConfig.from_pretrained.call_args.kwargs["cache_dir"] == str(tmp_path / "hf_home" / "hub")
|
|
assert (tmp_path / "hf_home" / "hub").is_dir()
|
|
|
|
|
|
@pytest.mark.parametrize("message", ["Gated repo denied", "Access to model is restricted"])
|
|
def test_config_gated_repo_failure_is_actionable(loader_boundary, monkeypatch, message):
|
|
loader.AutoConfig.from_pretrained.side_effect = OSError(message)
|
|
with pytest.raises(RuntimeError, match=r"(?s)Accept the license.*HF_TOKEN"):
|
|
loader.load_model("owner/gated", skip_snapshot=True)
|
|
|
|
|
|
def test_non_gated_config_oserror_is_preserved(loader_boundary):
|
|
loader.AutoConfig.from_pretrained.side_effect = OSError("offline cache miss")
|
|
with pytest.raises(OSError, match="offline cache miss"):
|
|
loader.load_model("x", local_files_only=True)
|
|
|
|
|
|
@pytest.mark.parametrize("exc", [ValueError("unknown"), KeyError("model_type")])
|
|
def test_malformed_or_unknown_config_has_stable_guidance(loader_boundary, exc):
|
|
loader.AutoConfig.from_pretrained.side_effect = exc
|
|
with pytest.raises(RuntimeError, match=r"(?s)not recognized by transformers.*trust_remote_code=True"):
|
|
loader.load_model("new/model")
|
|
|
|
|
|
def test_task_model_selection_and_gemma_contract(monkeypatch):
|
|
causal = object()
|
|
classification = object()
|
|
image_text = object()
|
|
monkeypatch.setitem(loader.TASK_MODEL_MAP, "causal_lm", causal)
|
|
monkeypatch.setitem(loader.TASK_MODEL_MAP, "classification", classification)
|
|
monkeypatch.setattr(loader, "AutoModelForImageTextToText", image_text)
|
|
assert loader._select_model_class("causal_lm", _config()) is causal
|
|
assert loader._select_model_class("classification", _config()) is classification
|
|
assert loader._select_model_class("causal_lm", _config(model_type="gemma4")) is image_text
|
|
assert loader._select_model_class(
|
|
"causal_lm",
|
|
_config(model_type="unknown", architectures=["Gemma4ForConditionalGeneration"]),
|
|
) is image_text
|
|
monkeypatch.setattr(loader, "AutoModelForImageTextToText", None)
|
|
with pytest.raises(RuntimeError, match="Upgrade transformers"):
|
|
loader._select_model_class("causal_lm", _config(model_type="gemma4"))
|
|
with pytest.raises(ValueError, match="Unknown task"):
|
|
loader._select_model_class("embedding", _config())
|
|
|
|
|
|
def test_model_handle_metadata_snapshot_restore_summary_and_cleanup(tmp_path):
|
|
model = _model()
|
|
nested = SimpleNamespace(
|
|
num_hidden_layers=3,
|
|
num_attention_heads=6,
|
|
hidden_size=12,
|
|
intermediate_size=24,
|
|
)
|
|
config = _config(
|
|
num_hidden_layers=0,
|
|
num_attention_heads=0,
|
|
hidden_size=0,
|
|
intermediate_size=0,
|
|
text_config=nested,
|
|
)
|
|
offload = tmp_path / "offload"
|
|
offload.mkdir()
|
|
(offload / "weight").write_text("x")
|
|
handle = loader.ModelHandle(
|
|
model,
|
|
SimpleNamespace(),
|
|
config,
|
|
"x",
|
|
"causal_lm",
|
|
_offload_dir=str(offload),
|
|
_owns_offload_dir=True,
|
|
)
|
|
assert (handle.num_layers, handle.num_heads, handle.hidden_size, handle.intermediate_size) == (3, 6, 12, 24)
|
|
with pytest.raises(RuntimeError, match="call .snapshot"):
|
|
handle.restore()
|
|
handle.snapshot()
|
|
model.load_state_dict.return_value = SimpleNamespace(missing_keys=[], unexpected_keys=[])
|
|
handle.restore()
|
|
model.load_state_dict.assert_called_once()
|
|
restored = model.load_state_dict.call_args.args[0]
|
|
assert torch.equal(restored["weight"], torch.ones(2))
|
|
assert model.load_state_dict.call_args.kwargs == {"strict": False}
|
|
assert handle.summary() == {
|
|
"model_name": "x",
|
|
"architecture": "gpt2",
|
|
"task": "causal_lm",
|
|
"num_layers": 3,
|
|
"num_heads": 6,
|
|
"hidden_size": 12,
|
|
"intermediate_size": 24,
|
|
"total_params": 5,
|
|
}
|
|
handle.cleanup()
|
|
assert not offload.exists()
|
|
assert handle._offload_dir is None
|
|
|
|
|
|
def test_model_handle_restore_tolerates_quantization_metadata():
|
|
model = _model()
|
|
handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
|
|
handle._original_state = {"weight": torch.ones(2)}
|
|
model.load_state_dict.return_value = SimpleNamespace(
|
|
missing_keys=["layer.weight.absmax"],
|
|
unexpected_keys=["layer.weight.bitsandbytes__nf4"],
|
|
)
|
|
handle.restore()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("missing", "unexpected", "message"),
|
|
[
|
|
(["layer.bias"], [], "missing keys: \\['layer.bias'\\]"),
|
|
([], ["layer.other_weight"], "unexpected keys: \\['layer.other_weight'\\]"),
|
|
(
|
|
["layer.bias"],
|
|
["layer.other_weight"],
|
|
r"missing keys: \['layer.bias'\].*unexpected keys: \['layer.other_weight'\]",
|
|
),
|
|
],
|
|
)
|
|
def test_model_handle_restore_rejects_parameter_mismatches(missing, unexpected, message):
|
|
model = _model()
|
|
handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
|
|
handle._original_state = {"weight": torch.ones(2)}
|
|
model.load_state_dict.return_value = SimpleNamespace(
|
|
missing_keys=missing,
|
|
unexpected_keys=unexpected,
|
|
)
|
|
with pytest.raises(RuntimeError, match=message):
|
|
handle.restore()
|
|
|
|
|
|
def test_model_handle_restore_preserves_loader_shape_failures():
|
|
model = _model()
|
|
handle = loader.ModelHandle(model, SimpleNamespace(), _config(), "x", "causal_lm")
|
|
handle._original_state = {"weight": torch.ones(2)}
|
|
model.load_state_dict.side_effect = RuntimeError("size mismatch for weight")
|
|
with pytest.raises(RuntimeError, match="size mismatch for weight"):
|
|
handle.restore()
|
|
|
|
|
|
def test_model_memory_estimation_handles_dense_moe_nested_and_unknown():
|
|
dense = loader._estimate_model_memory_gb(_config(), torch.float32)
|
|
moe = loader._estimate_model_memory_gb(_config(num_local_experts=4), torch.float32)
|
|
assert dense > 0
|
|
assert moe > dense
|
|
nested = _config(hidden_size=0, num_hidden_layers=0, intermediate_size=0, vocab_size=0)
|
|
nested.text_config = _config(hidden_size=8, num_hidden_layers=2, intermediate_size=16, vocab_size=32)
|
|
assert loader._estimate_model_memory_gb(nested, torch.float16) > 0
|
|
assert loader._estimate_model_memory_gb(_config(hidden_size=0), torch.float16) == 0
|
|
|
|
|
|
@pytest.mark.parametrize("quantization", ["4bit", "8bit"])
|
|
def test_quantization_rejects_non_cuda_instead_of_silent_degradation(loader_boundary, quantization):
|
|
with pytest.raises(RuntimeError, match="requires an available NVIDIA CUDA device"):
|
|
loader.load_model("x", quantization=quantization)
|
|
loader_boundary.model_class.from_pretrained.assert_not_called()
|
|
|
|
|
|
def test_quantization_requires_bitsandbytes(loader_boundary, monkeypatch):
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
|
|
monkeypatch.setattr(loader.dev, "supports_bitsandbytes", lambda _device=None: True)
|
|
real_import = builtins.__import__
|
|
|
|
def reject_bitsandbytes(name, *args, **kwargs):
|
|
if name == "bitsandbytes":
|
|
raise ImportError
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", reject_bitsandbytes)
|
|
with pytest.raises(RuntimeError, match="pip install -U bitsandbytes"):
|
|
loader.load_model("x", quantization="4bit")
|
|
|
|
|
|
@pytest.mark.parametrize("quantization", ["4bit", "8bit"])
|
|
def test_cuda_quantization_builds_explicit_config(loader_boundary, monkeypatch, quantization):
|
|
made = []
|
|
|
|
def fake_bnb(**kwargs):
|
|
made.append(kwargs)
|
|
return kwargs
|
|
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
|
|
monkeypatch.setattr(loader.dev, "supports_bitsandbytes", lambda _device=None: True)
|
|
monkeypatch.setattr("transformers.BitsAndBytesConfig", fake_bnb)
|
|
handle = loader.load_model("x", quantization=quantization, skip_snapshot=True)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert kwargs["device_map"] == "auto"
|
|
assert kwargs["quantization_config"] == made[0]
|
|
assert made[0][f"load_in_{quantization}"] is True
|
|
assert made[0]["llm_int8_enable_fp32_cpu_offload"] is True
|
|
handle.cleanup()
|
|
|
|
|
|
def test_native_quantization_wins_and_skips_dtype(loader_boundary):
|
|
loader_boundary.config.quantization_config = SimpleNamespace()
|
|
handle = loader.load_model("x", quantization="4bit", skip_snapshot=True)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert "torch_dtype" not in kwargs
|
|
assert kwargs["device_map"] == "auto"
|
|
assert "quantization_config" not in kwargs
|
|
handle.cleanup()
|
|
|
|
|
|
def test_cuda_auto_map_has_bounded_memory_and_requested_offload(loader_boundary, monkeypatch, tmp_path):
|
|
gib = 1024**3
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
|
|
monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: True)
|
|
monkeypatch.setattr(loader.dev, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(loader.dev, "device_count", lambda: 2)
|
|
monkeypatch.setattr(loader.dev, "_system_memory_gb", lambda: (64.0, 40.0))
|
|
monkeypatch.setattr(
|
|
loader.torch.cuda,
|
|
"get_device_properties",
|
|
lambda _index: SimpleNamespace(total_memory=20 * gib),
|
|
)
|
|
handle = loader.load_model("x", offload_folder=str(tmp_path), skip_snapshot=True)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert kwargs["offload_folder"] == str(tmp_path)
|
|
assert kwargs["max_memory"] == {0: "17408MiB", 1: "17408MiB", "cpu": "54GiB"}
|
|
handle.cleanup()
|
|
assert tmp_path.exists(), "cleanup must not remove an operator-owned directory"
|
|
|
|
|
|
def test_explicit_gpu_memory_utilization_overrides_default_headroom(
|
|
loader_boundary, monkeypatch,
|
|
):
|
|
gib = 1024**3
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
|
|
monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: True)
|
|
monkeypatch.setattr(loader.dev, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(loader.dev, "device_count", lambda: 1)
|
|
monkeypatch.setattr(loader.dev, "_system_memory_gb", lambda: (64.0, 40.0))
|
|
monkeypatch.setattr(
|
|
loader.torch.cuda,
|
|
"get_device_properties",
|
|
lambda _index: SimpleNamespace(total_memory=20 * gib),
|
|
)
|
|
handle = loader.load_model(
|
|
"x", gpu_memory_utilization=0.95, skip_snapshot=True,
|
|
)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert kwargs["max_memory"] == {0: "19456MiB", "cpu": "54GiB"}
|
|
handle.cleanup()
|
|
|
|
|
|
@pytest.mark.parametrize("device", ["cpu", "mps"])
|
|
def test_gpu_memory_utilization_is_safe_when_accelerate_auto_map_is_unavailable(
|
|
loader_boundary, monkeypatch, device,
|
|
):
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": device)
|
|
handle = loader.load_model(
|
|
"x", device=device, gpu_memory_utilization=0.95, skip_snapshot=True,
|
|
)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert "device_map" not in kwargs
|
|
assert "max_memory" not in kwargs
|
|
loader_boundary.model.to.assert_called_once_with(device)
|
|
handle.cleanup()
|
|
|
|
|
|
def _enable_cuda_quantization(monkeypatch, estimate_gb: float) -> None:
|
|
gib = 1024**3
|
|
monkeypatch.setattr(loader, "_estimate_model_memory_gb", lambda *_args: estimate_gb)
|
|
monkeypatch.setattr(loader.dev, "get_device", lambda _preference="auto": "cuda")
|
|
monkeypatch.setattr(loader.dev, "supports_bitsandbytes", lambda _device=None: True)
|
|
monkeypatch.setattr(loader.dev, "supports_device_map_auto", lambda _device=None: True)
|
|
monkeypatch.setattr(loader.dev, "get_total_free_gb", lambda: 16.0)
|
|
monkeypatch.setattr(loader.dev, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(loader.dev, "device_count", lambda: 1)
|
|
monkeypatch.setattr(loader.dev, "_system_memory_gb", lambda: (64.0, 40.0))
|
|
monkeypatch.setattr(
|
|
loader.torch.cuda,
|
|
"get_device_properties",
|
|
lambda _index: SimpleNamespace(total_memory=16 * gib),
|
|
)
|
|
monkeypatch.setattr("transformers.BitsAndBytesConfig", lambda **kwargs: kwargs)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("quantization", "estimate_gb", "expects_budget"),
|
|
[
|
|
("4bit", 40.0, False),
|
|
("4bit", 48.0, True),
|
|
("8bit", 20.0, False),
|
|
("8bit", 24.0, True),
|
|
],
|
|
)
|
|
def test_quantized_single_gpu_memory_budget_uses_effective_weight_size(
|
|
loader_boundary, monkeypatch, quantization, estimate_gb, expects_budget,
|
|
):
|
|
_enable_cuda_quantization(monkeypatch, estimate_gb)
|
|
handle = loader.load_model("x", quantization=quantization, skip_snapshot=True)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert ("max_memory" in kwargs) is expects_budget
|
|
if expects_budget:
|
|
assert kwargs["max_memory"] == {0: "13926MiB", "cpu": "54GiB"}
|
|
handle.cleanup()
|
|
|
|
|
|
def test_explicit_gpu_budget_overrides_quantized_fit_shortcut(loader_boundary, monkeypatch):
|
|
_enable_cuda_quantization(monkeypatch, estimate_gb=40.0)
|
|
handle = loader.load_model(
|
|
"x",
|
|
quantization="4bit",
|
|
gpu_memory_utilization=0.5,
|
|
skip_snapshot=True,
|
|
)
|
|
kwargs = loader_boundary.model_class.from_pretrained.call_args.kwargs
|
|
assert kwargs["max_memory"] == {0: "8192MiB", "cpu": "54GiB"}
|
|
handle.cleanup()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("quantization", "estimate_gb", "snapshots"),
|
|
[("4bit", 16.0, 1), ("4bit", 40.0, 0), ("8bit", 8.0, 1), ("8bit", 20.0, 0)],
|
|
)
|
|
def test_quantized_snapshot_policy_uses_effective_weight_size(
|
|
loader_boundary, monkeypatch, quantization, estimate_gb, snapshots,
|
|
):
|
|
_enable_cuda_quantization(monkeypatch, estimate_gb)
|
|
snapshot = Mock()
|
|
monkeypatch.setattr(loader.ModelHandle, "snapshot", snapshot)
|
|
handle = loader.load_model("x", quantization=quantization)
|
|
assert snapshot.call_count == snapshots
|
|
handle.cleanup()
|
|
|
|
|
|
def test_model_permission_error_retries_and_explicit_device_moves(loader_boundary, monkeypatch, tmp_path):
|
|
monkeypatch.setattr(loader.tempfile, "gettempdir", lambda: str(tmp_path))
|
|
loader_boundary.model_class.from_pretrained.side_effect = [PermissionError("cache"), loader_boundary.model]
|
|
loader.load_model("x", device="cpu", skip_snapshot=True)
|
|
assert loader_boundary.model_class.from_pretrained.call_count == 2
|
|
assert loader_boundary.model_class.from_pretrained.call_args.kwargs["cache_dir"].endswith("hf_home/hub")
|
|
loader_boundary.model.to.assert_called_once_with("cpu")
|
|
|
|
|
|
def test_model_gated_and_unknown_architecture_errors_are_actionable(loader_boundary):
|
|
loader_boundary.model_class.from_pretrained.side_effect = OSError("gated repo")
|
|
with pytest.raises(RuntimeError, match="Accept the license"):
|
|
loader.load_model("owner/gated")
|
|
|
|
loader_boundary.model_class.from_pretrained.side_effect = ValueError("does not recognize this architecture")
|
|
with pytest.raises(RuntimeError, match="pip install --upgrade transformers"):
|
|
loader.load_model("new/model")
|
|
|
|
loader_boundary.model_class.from_pretrained.side_effect = KeyError("unrelated")
|
|
with pytest.raises(KeyError, match="unrelated"):
|
|
loader.load_model("x")
|
|
|
|
|
|
def test_tokenizer_permission_retry_and_forced_snapshot(loader_boundary, monkeypatch, tmp_path):
|
|
monkeypatch.setattr(loader.tempfile, "gettempdir", lambda: str(tmp_path))
|
|
loader.AutoTokenizer.from_pretrained.side_effect = [PermissionError("cache"), loader_boundary.tokenizer]
|
|
handle = loader.load_model("x", skip_snapshot=False)
|
|
assert loader.AutoTokenizer.from_pretrained.call_count == 2
|
|
assert loader.AutoTokenizer.from_pretrained.call_args.kwargs["cache_dir"].endswith("hf_home/hub")
|
|
assert handle._original_state is not None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("native", "initial_free", "remaining_free", "snapshots"),
|
|
[
|
|
(True, 10.0, 3.0, 0),
|
|
(True, 10.0, 6.0, 1),
|
|
(False, 0.000001, 0.0, 0),
|
|
(False, 0.0, 0.0, 1),
|
|
],
|
|
)
|
|
def test_automatic_snapshot_memory_policy(
|
|
loader_boundary, monkeypatch, native, initial_free, remaining_free, snapshots,
|
|
):
|
|
if native:
|
|
loader_boundary.config.quantization_config = SimpleNamespace()
|
|
monkeypatch.setattr(loader.dev, "get_total_free_gb", Mock(side_effect=[initial_free, remaining_free]))
|
|
snapshot = Mock()
|
|
monkeypatch.setattr(loader.ModelHandle, "snapshot", snapshot)
|
|
handle = loader.load_model("x")
|
|
assert snapshot.call_count == snapshots
|
|
handle.cleanup()
|
|
|
|
|
|
def test_classification_load_sets_requested_label_count(loader_boundary):
|
|
handle = loader.load_model(
|
|
"x",
|
|
task="classification",
|
|
num_labels=7,
|
|
skip_snapshot=True,
|
|
)
|
|
assert loader_boundary.config.num_labels == 7
|
|
assert loader_boundary.model_class.from_pretrained.call_args.kwargs["config"] is loader_boundary.config
|
|
handle.cleanup()
|
|
|
|
|
|
def test_non_gated_model_oserror_is_preserved(loader_boundary):
|
|
loader_boundary.model_class.from_pretrained.side_effect = OSError("offline weight cache miss")
|
|
with pytest.raises(OSError, match="offline weight cache miss"):
|
|
loader.load_model("x", local_files_only=True)
|
|
|
|
|
|
def test_cached_path_compatibility_resolves_local_hub_and_fallback(monkeypatch, tmp_path):
|
|
local = tmp_path / "weights.bin"
|
|
local.write_bytes(b"weights")
|
|
assert loader._cached_path_shim(local) == str(local)
|
|
|
|
download = Mock(return_value="/cache/config.json")
|
|
monkeypatch.setattr("huggingface_hub.hf_hub_download", download)
|
|
assert loader._cached_path_shim("owner/model/config.json", cache_dir="/cache") == "/cache/config.json"
|
|
download.assert_called_once_with(
|
|
repo_id="owner/model",
|
|
filename="config.json",
|
|
cache_dir="/cache",
|
|
)
|
|
|
|
download.side_effect = RuntimeError("network unavailable")
|
|
assert loader._cached_path_shim("owner/model/config.json") == "owner/model/config.json"
|
|
download.reset_mock(side_effect=True)
|
|
assert loader._cached_path_shim("single-name") == "single-name"
|
|
download.assert_not_called()
|
|
|
|
|
|
def test_deferred_transformers_shims_are_additive_and_idempotent(monkeypatch):
|
|
fake_transformers = SimpleNamespace(_objects={})
|
|
monkeypatch.setitem(loader._sys.modules, "transformers", fake_transformers)
|
|
monkeypatch.setattr(loader, "_DEFERRED_SHIMS_APPLIED", False)
|
|
loader._apply_deferred_shims()
|
|
|
|
assert fake_transformers.LogitsWarper is fake_transformers._objects["LogitsWarper"]
|
|
assert fake_transformers.is_tf_available() is False
|
|
assert fake_transformers.is_flax_available() is False
|
|
assert fake_transformers.is_safetensors_available() is True
|
|
before = dict(fake_transformers.__dict__)
|
|
loader._apply_deferred_shims()
|
|
assert fake_transformers.__dict__ == before
|
|
|
|
|
|
def test_deferred_transformers_shims_tolerate_missing_module(monkeypatch):
|
|
monkeypatch.delitem(loader._sys.modules, "transformers", raising=False)
|
|
monkeypatch.setattr(loader, "_DEFERRED_SHIMS_APPLIED", False)
|
|
loader._apply_deferred_shims()
|
|
assert loader._DEFERRED_SHIMS_APPLIED is True
|
|
|
|
|
|
def test_deferred_transformers_shims_isolate_lazy_module_failures(monkeypatch):
|
|
class BrokenLazyModule:
|
|
def __getattribute__(self, name):
|
|
if name in {"LogitsWarper", "is_tf_available"}:
|
|
raise RuntimeError("lazy import failed")
|
|
return object.__getattribute__(self, name)
|
|
|
|
monkeypatch.setitem(loader._sys.modules, "transformers", BrokenLazyModule())
|
|
monkeypatch.setattr(loader, "_DEFERRED_SHIMS_APPLIED", False)
|
|
loader._apply_deferred_shims()
|
|
assert loader._DEFERRED_SHIMS_APPLIED is True
|
|
|
|
|
|
def test_image_text_input_compatibility_shim_preserves_argument_order():
|
|
images = object()
|
|
text = object()
|
|
assert loader._validate_images_text_input_order(
|
|
images=images,
|
|
text=text,
|
|
ignored="compatibility",
|
|
) == (images, text)
|
|
|
|
|
|
def test_cleanup_is_best_effort_when_directory_removal_fails(monkeypatch, tmp_path):
|
|
offload = tmp_path / "offload"
|
|
offload.mkdir()
|
|
handle = loader.ModelHandle(
|
|
_model(),
|
|
SimpleNamespace(),
|
|
_config(),
|
|
"x",
|
|
"causal_lm",
|
|
_offload_dir=str(offload),
|
|
_owns_offload_dir=True,
|
|
)
|
|
monkeypatch.setattr("shutil.rmtree", Mock(side_effect=OSError("busy")))
|
|
handle.cleanup()
|
|
assert handle._offload_dir is None
|
|
assert handle._owns_offload_dir is False
|
|
assert offload.exists()
|