diff --git a/packages/sie_sdk/src/sie_sdk/cache.py b/packages/sie_sdk/src/sie_sdk/cache.py index 53bb81f1f..5e4849167 100644 --- a/packages/sie_sdk/src/sie_sdk/cache.py +++ b/packages/sie_sdk/src/sie_sdk/cache.py @@ -98,6 +98,8 @@ def is_model_cached( model_id: str, config: CacheConfig | None = None, revision: str | None = None, + *, + required_files: list[str] | None = None, ) -> bool: """Check if a model is already in local cache. @@ -108,6 +110,8 @@ def is_model_cached( config: Cache configuration. If None, reads from environment. revision: Exact snapshot commit, branch, or tag to require. Branches and tags resolve through the local Hugging Face ``refs`` cache. + required_files: Safe relative files required for a cache hit. By + default, any file preserves the existing model-cache semantics. Returns: True if model appears to be cached locally. @@ -115,6 +119,15 @@ def is_model_cached( if config is None: config = get_cache_config() + if required_files is not None: + for name in required_files: + if ( + not isinstance(name, str) + or any(character in name for character in ("\\", "\x00", ":")) + or any(part in {"", ".", ".."} for part in name.split("/")) + ): + raise ValueError("Required cache files must be safe relative paths") + # HF Hub cache structure: models--{org}--{model}/snapshots/{revision}/ cache_dir = config.local_cache / f"models--{model_id.replace('/', '--')}" if not cache_dir.exists(): @@ -143,15 +156,27 @@ def is_model_cached( ): return False snapshot = snapshots_dir / resolved_revision - return snapshot.is_dir() and any(snapshot.iterdir()) + return ( + snapshot.is_dir() + and any(snapshot.iterdir()) + and (required_files is None or all((snapshot / name).is_file() for name in required_files)) + ) - return any(snapshot.is_dir() and any(snapshot.iterdir()) for snapshot in snapshots_dir.iterdir()) + return any( + snapshot.is_dir() + and any(snapshot.iterdir()) + and (required_files is None or all((snapshot / name).is_file() for name in required_files)) + for snapshot in snapshots_dir.iterdir() + ) def ensure_model_cached( model_id: str, config: CacheConfig | None = None, revision: str | None = None, + *, + allow_patterns: list[str] | None = None, + required_files: list[str] | None = None, ) -> Path: """Ensure model weights are in local cache. @@ -164,6 +189,10 @@ def ensure_model_cached( model_id: HuggingFace model ID (e.g., "BAAI/bge-m3"). config: Cache configuration. If None, reads from environment. revision: Specific model revision (commit hash, branch, or tag). + allow_patterns: Optional file filter for Hub downloads only. Local and + cluster cache lookup and population remain unchanged. + required_files: Files required for a local or cluster hit. A partial + snapshot continues through the existing fallback tiers. Returns: Path to cached model in local cache. @@ -176,20 +205,22 @@ def ensure_model_cached( config = get_cache_config() # Check local cache first - if is_model_cached(model_id, config, revision=revision): + if is_model_cached(model_id, config, revision=revision, required_files=required_files): logger.debug("Model %s found in local cache", model_id) return _get_model_cache_path(model_id, config) # Try cluster cache if configured if config.cluster_cache: - if _download_from_cluster_cache(model_id, config) and is_model_cached(model_id, config, revision=revision): + if _download_from_cluster_cache(model_id, config) and is_model_cached( + model_id, config, revision=revision, required_files=required_files + ): logger.info("Downloaded %s from cluster cache", model_id) return _get_model_cache_path(model_id, config) # Download from HuggingFace Hub if fallback enabled if config.hf_fallback: logger.info("Downloading %s from HuggingFace Hub", model_id) - return _download_from_huggingface(model_id, config, revision) + return _download_from_huggingface(model_id, config, revision, allow_patterns=allow_patterns) # Not in any cache and HF fallback disabled raise RuntimeError( @@ -344,6 +375,8 @@ def _download_from_huggingface( model_id: str, config: CacheConfig, revision: str | None = None, + *, + allow_patterns: list[str] | None = None, ) -> Path: """Download model from HuggingFace Hub to local cache. @@ -354,6 +387,7 @@ def _download_from_huggingface( model_id: HuggingFace model ID. config: Cache configuration. revision: Specific model revision. + allow_patterns: Optional file filter passed to snapshot_download. Returns: Path to the model in local cache. @@ -379,6 +413,7 @@ def _download_from_huggingface( revision=revision, token=token, cache_dir=str(config.local_cache), # Points to HF_HOME/hub + allow_patterns=allow_patterns, ignore_patterns=_get_hf_ignore_patterns(model_id, revision, token), ) return _get_model_cache_path(model_id, config) diff --git a/packages/sie_sdk/tests/test_cache.py b/packages/sie_sdk/tests/test_cache.py index 582bc90ca..073c85e3e 100644 --- a/packages/sie_sdk/tests/test_cache.py +++ b/packages/sie_sdk/tests/test_cache.py @@ -778,3 +778,73 @@ def test_cache_dir_uses_local_cache_path(self, tmp_path: Path) -> None: mock_download.assert_called_once() call_kwargs = mock_download.call_args[1] assert call_kwargs["cache_dir"] == str(tmp_path) + + +@pytest.mark.parametrize("allow_patterns", [None, ["adapter_config.json", "adapter_model.safetensors"]]) +def test_hub_download_pattern_filter_is_optional_and_forwarded( + tmp_path: Path, allow_patterns: list[str] | None +) -> None: + config = CacheConfig(local_cache=tmp_path, hf_fallback=True) + revision = "a" * 40 + with patch("huggingface_hub.snapshot_download") as download: + if allow_patterns is None: + result = ensure_model_cached("org/model", config, revision=revision) + else: + result = ensure_model_cached("org/model", config, revision=revision, allow_patterns=allow_patterns) + assert result == tmp_path / "models--org--model" + download.assert_called_once() + assert download.call_args.args == ("org/model",) + assert download.call_args.kwargs["revision"] == revision + assert download.call_args.kwargs["cache_dir"] == str(tmp_path) + assert download.call_args.kwargs.get("allow_patterns") == allow_patterns + + +@pytest.mark.parametrize("revision", [None, "a" * 40]) +@pytest.mark.parametrize("required_files", [None, [], ["adapter_config.json"], ["adapter_model.safetensors"]]) +def test_opt_in_cache_completeness_keeps_ordinary_semantics( + tmp_path: Path, revision: str | None, required_files: list[str] | None +) -> None: + directory = tmp_path / "models--org--model" / "snapshots" / ("a" * 40) + directory.mkdir(parents=True) + (directory / "adapter_config.json").write_bytes(b"{}") + config = CacheConfig(local_cache=tmp_path, hf_fallback=False) + + assert is_model_cached("org/model", config, revision=revision) is True + assert is_model_cached("org/model", config, revision=revision, required_files=required_files) is ( + required_files != ["adapter_model.safetensors"] + ) + + +@pytest.mark.parametrize("lookup", [is_model_cached, ensure_model_cached]) +@pytest.mark.parametrize( + "name", + [ + "", + ".", + "..", + "/absolute", + "../weights", + "nested/../../weights", + "nested//weights", + "nested\\weights", + "C:weights", + "file:stream", + "bad\x00name", + "./weights", + ], +) +def test_required_cache_files_refuse_unsafe_relative_names(tmp_path: Path, lookup, name: str) -> None: + config = CacheConfig(local_cache=tmp_path, cluster_cache="s3://must-not-contact", hf_fallback=True) + with ( + patch("sie_sdk.cache.get_storage_backend", side_effect=AssertionError("storage must not be contacted")), + patch("huggingface_hub.snapshot_download", side_effect=AssertionError("Hub must not be contacted")), + pytest.raises(ValueError, match="safe relative paths"), + ): + lookup("org/model", config, required_files=[name]) + + +def test_required_cache_files_support_safe_nested_files(tmp_path: Path) -> None: + directory = tmp_path / "models--org--model" / "snapshots" / ("a" * 40) + (directory / "nested").mkdir(parents=True) + (directory / "nested" / "config.json").write_bytes(b"{}") + assert is_model_cached("org/model", CacheConfig(local_cache=tmp_path), required_files=["nested/config.json"]) diff --git a/packages/sie_server/src/sie_server/adapters/_lora_source.py b/packages/sie_server/src/sie_server/adapters/_lora_source.py new file mode 100644 index 000000000..686d5e40f --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/_lora_source.py @@ -0,0 +1,95 @@ +"""Resolve admitted LoRA sources through the configured model-cache policy. + +This module neither declares nor authorizes adapters. Callers must resolve their +served alias and admit its source/revision before calling this load-time helper. +It has no PEFT, GLiNER or model-runtime dependency. +""" + +from pathlib import Path + +from huggingface_hub.utils import validate_repo_id +from sie_sdk.cache import ensure_model_cached, get_cache_config + +from sie_server.config.model import is_immutable_revision + +_FILES = ("adapter_config.json", "adapter_model.safetensors") + + +def _check_files(directory: Path, *, blobs: Path | None = None) -> Path: + for filename in _FILES: + path = directory / filename + resolved = path.resolve() + if not resolved.is_relative_to(directory) and (blobs is None or not resolved.is_relative_to(blobs)): + raise ValueError("LoRA artifact escapes its declared directory or repository blobs") + if not path.is_file(): + raise ValueError(f"LoRA source ships no {filename}") + return directory + + +def _check_repository(repository: Path) -> None: + # The configured cache root may intentionally be a symlink; it is already + # canonicalized. Derived repositories and layout directories may not alias + # another repository or snapshot. + for directory in (repository, *(repository / name for name in ("snapshots", "refs", "blobs"))): + if directory.resolve() != directory: + raise ValueError("LoRA cache repository directories cannot be aliases") + + +def resolve_lora_source(source: str, revision: str | None = None) -> Path: + """Return the admitted two-file adapter from local, cluster or allowed Hub storage. + + Explicit local directories (including a selected root symlink) are supported + without a Hub revision. Hub pins must be immutable commits. An unpinned Hub + source uses only its cached ``refs/main`` commit, never an arbitrary snapshot. + Standard snapshot-file links to this repository's blobs remain supported. + Missing files, redirected paths or a missing exact revision fail closed. + """ + directory = Path(source) + if directory.is_dir(): + if revision is not None: + raise ValueError("A local LoRA directory cannot honor a Hugging Face revision pin") + directory = directory.resolve() + repository = directory.parent.parent + if ( + directory.parent.name == "snapshots" + and is_immutable_revision(directory.name) + and repository.name.startswith("models--") + ): + _check_repository(repository) + return _check_files(directory, blobs=repository / "blobs") + return _check_files(directory) + + validate_repo_id(source) + if revision is not None and not is_immutable_revision(revision): + raise ValueError("LoRA revision must be an immutable 40-hex commit SHA") + config = get_cache_config() + repository = config.local_cache.resolve() / f"models--{source.replace('/', '--')}" + _check_repository(repository) + if revision is None: + reference = repository / "refs" / "main" + if reference.is_symlink(): + raise ValueError("LoRA default revision cannot be a file alias") + if reference.is_file() and not is_immutable_revision(reference.read_text().strip()): + raise ValueError("LoRA cached default revision must be an immutable commit SHA") + cached = ensure_model_cached( + source, + config, + revision=revision or "main", + allow_patterns=list(_FILES), + required_files=list(_FILES), + ) + if cached.resolve() != repository: + raise ValueError("LoRA cache returned a different repository") + _check_repository(repository) + + if revision is None: + reference = repository / "refs" / "main" + if reference.is_symlink(): + raise ValueError("LoRA default revision cannot be a file alias") + revision = reference.read_text().strip() + if not is_immutable_revision(revision): + raise ValueError("LoRA cached default revision must be an immutable commit SHA") + directory = repository / "snapshots" / revision + if directory.resolve() != directory: + raise ValueError("LoRA snapshot cannot alias another revision or repository") + return _check_files(directory, blobs=repository / "blobs") diff --git a/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py b/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py index e95818e46..1ef5fd446 100644 --- a/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py +++ b/packages/sie_server/src/sie_server/adapters/cross_encoder/_lora.py @@ -19,7 +19,6 @@ from typing import Any import torch -from huggingface_hub import snapshot_download from peft import LoraConfig, PeftConfig from peft.functional import delete_adapter from peft.tuners.lora.layer import Linear as LoraLinear @@ -27,6 +26,7 @@ from sentence_transformers import CrossEncoder from transformers import BertForSequenceClassification, XLMRobertaForSequenceClassification +from sie_server.adapters._lora_source import resolve_lora_source from sie_server.adapters.lora import LoraRuntimeConfig, LoraStateUncertainError, LoraUnavailableError from sie_server.adapters.peft_lora_mixin import PEFTLoRAMixin @@ -211,12 +211,7 @@ def fetch(self, name: str, revision: str | None = None) -> str: generation = self._generation if source in self._prepared: return self._prepared[source].directory.name - directory = Path(source) - if directory.is_dir(): - if pinned is not None: - raise ValueError("A local LoRA directory cannot honor a Hugging Face revision pin") - else: - directory = Path(snapshot_download(source, revision=pinned, allow_patterns=list(_FILES))) + directory = resolve_lora_source(source, pinned) temporary = tempfile.TemporaryDirectory(prefix="sie-cross-encoder-lora-") try: # A private snapshot prevents later source edits or alternative pickle files diff --git a/packages/sie_server/src/sie_server/adapters/peft_lora_mixin.py b/packages/sie_server/src/sie_server/adapters/peft_lora_mixin.py index 0745286fd..4283c6ba7 100644 --- a/packages/sie_server/src/sie_server/adapters/peft_lora_mixin.py +++ b/packages/sie_server/src/sie_server/adapters/peft_lora_mixin.py @@ -18,12 +18,12 @@ def load(self, device: str) -> None: import hashlib import logging import re -from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar -from huggingface_hub import snapshot_download from safetensors import safe_open +from sie_server.adapters._lora_source import resolve_lora_source + if TYPE_CHECKING: from collections.abc import Collection @@ -36,7 +36,6 @@ def load(self, device: str) -> None: # OLoRA, LoftQ and LoRA-GA instead rewrite the base layer's weights. _BASE_PRESERVING_INITIALIZERS = ("gaussian", "eva", "orthogonal") -_ADAPTER_CONFIG_FILE = "adapter_config.json" _ADAPTER_WEIGHTS_FILE = "adapter_model.safetensors" # The tensor names PEFT saves for a LoRA's own parameters. PEFT loads any @@ -219,9 +218,10 @@ def supports_hot_lora_reload(self) -> bool: def fetch_lora(self, lora_path: str, revision: str | None = None) -> str: """Resolve a LoRA to one local directory, and refuse it there unless it is safe to attach. - A Hub id is downloaded once, at one commit, with only its - ``adapter_config.json`` and ``adapter_model.safetensors``; a local - directory is used as it is. :meth:`load_lora` then reads both files + A Hub id uses the configured local/cluster cache and allowed Hub + fallback at one commit; a local directory is used as it is. The + required files are ``adapter_config.json`` and + ``adapter_model.safetensors``. :meth:`load_lora` reads both files from the returned directory. Network and disk I/O happen here, so a caller can run this outside inference. @@ -238,21 +238,13 @@ def fetch_lora(self, lora_path: str, revision: str | None = None) -> str: the adapter ships no ``adapter_model.safetensors``, or its weights file has a tensor that is not a LoRA parameter. """ - directory = Path(lora_path) - if not directory.is_dir(): - directory = Path( - snapshot_download( - lora_path, - revision=revision, - allow_patterns=[_ADAPTER_CONFIG_FILE, _ADAPTER_WEIGHTS_FILE], - ) - ) + directory = resolve_lora_source(lora_path, revision) try: from peft import PeftConfig except ImportError as e: msg = "PEFT is required for LoRA support. Install with: pip install peft" raise RuntimeError(msg) from e - self._check_lora_config(lora_path, PeftConfig.from_pretrained(str(directory))) + self._check_lora_config(lora_path, PeftConfig.from_pretrained(str(directory), local_files_only=True)) weights = directory / _ADAPTER_WEIGHTS_FILE if not weights.is_file(): msg = f"LoRA '{lora_path}' ships no {_ADAPTER_WEIGHTS_FILE}; adapter weights in other formats are refused." diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_lora.py b/packages/sie_server/src/sie_server/adapters/sglang/_lora.py index 27e387ece..f1fe8afb6 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_lora.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_lora.py @@ -14,7 +14,6 @@ from typing import Any, cast import requests -from huggingface_hub import snapshot_download from safetensors import safe_open from sie_server.adapters._generation_base import ( @@ -22,6 +21,7 @@ GenerationInvalidRequestError, GenerationLoraLoadingError, ) +from sie_server.adapters._lora_source import resolve_lora_source from sie_server.adapters.lora import LoraRuntimeConfig, LoraStateUncertainError, LoraUnavailableError from sie_server.adapters.peft_lora_mixin import is_lora_weight_key @@ -222,9 +222,7 @@ def fetch(self, name: str, revision: str | None = None) -> str: raise RuntimeError("The LoRA engine was unloaded during fetch") if source in self._snapshots: return self._snapshots[source].name - directory = Path(source) - if not directory.is_dir(): - directory = Path(snapshot_download(source, revision=declared_revision or revision, allow_patterns=_FILES)) + directory = resolve_lora_source(source, declared_revision or revision) config = json.loads((directory / _FILES[0]).read_text()) if not isinstance(config, dict) or config.get("peft_type") != "LORA": raise ValueError("SGLang requires a PEFT LoRA adapter") diff --git a/packages/sie_server/tests/adapters/test_cross_encoder_lora.py b/packages/sie_server/tests/adapters/test_cross_encoder_lora.py index 046a2f67f..161713d05 100644 --- a/packages/sie_server/tests/adapters/test_cross_encoder_lora.py +++ b/packages/sie_server/tests/adapters/test_cross_encoder_lora.py @@ -6,6 +6,7 @@ import builtins import gc import json +import socket import subprocess import sys import time @@ -13,6 +14,7 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock import httpx import numpy as np @@ -328,6 +330,37 @@ def test_real_loader_routes_only_the_explicit_profile_and_preserves_base_options load_adapter(ModelConfig.model_validate(invalid), checkpoint, device="cpu") +def test_pinned_cluster_adapter_preserves_lifecycle_with_network_denied(checkpoint, tmp_path, monkeypatch): + cluster = tmp_path / "cluster" + revision = "a" * 40 + cached = cluster / "models--org--adapter" / "snapshots" / revision + cached.parent.mkdir(parents=True) + save_adapter(checkpoint, cached) + monkeypatch.setenv("SIE_LOCAL_CACHE", str(tmp_path / "local")) + monkeypatch.setenv("SIE_CLUSTER_CACHE", str(cluster)) + monkeypatch.setenv("SIE_HF_FALLBACK", "false") + monkeypatch.setattr(socket.socket, "connect", Mock(side_effect=AssertionError("network forbidden"))) + monkeypatch.setattr("huggingface_hub.snapshot_download", Mock(side_effect=AssertionError("Hub forbidden"))) + adapter = CrossEncoderLoraAdapter( + checkpoint, lora_paths={"a": "org/adapter"}, lora_revisions={"org/adapter": revision}, max_loaded_loras=1 + ) + adapter.load("cpu") + try: + baseline = scores(adapter) + original = base_snapshot(adapter) + prepared = adapter.fetch_lora("a") + adapter.load_lora("a", source=prepared) + adapter.set_active_lora("a") + assert not np.array_equal(scores(adapter), baseline) + adapter.set_active_lora(None) + np.testing.assert_array_equal(scores(adapter), baseline) + adapter.unload_lora("a") + assert_base_preserved(adapter, original, unwrapped=True) + finally: + adapter.unload() + assert not Path(prepared).exists() + + def test_actual_effect_restoration_head_and_final_module_identity(adapter, checkpoint): snapshot = base_snapshot(adapter) baseline = scores(adapter) diff --git a/packages/sie_server/tests/adapters/test_lora.py b/packages/sie_server/tests/adapters/test_lora.py index f42459a3a..135aa7677 100644 --- a/packages/sie_server/tests/adapters/test_lora.py +++ b/packages/sie_server/tests/adapters/test_lora.py @@ -16,6 +16,7 @@ import asyncio import json +import socket import threading import time from collections import OrderedDict @@ -1435,13 +1436,13 @@ async def test_a_hub_lora_is_fetched_once_at_its_revision_and_loaded_from_that_s self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: snapshot = _save_plain_lora(tmp_path / "snapshot") - calls: list[tuple[str, dict[str, Any]]] = [] + calls: list[tuple[str, str | None]] = [] - def snapshot_download(repo_id: str, **kwargs: Any) -> str: - calls.append((repo_id, kwargs)) - return snapshot + def resolve_source(repo_id: str, revision: str | None = None) -> Path: + calls.append((repo_id, revision)) + return Path(snapshot) - monkeypatch.setattr(peft_lora_mixin, "snapshot_download", snapshot_download) + monkeypatch.setattr(peft_lora_mixin, "resolve_lora_source", resolve_source) adapter = TinyLoRAAdapter() adapter.load("cpu") expected = adapter._model(_probe()).detach().clone() @@ -1450,15 +1451,35 @@ def snapshot_download(repo_id: str, **kwargs: Any) -> str: adapter.load_lora("org/hub-lora", revision=revision) adapter.set_active_lora("org/hub-lora") - assert calls == [ - ( - "org/hub-lora", - {"revision": revision, "allow_patterns": ["adapter_config.json", "adapter_model.safetensors"]}, - ) - ] + assert calls == [("org/hub-lora", revision)] assert not torch.allclose(adapter._model(_probe()), expected) +def test_peft_loads_pinned_cluster_adapter_with_network_denied(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + revision = "c" * 40 + cluster = tmp_path / "cluster" + cached = cluster / "models--org--adapter" / "snapshots" / revision + cached.parent.mkdir(parents=True) + _save_plain_lora(cached) + monkeypatch.setenv("SIE_LOCAL_CACHE", str(tmp_path / "local")) + monkeypatch.setenv("SIE_CLUSTER_CACHE", str(cluster)) + monkeypatch.setenv("SIE_HF_FALLBACK", "false") + monkeypatch.setattr(socket.socket, "connect", MagicMock(side_effect=AssertionError("network forbidden"))) + monkeypatch.setattr("huggingface_hub.snapshot_download", MagicMock(side_effect=AssertionError("Hub forbidden"))) + adapter = TinyLoRAAdapter() + adapter.load("cpu") + try: + baseline = adapter._model(_probe()).detach().clone() + adapter.load_lora("org/adapter", revision=revision) + adapter.set_active_lora("org/adapter") + assert not torch.allclose(adapter._model(_probe()), baseline) + adapter.set_active_lora(None) + assert torch.equal(adapter._model(_probe()), baseline) + adapter.unload_lora("org/adapter") + finally: + adapter.unload() + + class TestLoraWeightKeys: @pytest.mark.parametrize( "key", diff --git a/packages/sie_server/tests/adapters/test_lora_source.py b/packages/sie_server/tests/adapters/test_lora_source.py new file mode 100644 index 000000000..094a0b80d --- /dev/null +++ b/packages/sie_server/tests/adapters/test_lora_source.py @@ -0,0 +1,355 @@ +"""LoRA cache resolution with every network connection denied.""" + +import socket +from pathlib import Path +from unittest.mock import Mock + +import huggingface_hub +import pytest +from huggingface_hub.utils import HFValidationError +from sie_server.adapters._lora_source import resolve_lora_source + +SHA = "a" * 40 +OTHER_SHA = "b" * 40 +FILES = ("adapter_config.json", "adapter_model.safetensors") + + +def snapshot(root: Path, source: str = "tenant-a/adapter", revision: str = SHA) -> Path: + repository = root / f"models--{source.replace('/', '--')}" + directory = repository / "snapshots" / revision + directory.mkdir(parents=True) + for filename in FILES: + (directory / filename).write_bytes(f"{source}:{revision}:{filename}".encode()) + (repository / "refs").mkdir(exist_ok=True) + (repository / "refs" / "main").write_text(revision) + return directory + + +@pytest.fixture +def cache_policy(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + local, cluster = tmp_path / "local", tmp_path / "cluster" + monkeypatch.setenv("SIE_LOCAL_CACHE", str(local)) + monkeypatch.setenv("SIE_CLUSTER_CACHE", str(cluster)) + monkeypatch.setenv("SIE_HF_FALLBACK", "false") + monkeypatch.setattr(socket.socket, "connect", Mock(side_effect=AssertionError("network forbidden"))) + monkeypatch.setattr(huggingface_hub, "snapshot_download", Mock(side_effect=AssertionError("Hub forbidden"))) + monkeypatch.setattr(huggingface_hub, "file_exists", Mock(side_effect=AssertionError("Hub metadata forbidden"))) + return local, cluster + + +@pytest.mark.parametrize("revision", [None, SHA]) +@pytest.mark.parametrize("tier", ["local", "cluster"]) +@pytest.mark.parametrize("fallback", ["false", "true"]) +def test_exact_adapter_from_cache_without_hub( + cache_policy: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, tier: str, revision: str | None, fallback: str +) -> None: + local, cluster = cache_policy + monkeypatch.setenv("SIE_HF_FALLBACK", fallback) + original = snapshot(local if tier == "local" else cluster) + + directory = resolve_lora_source("tenant-a/adapter", revision) + + assert directory == local / "models--tenant-a--adapter" / "snapshots" / SHA + assert [path.read_bytes() for path in (directory / name for name in FILES)] == [ + (original / name).read_bytes() for name in FILES + ] + + +def test_other_revision_cannot_satisfy_pin(cache_policy: tuple[Path, Path]) -> None: + local, cluster = cache_policy + snapshot(local, revision=OTHER_SHA) + snapshot(cluster, revision=OTHER_SHA) + with pytest.raises(RuntimeError, match="fallback is disabled"): + resolve_lora_source("tenant-a/adapter", SHA) + + +@pytest.mark.parametrize("filename", FILES) +def test_incomplete_snapshot_fails_closed(cache_policy: tuple[Path, Path], filename: str) -> None: + local, _ = cache_policy + directory = snapshot(local) + (directory / filename).unlink() + with pytest.raises(RuntimeError, match="fallback is disabled"): + resolve_lora_source("tenant-a/adapter", SHA) + + +def test_default_branch_never_chooses_an_arbitrary_snapshot(cache_policy: tuple[Path, Path]) -> None: + local, _ = cache_policy + directory = snapshot(local) + (directory.parent.parent / "refs" / "main").unlink() + with pytest.raises(RuntimeError, match="fallback is disabled"): + resolve_lora_source("tenant-a/adapter") + + +def test_default_reference_cannot_alias_another_tenant(cache_policy: tuple[Path, Path]) -> None: + local, _ = cache_policy + directory = snapshot(local) + other = snapshot(local, source="tenant-b/adapter", revision=OTHER_SHA) + reference = directory.parent.parent / "refs" / "main" + reference.unlink() + reference.symlink_to(other.parent.parent / "refs" / "main") + with pytest.raises(ValueError, match="file alias"): + resolve_lora_source("tenant-a/adapter") + + +def test_configured_root_symlink_preserves_exact_repository( + cache_policy: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + local, _ = cache_policy + directory = snapshot(local) + selected = tmp_path / "selected" + selected.symlink_to(local, target_is_directory=True) + monkeypatch.setenv("SIE_LOCAL_CACHE", str(selected)) + assert resolve_lora_source("tenant-a/adapter", SHA) == directory + + +@pytest.mark.parametrize("revision", ["main", "../other", "a" * 39, "A" * 40]) +def test_mutable_or_malformed_pin_is_refused(cache_policy: tuple[Path, Path], revision: str) -> None: + with pytest.raises(ValueError, match="immutable"): + resolve_lora_source("tenant-a/adapter", revision) + + +@pytest.mark.parametrize("source", ["../tenant-b/adapter", "a/b/c", "s3://bucket/adapter", "tenant\\adapter"]) +def test_invalid_repository_is_refused(cache_policy: tuple[Path, Path], source: str) -> None: + with pytest.raises(HFValidationError): + resolve_lora_source(source) + + +def test_local_root_symlink_is_supported_but_pin_is_not(cache_policy: tuple[Path, Path], tmp_path: Path) -> None: + actual = snapshot(tmp_path / "fixtures") + selected = tmp_path / "selected" + selected.symlink_to(actual, target_is_directory=True) + assert resolve_lora_source(str(selected)) == actual + with pytest.raises(ValueError, match="local LoRA directory cannot honor"): + resolve_lora_source(str(selected), SHA) + + +def test_local_file_escape_is_refused(cache_policy: tuple[Path, Path], tmp_path: Path) -> None: + actual = snapshot(tmp_path / "fixtures") + outside = tmp_path / "outside" + outside.write_bytes(b"not this adapter") + (actual / FILES[1]).unlink() + (actual / FILES[1]).symlink_to(outside) + with pytest.raises(ValueError, match="escapes"): + resolve_lora_source(str(actual)) + + +def test_repository_cannot_alias_another_tenant(cache_policy: tuple[Path, Path]) -> None: + local, _ = cache_policy + other = snapshot(local, source="tenant-b/adapter") + (local / "models--tenant-a--adapter").symlink_to(other.parent.parent, target_is_directory=True) + with pytest.raises(ValueError, match="aliases"): + resolve_lora_source("tenant-a/adapter", SHA) + + +@pytest.mark.parametrize("part", ["snapshots", "refs", "blobs"]) +def test_repository_layout_cannot_alias_another_tenant(cache_policy: tuple[Path, Path], part: str) -> None: + local, _ = cache_policy + other = snapshot(local, source="tenant-b/adapter").parent.parent + (other / "blobs").mkdir() + repository = local / "models--tenant-a--adapter" + repository.mkdir() + (repository / part).symlink_to(other / part, target_is_directory=True) + with pytest.raises(ValueError, match="aliases"): + resolve_lora_source("tenant-a/adapter", SHA) + + +def test_snapshot_cannot_alias_another_revision(cache_policy: tuple[Path, Path]) -> None: + local, _ = cache_policy + other = snapshot(local, revision=OTHER_SHA) + (other.parent / SHA).symlink_to(other, target_is_directory=True) + with pytest.raises(ValueError, match="another revision"): + resolve_lora_source("tenant-a/adapter", SHA) + + +@pytest.mark.parametrize("other_tenant", [False, True]) +def test_hf_file_links_are_limited_to_this_repository_blobs( + cache_policy: tuple[Path, Path], other_tenant: bool +) -> None: + local, _ = cache_policy + directory = snapshot(local) + repository = snapshot(local, source="tenant-b/adapter").parent.parent if other_tenant else directory.parent.parent + blob = repository / "blobs" / "digest" + blob.parent.mkdir() + blob.write_bytes(b"blob") + (directory / FILES[1]).unlink() + (directory / FILES[1]).symlink_to(blob) + if other_tenant: + with pytest.raises(ValueError, match="escapes"): + resolve_lora_source("tenant-a/adapter", SHA) + else: + assert resolve_lora_source("tenant-a/adapter", SHA) == directory + + +def test_hub_fallback_uses_the_configured_cache_and_exact_pin( + cache_policy: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + local, _ = cache_policy + monkeypatch.setenv("SIE_HF_FALLBACK", "true") + download = Mock(side_effect=lambda *_args, **_kwargs: str(snapshot(local))) + monkeypatch.setattr(huggingface_hub, "snapshot_download", download) + monkeypatch.setattr(huggingface_hub, "file_exists", Mock(return_value=False)) + + assert resolve_lora_source("tenant-a/adapter", SHA) == local / "models--tenant-a--adapter" / "snapshots" / SHA + download.assert_called_once() + assert download.call_args.args == ("tenant-a/adapter",) + assert download.call_args.kwargs["revision"] == SHA + assert download.call_args.kwargs["cache_dir"] == str(local) + assert download.call_args.kwargs["allow_patterns"] == list(FILES) + + +@pytest.mark.parametrize("selected_alias", [False, True]) +def test_explicit_local_hf_snapshot_supports_owned_blob_links( + cache_policy: tuple[Path, Path], tmp_path: Path, selected_alias: bool +) -> None: + directory = snapshot(tmp_path / "explicit-cache") + repository = directory.parent.parent + blobs = repository / "blobs" + blobs.mkdir() + for filename in FILES: + original = (directory / filename).read_bytes() + (blobs / filename).write_bytes(original) + (directory / filename).unlink() + (directory / filename).symlink_to(Path("../../blobs") / filename) + selected = directory + if selected_alias: + selected = tmp_path / "selected" + selected.symlink_to(directory, target_is_directory=True) + + assert resolve_lora_source(str(selected)) == directory + with pytest.raises(ValueError, match="local LoRA directory cannot honor"): + resolve_lora_source(str(selected), SHA) + + +@pytest.mark.parametrize("escape", ["other-repository", "blobs-directory-alias", "non-hf-layout"]) +def test_explicit_local_snapshot_does_not_allow_unowned_blob_links( + cache_policy: tuple[Path, Path], tmp_path: Path, escape: str +) -> None: + directory = snapshot(tmp_path / "explicit-cache") + repository = directory.parent.parent + outside = snapshot(tmp_path / "explicit-cache", source="tenant-b/adapter").parent.parent / "blobs" + outside.mkdir() + (outside / "weights").write_bytes(b"other repository") + if escape == "blobs-directory-alias": + (repository / "blobs").symlink_to(outside, target_is_directory=True) + target = repository / "blobs" / "weights" + elif escape == "non-hf-layout": + plain = tmp_path / "plain" + directory.rename(plain) + directory = plain + target = outside / "weights" + else: + target = outside / "weights" + (directory / FILES[1]).unlink() + (directory / FILES[1]).symlink_to(target) + with pytest.raises(ValueError, match=r"escapes|aliases"): + resolve_lora_source(str(directory)) + + +def complete_hub_snapshot(local: Path) -> Mock: + def download(*_args: object, **_kwargs: object) -> str: + repository = local / "models--tenant-a--adapter" + directory = repository / "snapshots" / SHA + directory.mkdir(parents=True, exist_ok=True) + for filename in FILES: + (directory / filename).write_bytes(f"complete:{filename}".encode()) + (repository / "refs").mkdir(exist_ok=True) + (repository / "refs" / "main").write_text(SHA) + return str(directory) + + return Mock(side_effect=download) + + +@pytest.mark.parametrize("missing", FILES) +@pytest.mark.parametrize("revision", [None, SHA]) +@pytest.mark.parametrize("tier", ["cluster", "hub"]) +def test_partial_local_snapshot_continues_to_correct_tier( + cache_policy: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + missing: str, + revision: str | None, + tier: str, +) -> None: + local, cluster = cache_policy + partial = snapshot(local) + (partial / missing).unlink() + download = None + if tier == "cluster": + expected = snapshot(cluster) + else: + monkeypatch.setenv("SIE_HF_FALLBACK", "true") + download = complete_hub_snapshot(local) + monkeypatch.setattr(huggingface_hub, "snapshot_download", download) + monkeypatch.setattr(huggingface_hub, "file_exists", Mock(return_value=False)) + expected = partial + + resolved = resolve_lora_source("tenant-a/adapter", revision) + + assert resolved == partial + assert all((resolved / name).is_file() for name in FILES) + if download is not None: + download.assert_called_once() + assert download.call_args.kwargs["revision"] == (revision or "main") + assert download.call_args.kwargs["allow_patterns"] == list(FILES) + else: + assert [(resolved / name).read_bytes() for name in FILES] == [(expected / name).read_bytes() for name in FILES] + + +@pytest.mark.parametrize("cached_revision", [SHA, OTHER_SHA]) +@pytest.mark.parametrize("tier", ["cluster", "hub"]) +def test_default_ref_miss_continues_to_correct_tier( + cache_policy: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, cached_revision: str, tier: str +) -> None: + local, cluster = cache_policy + old = snapshot(local, revision=cached_revision) + old_bytes = [(old / name).read_bytes() for name in FILES] + (old.parent.parent / "refs" / "main").unlink() + download = None + if tier == "cluster": + snapshot(cluster) + else: + monkeypatch.setenv("SIE_HF_FALLBACK", "true") + download = complete_hub_snapshot(local) + monkeypatch.setattr(huggingface_hub, "snapshot_download", download) + monkeypatch.setattr(huggingface_hub, "file_exists", Mock(return_value=False)) + + resolved = resolve_lora_source("tenant-a/adapter") + + assert resolved == old.parent / SHA + assert all((resolved / name).is_file() for name in FILES) + if cached_revision != SHA: + assert [(old / name).read_bytes() for name in FILES] == old_bytes + if download is not None: + download.assert_called_once() + assert download.call_args.kwargs["revision"] == "main" + + +@pytest.mark.parametrize("fallback", [False, True]) +def test_partial_cluster_hit_requires_complete_files_before_return( + cache_policy: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, fallback: bool +) -> None: + local, cluster = cache_policy + partial = snapshot(cluster) + (partial / FILES[1]).unlink() + if not fallback: + with pytest.raises(RuntimeError, match="fallback is disabled"): + resolve_lora_source("tenant-a/adapter", SHA) + assert (local / partial.relative_to(cluster) / FILES[0]).is_file() + assert not (local / partial.relative_to(cluster) / FILES[1]).exists() + else: + monkeypatch.setenv("SIE_HF_FALLBACK", "true") + download = complete_hub_snapshot(local) + monkeypatch.setattr(huggingface_hub, "snapshot_download", download) + monkeypatch.setattr(huggingface_hub, "file_exists", Mock(return_value=False)) + resolved = resolve_lora_source("tenant-a/adapter", SHA) + assert all((resolved / name).is_file() for name in FILES) + download.assert_called_once() + + +@pytest.mark.parametrize("filename", FILES) +def test_explicit_incomplete_directory_still_fails_locally( + cache_policy: tuple[Path, Path], tmp_path: Path, filename: str +) -> None: + directory = snapshot(tmp_path / "explicit") + (directory / filename).unlink() + with pytest.raises(ValueError, match=f"ships no {filename}"): + resolve_lora_source(str(directory)) diff --git a/packages/sie_server/tests/adapters/test_sglang_lora.py b/packages/sie_server/tests/adapters/test_sglang_lora.py index eb916b9fa..2ae0d6f41 100644 --- a/packages/sie_server/tests/adapters/test_sglang_lora.py +++ b/packages/sie_server/tests/adapters/test_sglang_lora.py @@ -2,6 +2,7 @@ import asyncio import json +import socket import threading from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -103,7 +104,7 @@ def test_startup_only_capacity_does_not_mask_invalid_limits(field, invalid): def test_undeclared_source_never_fetches_or_calls_engine(monkeypatch): fetch = Mock() post = Mock() - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", fetch) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", fetch) monkeypatch.setattr("sie_server.adapters.sglang._lora.requests.post", post) loras = manager() for operation in (loras.fetch, loras.load, loras.unload): @@ -119,8 +120,8 @@ def test_fetch_uses_immutable_revision_and_only_adapter_files(tmp_path, monkeypa {"base_model.model.q_proj.lora_A.weight": np.ones((4, 2), dtype=np.float32)}, str(tmp_path / "adapter_model.safetensors"), ) - fetch = Mock(return_value=str(tmp_path)) - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", fetch) + fetch = Mock(return_value=tmp_path) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", fetch) (tmp_path / "unvalidated.safetensors").write_bytes(b"must not reach the engine") loras = manager() snapshot = Path(loras.fetch("a")) @@ -129,9 +130,7 @@ def test_fetch_uses_immutable_revision_and_only_adapter_files(tmp_path, monkeypa assert (snapshot / "adapter_model.safetensors").read_bytes() == ( tmp_path / "adapter_model.safetensors" ).read_bytes() - fetch.assert_called_once_with( - "org/a", revision="a" * 40, allow_patterns=["adapter_config.json", "adapter_model.safetensors"] - ) + fetch.assert_called_once_with("org/a", "a" * 40) with pytest.raises(ValueError, match="revision"): loras.fetch("a", "b" * 40) loras.close() @@ -139,6 +138,32 @@ def test_fetch_uses_immutable_revision_and_only_adapter_files(tmp_path, monkeypa assert (tmp_path / "adapter_model.safetensors").exists() +def test_sglang_prepares_pinned_cluster_adapter_with_network_denied(tmp_path, monkeypatch): + cluster = tmp_path / "cluster" + cached = cluster / "models--org--a" / "snapshots" / ("a" * 40) + cached.mkdir(parents=True) + (cached / "adapter_config.json").write_text('{"peft_type":"LORA","r":4,"target_modules":["q_proj"]}') + save_file( + {"base.q_proj.lora_A.weight": np.ones((4, 2), dtype=np.float32)}, str(cached / "adapter_model.safetensors") + ) + monkeypatch.setenv("SIE_LOCAL_CACHE", str(tmp_path / "local")) + monkeypatch.setenv("SIE_CLUSTER_CACHE", str(cluster)) + monkeypatch.setenv("SIE_HF_FALLBACK", "false") + monkeypatch.setattr(socket.socket, "connect", Mock(side_effect=AssertionError("network forbidden"))) + monkeypatch.setattr("huggingface_hub.snapshot_download", Mock(side_effect=AssertionError("Hub forbidden"))) + loras = manager() + try: + prepared = Path(loras.fetch("a")) + assert prepared != cached + assert {path.name for path in prepared.iterdir()} == {"adapter_config.json", "adapter_model.safetensors"} + assert (prepared / "adapter_model.safetensors").read_bytes() == ( + cached / "adapter_model.safetensors" + ).read_bytes() + finally: + loras.close() + assert not prepared.exists() + + @pytest.mark.parametrize( "override", [ @@ -153,7 +178,7 @@ def test_fetch_uses_immutable_revision_and_only_adapter_files(tmp_path, monkeypa def test_fetch_refuses_unsupported_transform_or_engine_capacity(tmp_path, monkeypatch, override): config = {"peft_type": "LORA", "r": 4, "target_modules": ["q_proj"], **override} (tmp_path / "adapter_config.json").write_text(json.dumps(config)) - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", lambda *args, **kwargs: tmp_path) with pytest.raises(ValueError, match="LoRA"): manager().fetch("a") @@ -168,9 +193,9 @@ def test_fetch_finishing_after_model_close_does_not_retain_source(tmp_path, monk def download(*args, **kwargs): entered.set() assert release.wait(2) - return str(tmp_path) + return tmp_path - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", download) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", download) loras = manager() with ThreadPoolExecutor() as pool: pending = pool.submit(loras.fetch, "a") @@ -187,7 +212,7 @@ def test_discard_source_keeps_attached_or_uncertain_views_until_model_close(tmp_ save_file( {"base.q_proj.lora_A.weight": np.ones((4, 2), dtype=np.float32)}, str(tmp_path / "adapter_model.safetensors") ) - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", lambda *args, **kwargs: tmp_path) loras = manager() unused = Path(loras.fetch("a")) loras.discard_source("a") @@ -251,7 +276,7 @@ def test_fetch_rejects_shared_base_tensor(tmp_path, monkeypatch): {"base_model.model.q_proj.weight": np.ones((4, 2), dtype=np.float32)}, str(tmp_path / "adapter_model.safetensors"), ) - monkeypatch.setattr("sie_server.adapters.sglang._lora.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr("sie_server.adapters.sglang._lora.resolve_lora_source", lambda *args, **kwargs: tmp_path) with pytest.raises(ValueError, match="outside the adapter"): manager().fetch("a")