"""Tests for the local in-process backend — catalogue, gate, registry, router dispatch. The ``local`` backend (ADR-0033) runs a small ``transformers`` model in-process on the host GPU behind ``@spaces.GPU`` — hardware-agnostic (ZeroGPU or a dedicated GPU), with no HTTP endpoint. These tests cover the deterministic, offline-safe surface: the catalogue data, the capability gate (env signals + an injectable CUDA probe), the unified-registry integration, and that the router dispatches a ``local:`` key to the in-process provider rather than the HTTP gateway. The actual GPU forward pass is integration-only (it needs a GPU and weights), exactly as the HTTP provider's live call is — so nothing here downloads a model or touches CUDA. """ from __future__ import annotations import os import pytest from src.models import inference, local_catalogue from src.models.local_provider import LocalTransformersProvider from src.models.router import ModelRouter # ── catalogue ───────────────────────────────────────────────────────────────────── def test_one_sponsor_model_per_tier_and_sizes_stay_small(): # Each tier maps to a *distinct* sponsor model (the multi-track cast), so one show spans # NVIDIA · OpenBMB · Cohere · JetBrains. Every model honours the ≤32B rule and the tiny # default keeps the Tiny-Titan ≤4B band. tagged = {m.profile: m for m in local_catalogue.LOCAL_MODELS if m.profile is not None} assert set(tagged) == {"tiny", "fast", "balanced", "strong"} assert all(m.params_b is None or m.params_b <= 32 for m in local_catalogue.LOCAL_MODELS) assert tagged["tiny"].params_b <= 4 # Tiny-Titan band assert len({m.source for m in tagged.values()}) == 4 # four sponsor families def test_every_tier_resolves_to_its_sponsor_model(): assert local_catalogue.default_key_for_profile("tiny") == "nvidia/Nemotron-Mini-4B-Instruct" # OpenBMB lane uses MiniCPM5 (native llama arch) — the MiniCPM 4.x custom code mis-computes # under the transformers 5.x floor (KV-cache crash / gibberish), so it is deliberately avoided. assert local_catalogue.default_key_for_profile("fast") == "openbmb/MiniCPM5-1B" assert local_catalogue.default_key_for_profile("balanced") == "CohereLabs/aya-expanse-8b" assert local_catalogue.default_key_for_profile("strong") == "JetBrains/Mellum2-12B-A2.5B-Instruct" # the tiny model is listed first, so an untagged/unknown tier falls back to the cheapest. assert local_catalogue.LOCAL_MODELS[0].profile == "tiny" def test_catalogue_cast_is_all_native_arch_and_field_still_plumbs_custom_code(): # The whole live cast loads with the stock AutoModelForCausalLM — no trust_remote_code, # no custom-code-only KV-cache workaround — which is why output is correct under # transformers 5.x (MiniCPM5 replaced the 4.x custom-code model for exactly this reason). for m in local_catalogue.LOCAL_MODELS: assert m.trust_remote_code is False, f"{m.repo_id} unexpectedly needs trust_remote_code" assert m.use_cache is True, f"{m.repo_id} unexpectedly disables the KV cache" assert local_catalogue.model_by_key("openbmb/MiniCPM5-1B").trust_remote_code is False assert local_catalogue.model_by_key("does/not-exist") is None # The fields still plumb a non-default value, so a future custom-code model is one append away. custom = local_catalogue.LocalModel(repo_id="acme/custom", trust_remote_code=True, use_cache=False) assert custom.trust_remote_code is True and custom.use_cache is False def test_binding_is_a_bare_repo_id_with_no_endpoint(): # In-process: the binding carries the raw transformers repo id (no openai/ prefix) and # neither a base_url nor an api_key — the router builds the in-process provider from it. binding = local_catalogue.binding_for("nvidia/Nemotron-Mini-4B-Instruct") assert binding["model"] == "nvidia/Nemotron-Mini-4B-Instruct" assert binding["base_url"] == "" assert binding["api_key"] == "" def test_binding_unknown_key_raises(): with pytest.raises(KeyError): local_catalogue.binding_for("nobody/here") # ── capability gate ───────────────────────────────────────────────────────────────── def test_gate_explicit_env_is_deterministic_without_a_probe(): # An explicit env dict is the whole story — no torch import, no host probe. assert local_catalogue.has_credentials(env={}) is False assert local_catalogue.has_credentials(env={"SPACES_ZERO_GPU": "true"}) is True assert local_catalogue.has_credentials(env={"LOCAL_INFERENCE": "1"}) is True def test_gate_accepts_common_truthy_spellings(): for val in ("1", "true", "TRUE", "yes", "on"): assert local_catalogue.has_credentials(env={"LOCAL_INFERENCE": val}) is True for val in ("0", "false", "", "no"): assert local_catalogue.has_credentials(env={"LOCAL_INFERENCE": val}) is False def test_gate_uses_injected_cuda_probe_when_env_signals_absent(): # No env signal → fall through to the probe (auto-detect a dedicated GPU / local box). assert local_catalogue.has_credentials(env={}, cuda_probe=lambda: True) is True assert local_catalogue.has_credentials(env={}, cuda_probe=lambda: False) is False # An env signal short-circuits before the probe is ever consulted. assert local_catalogue.has_credentials(env={"SPACES_ZERO_GPU": "1"}, cuda_probe=lambda: False) is True def test_gate_auto_probes_only_against_the_real_environment(): # Passing os.environ itself opts into the host CUDA probe; an arbitrary dict does not, # keeping façade/test calls deterministic. We assert the boolean, whatever the host is. assert isinstance(local_catalogue.has_credentials(env=os.environ), bool) # ── unified registry integration ───────────────────────────────────────────────────── def test_local_backend_is_registered_and_qualified(): assert "local" in {b.key for b in inference.backends()} keys = {e["key"] for e in inference.entries("local")} assert keys and all(k.startswith("local:") for k in keys) def test_registry_default_and_binding_round_trip(): key = inference.default_key_for_profile("tiny", "local") assert key == "local:nvidia/Nemotron-Mini-4B-Instruct" binding = inference.binding_for(key) assert binding["model"] == "nvidia/Nemotron-Mini-4B-Instruct" assert binding["base_url"] == "" def test_backend_available_and_configured_backends_for_local(): assert inference.backend_available("local", env={"LOCAL_INFERENCE": "1"}) is True assert inference.backend_available("local", env={"SPACES_ZERO_GPU": "yes"}) is True assert inference.backend_available("local", env={}) is False configured = inference.configured_backends(env={"LOCAL_INFERENCE": "1"}) assert "local" in configured # ── router dispatch ────────────────────────────────────────────────────────────────── def test_router_dispatches_local_key_to_in_process_provider(): # A live router resolving a local: key must build the in-process provider (not LiteLLM), # bound to the bare repo id. Construction only — no GPU is touched. router = ModelRouter(offline=False) provider = router.for_profile("local:nvidia/Nemotron-Mini-4B-Instruct") assert isinstance(provider, LocalTransformersProvider) assert provider.model == "nvidia/Nemotron-Mini-4B-Instruct" assert provider.model_id == "nvidia/Nemotron-Mini-4B-Instruct" def test_catalogue_spec_tags_local_kind_and_others_litellm(): router = ModelRouter(offline=False) local_spec = router._catalogue_spec("local:nvidia/Nemotron-Mini-4B-Instruct") assert local_spec is not None and local_spec.kind == "local" # An HF key resolves through the same path but stays on the HTTP transport. hf_spec = router._catalogue_spec("hf:katanemo/Arch-Router-1.5B") assert hf_spec is not None and hf_spec.kind == "litellm" # ── provider (cheap, offline-safe surface) ─────────────────────────────────────────── def test_provider_reports_model_id_and_zeroed_usage_before_any_call(): provider = LocalTransformersProvider(model="nvidia/Nemotron-Mini-4B-Instruct") assert provider.model_id == "nvidia/Nemotron-Mini-4B-Instruct" assert provider.last_usage == {} # no call yet — matches the sibling providers provider._zero_usage() assert provider.last_usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} def test_provider_resolves_trust_remote_code_from_catalogue(): # The cast is all native-arch, so every catalogue model resolves to False; an # off-catalogue repo also defaults to the safe choice. assert LocalTransformersProvider(model="openbmb/MiniCPM5-1B")._trust_remote_code() is False assert LocalTransformersProvider(model="CohereLabs/aya-expanse-8b")._trust_remote_code() is False assert LocalTransformersProvider(model="some/random-repo")._trust_remote_code() is False def test_provider_resolves_use_cache_from_catalogue(): # Native-arch models keep the KV cache on (the fast path), and an off-catalogue repo # defaults to the cached path; no model in the current cast disables it. assert LocalTransformersProvider(model="openbmb/MiniCPM5-1B")._use_cache() is True assert LocalTransformersProvider(model="nvidia/Nemotron-Mini-4B-Instruct")._use_cache() is True assert LocalTransformersProvider(model="some/random-repo")._use_cache() is True # ── ZeroGPU contract: CUDA only inside @spaces.GPU, never in the parent ─────────────── # Regression guard for the production crash "Low-level CUDA init (torch._C._cuda_init) # reached … ZeroGPU's emulation did not intercept": the parent process gets no GPU, so any # CUDA placement (or a model load that places onto a device) outside the @spaces.GPU window # kills the worker. The forward pass can only be exercised with a GPU + weights # (integration), so we pin the *structural* invariant — where CUDA may be touched, and how # the model reaches the device — by source contract. def test_parent_warm_only_downloads_never_loads_or_initialises_cuda(): import ast import inspect from src.models import local_provider # _ensure_downloaded runs in the parent. It must only fetch weights to disk — never touch # CUDA, and never materialise the model in RAM (a multi-model cast would pin tens of GB). # Check the executable body with the docstring stripped (the docstring explains the # invariant in prose, so it legitimately mentions CUDA/RAM); the banned ops are device # moves, torch.cuda.* and any model instantiation. fn = ast.parse(inspect.getsource(local_provider._ensure_downloaded)).body[0] if ast.get_docstring(fn): fn.body = fn.body[1:] code = ast.unparse(fn) assert 'to("cuda")' not in code and "torch.cuda" not in code and ".cuda(" not in code # No weight materialisation in the parent — only the on-disk fetch. assert "from_pretrained" not in code and "AutoModel" not in code assert "snapshot_download" in code def test_worker_loads_onto_device_via_device_map_no_meta_prone_move(): import ast import inspect from src.models import local_provider # Regression guard for the ZeroGPU crash "Cannot copy out of meta tensor; no data!". # transformers 5.x always builds on the meta device and streams the checkpoint onto the # target; a bare from_pretrained(...).to("cuda") can leave a non-persistent buffer (e.g. # rotary inv_freq) or a tied/"missing" head on meta, and the move then dies # (transformers#41038/#30703). low_cpu_mem_usage no longer changes this (5.x drops the # kwarg). The fix: hand transformers the device via device_map so it materialises AND # places everything on-device in one step — no fragile post-hoc .to("cuda"). fn = ast.parse(inspect.getsource(local_provider._ensure_loaded_on_device)).body[0] if ast.get_docstring(fn): fn.body = fn.body[1:] code = ast.unparse(fn) # The supported placement path is used… assert "device_map" in code # …the GPU window never re-downloads (the parent already fetched the weights)… assert "local_files_only=True" in code # …and the meta-prone manual move / dead kwarg are gone. assert 'to("cuda")' not in code assert "low_cpu_mem_usage" not in code def test_v4_compat_shim_backfills_removed_remote_code_predicates(): # Regression guard for the ZeroGPU error "cannot import name 'is_torch_fx_available' # from transformers.utils.import_utils": transformers 5.x removed these predicates, but # MiniCPM's (and other) trust_remote_code modelling files still import them. The provider # back-fills them (all True at our torch floor) so the remote import succeeds. from src.models import local_provider local_provider._ensure_transformers_v4_symbols() from transformers.utils import import_utils # Every name the shim covers is importable from transformers.utils.import_utils and True. for name in local_provider._REMOVED_TORCH_PREDICATES: fn = getattr(import_utils, name) assert fn() is True # And the device loader runs the shim before touching any remote code. import inspect assert "_ensure_transformers_v4_symbols()" in inspect.getsource(local_provider._ensure_loaded_on_device) def test_device_placement_lives_inside_the_spaces_gpu_function(): from pathlib import Path from src.models import local_provider # _generate is wrapped by @spaces.GPU, so read the module source and isolate its block. module_src = Path(local_provider.__file__).read_text() gen_block = module_src.split("def _generate(", 1)[1].split("\ndef ", 1)[0] # The model reaches the device here (the one place ZeroGPU grants one) via the on-device # loader — never via a parent-side load… assert "_ensure_loaded_on_device(" in gen_block # …and the function carries the decorator the platform registers. assert "@spaces.GPU" in module_src.split("def _generate(", 1)[0].rsplit("\n\n", 1)[-1] # The parent path (complete) warms the on-disk cache only — it must not load on-device. complete_block = module_src.split("def complete(", 1)[1].split("\n def ", 1)[0] assert "_ensure_downloaded(" in complete_block assert "_ensure_loaded_on_device(" not in complete_block def test_generate_unpacks_batchencoding_never_passes_a_positional_dict(): # Regression guard for the production AttributeError "inputs_tensor.shape[0]" in # transformers.generate: in transformers 5.x apply_chat_template(return_tensors="pt") # defaults to a BatchEncoding *dict*, and passing that dict positionally into # model.generate(inputs) makes generate() do .shape on a dict. The fix: request the # dict explicitly (return_dict=True) and unpack it with ** so input_ids + attention_mask # are fed as kwargs. Pinned by AST so the call shape can't silently regress. import ast from pathlib import Path from src.models import local_provider tree = ast.parse(Path(local_provider.__file__).read_text()) gen = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "_generate") calls = [c for c in ast.walk(gen) if isinstance(c, ast.Call)] # apply_chat_template asks for the dict form explicitly (robust whatever the default). act = next(c for c in calls if isinstance(c.func, ast.Attribute) and c.func.attr == "apply_chat_template") assert any(k.arg == "return_dict" and k.value.value is True for k in act.keywords) # Reasoning models (e.g. MiniCPM5) are told not to think, so a block can't eat the # token budget and leave an empty spoken line; harmlessly ignored by non-reasoning templates. assert any(k.arg == "enable_thinking" and k.value.value is False for k in act.keywords) # model.generate(**inputs, …): the encoding is unpacked, never a positional dict. gen_call = next(c for c in calls if isinstance(c.func, ast.Attribute) and c.func.attr == "generate") assert not gen_call.args, "generate() must take no positional arg (the old bug passed the dict positionally)" assert any(k.arg is None and isinstance(k.value, ast.Name) and k.value.id == "inputs" for k in gen_call.keywords) # use_cache is threaded through so a model with broken 5.x cache handling (MiniCPM) can # disable it ("Key and Value must have the same sequence length"). assert any(k.arg == "use_cache" for k in gen_call.keywords)