diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml index f48631911..7b8dec2b8 100644 --- a/.github/workflows/run-tests.yaml +++ b/.github/workflows/run-tests.yaml @@ -112,6 +112,7 @@ jobs: tests/test_model_catalog.py \ tests/test_model_manager.py \ tests/test_model_manifest.py \ + tests/test_local_artifact_snapshot.py \ tests/test_multi_machine_qualification.py \ tests/test_native_credentials.py \ tests/test_node_config.py \ diff --git a/docs/MODEL_QUALIFICATION_V1.md b/docs/MODEL_QUALIFICATION_V1.md index 4cf3651e9..912c11bbc 100644 --- a/docs/MODEL_QUALIFICATION_V1.md +++ b/docs/MODEL_QUALIFICATION_V1.md @@ -61,6 +61,45 @@ with opaque labels. The report also records the normalized operating system, req device profile, observed worker device, dtype, attention implementation, and the source commit when it can resolve one from the checkout. +## Loading an existing read-only snapshot + +The fixed-manifest worker and text-peer CLIs accept `--artifact_root` for a +complete local snapshot, paired with an explicit, separate writable `--cache_dir`: + +```text +drift server --model_manifest \ + --artifact_root /models/verified-snapshot --cache_dir /run/communityai/cache \ + --block_indices 0:14 --identity_path +drift text-peer \ + --artifact_root /models/verified-snapshot --cache_dir /run/communityai/cache \ + --identity_path --initial_peers +``` + +Each declared artifact lives at its manifest-relative path below `artifact_root`. +This interface requires a materialized tree containing exactly the declared +regular files and their parent directories. Symlinks, Windows reparse points, +extra files (including loader sidecars, README/license files) and extra directories +are rejected. A conventional Hub snapshot containing blob symlinks therefore +needs a separate materialized, manifest-only snapshot for this interface. +Every declared file, including unassigned shards and tokenizer files, is checked +by size and SHA-256 before construction of a model loader; the text-peer CLI also +checks before identity/DHT startup. Missing or corrupt files fail without Hub or +custom HTTP fallback. The root reaches worker config, module and block loaders, +and the text-peer client loader. Runtime locks such as `blocks.lock` belong to +`cache_dir`; neither directory may contain the other. The caller supplies and +maintains the read-only snapshot mount. Managed workers retain all five manifest, +span, artifact-byte, artifact-set and writable-cache placement claims. +These examples apply to direct `drift server` and `drift text-peer` launches. +The `drift node` supervisor does not propagate this opt-in flag to its children. + +The Python `make_manifest_loader(..., artifact_root=..., cache_dir=...)` interface +uses the same boundary. `ManifestArtifactVerifier(..., cache_only=True)` separately +supports verified existing Hub/custom cache entries without starting a transfer +on a miss. Omitting these opt-in settings preserves the existing online behavior. +Content validation does not create signed catalog approval or authorize model or +provider execution. Direct fixed-manifest experiments and automatic signed-catalog +selection remain separate interfaces. + ## Cross-platform matrix gate Each real host must produce a complete artifact, parity, and local-failover report with diff --git a/src/drift/artifact_snapshot.py b/src/drift/artifact_snapshot.py new file mode 100644 index 000000000..8fa0c2fa0 --- /dev/null +++ b/src/drift/artifact_snapshot.py @@ -0,0 +1,72 @@ +"""Local artifact snapshot validation, independent of catalog or owner admission.""" + +import stat +from pathlib import Path, PurePosixPath +from typing import Optional, Union + +from drift.model_manifest import ManifestError, ModelManifest + + +def validate_artifact_snapshot( + manifest: Optional[ModelManifest], + artifact_root: Optional[Union[str, Path]], + *, + cache_dir: Optional[Union[str, Path]], +) -> Optional[Path]: + """Hash a materialized, manifest-only tree without writing to the snapshot. + + Runtime locks belong to an explicit, separate cache directory. This checks + content only; it does not approve a catalog entry or authorize a role start. + """ + if artifact_root is None: + return None + if not isinstance(manifest, ModelManifest): + raise ManifestError("artifact_root requires an exact model manifest") + if cache_dir is None or not str(cache_dir): + raise ManifestError("artifact_root requires a separate writable cache_dir") + try: + root = Path(artifact_root).expanduser().resolve(strict=True) + cache = Path(cache_dir).expanduser().resolve() + except (OSError, TypeError, ValueError) as exc: + raise ManifestError("Could not resolve artifact snapshot or cache directory") from exc + if not root.is_dir(): + raise ManifestError("artifact_root must be an existing snapshot directory") + if cache.exists() and not cache.is_dir(): + raise ManifestError("cache_dir must be a writable cache directory") + if root == cache or root.is_relative_to(cache) or cache.is_relative_to(root): + raise ManifestError("artifact_root and writable cache_dir must be disjoint directories") + _validate_snapshot_tree(manifest, root) + manifest.verify_artifacts(root) + return root + + +def _validate_snapshot_tree(manifest: ModelManifest, root: Path) -> None: + files = {artifact.path for artifact in manifest.artifacts} + directories = { + parent.as_posix() + for filename in files + for parent in PurePosixPath(filename).parents + if parent != PurePosixPath(".") + } + pending = [root] + try: + while pending: + for path in pending.pop().iterdir(): + relative = path.relative_to(root).as_posix() + metadata = path.lstat() + if stat.S_ISLNK(metadata.st_mode) or getattr(metadata, "st_file_attributes", 0) & getattr( + stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0 + ): + raise ManifestError(f"Artifact snapshot cannot contain links or reparse points: {relative}") + if not path.resolve(strict=True).is_relative_to(root): + raise ManifestError(f"Artifact escapes snapshot root: {relative}") + if stat.S_ISDIR(metadata.st_mode): + if relative not in directories: + raise ManifestError(f"Undeclared artifact snapshot directory: {relative}") + pending.append(path) + elif not stat.S_ISREG(metadata.st_mode): + raise ManifestError(f"Artifact snapshot entry must be a regular file: {relative}") + elif relative not in files: + raise ManifestError(f"Undeclared artifact snapshot file: {relative}") + except OSError as exc: + raise ManifestError("Could not inspect artifact snapshot tree") from exc diff --git a/src/drift/cli/run_server.py b/src/drift/cli/run_server.py index 991351b7c..e3a9484ac 100644 --- a/src/drift/cli/run_server.py +++ b/src/drift/cli/run_server.py @@ -250,6 +250,8 @@ def build_parser(*, bound_worker: bool = False) -> configargparse.ArgParser: parser.add_argument('--cache_dir', type=str, default=None, help='Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used.') + parser.add_argument('--artifact_root', type=str, default=None, + help='Read-only snapshot of every manifested artifact; requires --model_manifest and a separate writable --cache_dir. Missing or corrupt files fail without download.') parser.add_argument("--max_disk_space", type=str, default=None, help="Maximal disk space used for caches. Example: 50GB, 100GiB (GB != GiB here). " "Default: unlimited. " diff --git a/src/drift/cli/run_text_peer.py b/src/drift/cli/run_text_peer.py index d01d25318..982fc91c5 100644 --- a/src/drift/cli/run_text_peer.py +++ b/src/drift/cli/run_text_peer.py @@ -7,6 +7,7 @@ from hivemind import DHT +from drift.artifact_snapshot import validate_artifact_snapshot from drift.model_manifest import ModelManifest from drift.protocol_identity import NodeIdentity from drift.server.text_peer import TextPeerService @@ -18,12 +19,16 @@ def main(): parser.add_argument("--initial_peers", nargs="+", required=True) parser.add_argument("--identity_path", type=Path, required=True) parser.add_argument("--cache_dir", type=Path, required=True) + parser.add_argument( + "--artifact_root", type=Path, help="Complete read-only manifested snapshot; cache_dir owns locks" + ) parser.add_argument("--host_maddrs", nargs="+", default=["/ip4/0.0.0.0/tcp/31337"]) parser.add_argument("--announce_maddrs", nargs="+") parser.add_argument("--max_context_tokens", type=int, default=2048) parser.add_argument("--max_output_tokens", type=int, default=512) args = parser.parse_args() manifest = ModelManifest.load(args.manifest) + artifact_root = validate_artifact_snapshot(manifest, args.artifact_root, cache_dir=args.cache_dir) identity = NodeIdentity.ensure(args.identity_path) dht = DHT( initial_peers=args.initial_peers, @@ -39,6 +44,7 @@ def main(): manifest, initial_peers=args.initial_peers, cache_dir=str(args.cache_dir), + artifact_root=artifact_root, max_context_tokens=args.max_context_tokens, max_output_tokens=args.max_output_tokens, ).start() diff --git a/src/drift/model_manifest.py b/src/drift/model_manifest.py index c3c624d41..21e857abe 100644 --- a/src/drift/model_manifest.py +++ b/src/drift/model_manifest.py @@ -661,7 +661,8 @@ class ManifestArtifactVerifier: A verifier deliberately materializes only requested files. This preserves Petals' partial-checkpoint behavior: a worker verifies the shards containing its assigned blocks, while an API client verifies the tokenizer and the shards containing its local embeddings/head. Successful hashes are cached only while - the resolved path, size, and modification timestamp remain unchanged. + the resolved path, size, and modification timestamp remain unchanged. Explicit ``cache_only`` mode + accepts verified local files but never starts a transfer on a cache miss. """ manifest: ModelManifest @@ -672,6 +673,7 @@ class ManifestArtifactVerifier: max_disk_space: Optional[int] = None artifact_root: Optional[Union[str, os.PathLike]] = None allowed_paths: Optional[Iterable[str]] = None + cache_only: bool = False _verified: Dict[Tuple[str, int, int, str], bool] = field(default_factory=dict, init=False, repr=False) _snapshot_root: Optional[Path] = field(default=None, init=False, repr=False) _weight_map: Optional[Mapping[str, str]] = field(default=None, init=False, repr=False) @@ -690,6 +692,8 @@ def __post_init__(self) -> None: raise ManifestError( f"Artifact verifier revision is {self.revision!r}, expected {self.manifest.source.revision!r}" ) + if type(self.cache_only) is not bool: + raise ManifestError("Artifact verifier cache_only must be a boolean") if self.allowed_paths is not None: self.allowed_paths = self._normalize_allowed_paths(self.allowed_paths) if self.artifact_root is not None: @@ -842,6 +846,7 @@ def _ensure_path(self, path: str, *, allowed_roles: Optional[Iterable[str]] = No ) candidate = None + already_verified = False if self._snapshot_root is not None: candidate = _artifact_path_below_root(self._snapshot_root, artifact.path) if candidate is None or (not candidate.exists() and self.artifact_root is None): @@ -862,16 +867,24 @@ def _ensure_path(self, path: str, *, allowed_roles: Optional[Iterable[str]] = No resolved = None if resolved is None: - from drift.utils.disk_cache import allow_cache_writes, free_disk_space_for - - Path(self.cache_dir).mkdir(parents=True, exist_ok=True) - with allow_cache_writes(self.cache_dir): - free_disk_space_for( - max(0, artifact.size - self.partial_size(artifact.path)), - cache_dir=self.cache_dir, - max_disk_space=self.max_disk_space, - ) - resolved = self._resumable_hub_download(artifact, destination=candidate) + if self.cache_only: + _, cached_final, _ = self._resumable_paths(artifact) + if not cached_final.exists(): + raise ManifestError(f"Artifact {artifact.path} is not available in the local cache") + _verify_artifact_file(artifact, cached_final) + already_verified = True + resolved = str(cached_final) + else: + from drift.utils.disk_cache import allow_cache_writes, free_disk_space_for + + Path(self.cache_dir).mkdir(parents=True, exist_ok=True) + with allow_cache_writes(self.cache_dir): + free_disk_space_for( + max(0, artifact.size - self.partial_size(artifact.path)), + cache_dir=self.cache_dir, + max_disk_space=self.max_disk_space, + ) + resolved = self._resumable_hub_download(artifact, destination=candidate) except ManifestTransferInterrupted: raise except Exception as exc: @@ -893,7 +906,8 @@ def _ensure_path(self, path: str, *, allowed_roles: Optional[Iterable[str]] = No self._promote_cached_artifact(artifact, resolved_candidate, candidate) self._report(artifact, "verifying", received=artifact.size) - self._verify(artifact, candidate) + if not already_verified: + self._verify(artifact, candidate) return candidate def _promote_cached_artifact(self, artifact: ManifestArtifact, source: Path, destination: Path) -> None: @@ -954,6 +968,9 @@ def _resumable_hub_download(self, artifact: ManifestArtifact, *, destination: Op A partial is never exposed to Transformers. It is atomically promoted only after its exact declared size and SHA-256 pass. """ + if self.cache_only: + raise ManifestError("Artifact transfer is disabled in cache-only mode") + import requests from huggingface_hub import hf_hub_url from huggingface_hub.utils import build_hf_headers @@ -1096,11 +1113,12 @@ def _verify(self, artifact: ManifestArtifact, candidate: Path) -> None: except OSError as exc: raise ManifestError(f"Could not read artifact {artifact.path}: {exc}") from exc cache_key = (str(candidate), stat_result.st_size, stat_result.st_mtime_ns, artifact.sha256) - if cache_key in self._verified: + if not self.cache_only and cache_key in self._verified: return _verify_artifact_file(artifact, candidate) - self._verified = {key: value for key, value in self._verified.items() if key[0] != str(candidate)} - self._verified[cache_key] = True + if not self.cache_only: + self._verified = {key: value for key, value in self._verified.items() if key[0] != str(candidate)} + self._verified[cache_key] = True def _select_snapshot_artifacts(files: Iterable[str], root: Path | str) -> Dict[str, str]: diff --git a/src/drift/node/loading.py b/src/drift/node/loading.py index 160430756..4a19e9eac 100644 --- a/src/drift/node/loading.py +++ b/src/drift/node/loading.py @@ -8,6 +8,7 @@ from hivemind.utils.logging import get_logger +from drift.artifact_snapshot import validate_artifact_snapshot from drift.model_manifest import ManifestArtifactVerifier, ManifestError, ModelManifest from drift.node.model_manager import ModelRuntime from drift.node.route_health import sequence_manager_route_health @@ -63,6 +64,7 @@ def make_manifest_loader( initial_peers: Sequence[str], token: Optional[str] = None, cache_dir: Optional[str] = None, + artifact_root: Optional[str] = None, revocation_files: Sequence[str] = (), request_timeout: float = 30, max_retries: int = 3, @@ -70,6 +72,7 @@ def make_manifest_loader( """Return a lazy loader pinned to one exact manifest and swarm namespace.""" def load() -> ModelRuntime: + root = validate_artifact_snapshot(manifest, artifact_root, cache_dir=cache_dir) import torch from transformers import AutoTokenizer @@ -81,6 +84,8 @@ def load() -> ModelRuntime: revision=manifest.source.revision, token=token, cache_dir=cache_dir, + artifact_root=root, + cache_only=root is not None, ) verifier.ensure_startup_metadata(include_tokenizer=True) logger.info(f"Loading tokenizer and client-side weights for {manifest.digest_id}") diff --git a/src/drift/server/server.py b/src/drift/server/server.py index 0bab0251a..08a8a5427 100644 --- a/src/drift/server/server.py +++ b/src/drift/server/server.py @@ -22,6 +22,7 @@ from transformers import PretrainedConfig import drift +from drift.artifact_snapshot import validate_artifact_snapshot from drift.constants import DTYPE_MAP from drift.data_structures import CHAIN_DELIMITER, UID_DELIMITER, ModelInfo, ServerInfo, ServerState, parse_uid from drift.model_manifest import ManifestArtifactVerifier, ManifestError, ModelManifest @@ -181,6 +182,7 @@ def _scoped_manifest_artifact_verifier( cache_dir=cache_dir, max_disk_space=max_disk_space, artifact_root=artifact_root, + cache_only=artifact_root is not None, allowed_paths=startup_paths, ) artifact_plan = verifier.bind_block_artifact_plan( @@ -257,6 +259,7 @@ def __init__( admission_policy: Optional[AdmissionPolicy] = None, revocation_files: Sequence[str] = (), cache_dir: Optional[str] = None, + artifact_root: Optional[str] = None, max_disk_space: Optional[int] = None, max_device_memory: Optional[int] = None, max_processing_percent: float = 100, @@ -292,6 +295,7 @@ def __init__( self._managed_loading_session = managed_loading_session self._managed_lifecycle_started = False self._managed_ready = False + self.artifact_root = validate_artifact_snapshot(model_manifest, artifact_root, cache_dir=cache_dir) converted_model_name_or_path = get_compatible_model_repo(converted_model_name_or_path) self.converted_model_name_or_path = converted_model_name_or_path @@ -315,6 +319,8 @@ def __init__( token=token, cache_dir=cache_dir, max_disk_space=max_disk_space, + artifact_root=self.artifact_root, + cache_only=self.artifact_root is not None, ) config_source = artifact_verifier.ensure_startup_metadata() else: @@ -732,6 +738,7 @@ def _create_module_container(self, block_indices: List[int]) -> ModuleContainer: inference_max_length=self.inference_max_length, torch_dtype=self.torch_dtype, cache_dir=self.cache_dir, + artifact_root=self.artifact_root, max_disk_space=self.max_disk_space, device=self.device, compression=self.compression, @@ -919,6 +926,7 @@ def create( torch_dtype: torch.dtype, cache_dir: str, max_disk_space: int, + artifact_root: Optional[str] = None, device: Union[str, torch.device], compression: CompressionType, update_period: float, @@ -954,6 +962,7 @@ def create( expected_artifact_bytes=expected_artifact_bytes, expected_artifact_set_digest=expected_artifact_set_digest, expected_cache_root=expected_cache_root, + artifact_root=validate_artifact_snapshot(model_manifest, artifact_root, cache_dir=cache_dir), ) memory_cache = MemoryCache(attn_cache_bytes, max_alloc_timeout, paged=paged_cache, page_size=page_size) diff --git a/src/drift/server/text_peer.py b/src/drift/server/text_peer.py index aeff1210f..917b3ad02 100644 --- a/src/drift/server/text_peer.py +++ b/src/drift/server/text_peer.py @@ -23,6 +23,7 @@ def __init__( *, initial_peers, cache_dir, + artifact_root=None, max_context_tokens=2048, max_output_tokens=512, request_timeout=900 @@ -33,6 +34,7 @@ def __init__( raise ValueError("Invalid text peer context/output limits") self.dht, self.identity, self.manifest = dht, identity, manifest self.initial_peers, self.cache_dir = initial_peers, cache_dir + self.artifact_root = artifact_root self.max_context_tokens, self.max_output_tokens = max_context_tokens, max_output_tokens self.request_timeout = request_timeout self._stop = threading.Event() @@ -60,6 +62,7 @@ def _run(self): self.manifest, initial_peers=self.initial_peers, cache_dir=self.cache_dir, + artifact_root=self.artifact_root, request_timeout=180, max_retries=1, )() diff --git a/tests/test_local_artifact_snapshot.py b/tests/test_local_artifact_snapshot.py new file mode 100644 index 000000000..8fb65b8e3 --- /dev/null +++ b/tests/test_local_artifact_snapshot.py @@ -0,0 +1,486 @@ +"""Actual loader seams with tiny synthetic artifacts and no model/DHT/network.""" + +import json +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +from test_model_manifest import create_block_plan_snapshot + +import drift +from drift.artifact_snapshot import validate_artifact_snapshot +from drift.cli import run_server, run_text_peer +from drift.model_manifest import ManifestArtifactVerifier, ManifestError, select_manifest_block_artifacts +from drift.node.loading import make_manifest_loader +from drift.server import from_pretrained, server +from drift.server.text_peer import TextPeerService + + +class ConstructionReached(Exception): + pass + + +@pytest.fixture(autouse=True) +def no_external_effects(monkeypatch): + def forbidden(*args, **kwargs): + raise AssertionError("network, identity or model construction is forbidden") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", forbidden) + monkeypatch.setattr("requests.Session.request", forbidden) + monkeypatch.setattr(ManifestArtifactVerifier, "_resumable_hub_download", forbidden) + monkeypatch.setattr(server, "DHT", forbidden) + monkeypatch.setattr(server.NodeIdentity, "ensure", forbidden) + monkeypatch.setattr(server.AutoDistributedConfig, "from_pretrained", forbidden) + monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", forbidden) + monkeypatch.setattr(drift.AutoDistributedModelForCausalLM, "from_pretrained", forbidden) + + +@pytest.fixture +def snapshot(tmp_path): + root = tmp_path / "artifacts" + root.mkdir() + manifest, payloads = create_block_plan_snapshot(root) + cache = tmp_path / "writable-cache" + path = tmp_path / "manifest.json" + path.write_text(manifest.canonical_json(), encoding="utf-8") + return SimpleNamespace(root=root, cache=cache, manifest=manifest, payloads=payloads, path=path) + + +def server_kwargs(s): + return dict( + initial_peers=[], + dht_prefix=s.manifest.dht_prefix, + converted_model_name_or_path=s.manifest.source.repository, + throughput=0.01, + model_manifest=s.manifest, + revision=s.manifest.source.revision, + artifact_root=str(s.root), + cache_dir=str(s.cache), + ) + + +def construct_from_snapshot(s, monkeypatch, entry): + if entry == "server": + return server.Server(**server_kwargs(s)) + if entry == "client": + return make_manifest_loader(s.manifest, initial_peers=[], cache_dir=str(s.cache), artifact_root=str(s.root))() + if entry == "module": + return server.ModuleContainer.create(**module_kwargs(s)) + monkeypatch.setattr( + sys, + "argv", + [ + "text-peer", + str(s.path), + "--cache_dir", + str(s.cache), + "--artifact_root", + str(s.root), + "--identity_path", + "unused.key", + "--initial_peers", + "unused", + ], + ) + return run_text_peer.main() + + +@pytest.mark.parametrize("entry", ["server", "client", "text-cli", "module"]) +@pytest.mark.parametrize("corrupt", [False, True]) +def test_complete_snapshot_rejected_before_constructors(snapshot, monkeypatch, entry, corrupt): + s = snapshot + # Neither this unassigned shard nor tokenizer is worker startup metadata. + target = s.root / ("layer-2.safetensors" if entry in {"server", "module"} else "tokenizer.json") + if corrupt: + target.write_bytes(b"X" * target.stat().st_size) + else: + target.unlink() + with pytest.raises(ManifestError): + construct_from_snapshot(s, monkeypatch, entry) + assert not s.cache.exists() + + +def create_symlink_or_skip(path, target, *, directory=False): + try: + path.symlink_to(target, target_is_directory=directory) + except OSError as exc: + pytest.skip(f"Host cannot create symbolic links: {exc}") + + +@pytest.mark.parametrize("entry", ["server", "client", "text-cli", "module"]) +@pytest.mark.parametrize("violation", ["escaping-declared-symlink", "tokenizer-sidecar", "nested-sidecar"]) +def test_snapshot_tree_rejected_before_constructors(snapshot, monkeypatch, entry, violation): + s = snapshot + if violation == "escaping-declared-symlink": + external = s.root.parent / "external-shard" + external.write_bytes(s.payloads["shared.safetensors"]) + declared = s.root / "shared.safetensors" + declared.unlink() + create_symlink_or_skip(declared, external) + # The old hash-only preflight accepted this mutable external target. + s.manifest.verify_artifacts(s.root) + message = "links or reparse points" + elif violation == "tokenizer-sidecar": + (s.root / "tokenizer_config.json").write_text('{"tokenizer_class":"UnexpectedTokenizer"}') + message = "Undeclared artifact snapshot file" + else: + templates = s.root / "additional_chat_templates" + templates.mkdir() + (templates / "default.jinja").write_text("undeclared template") + message = "Undeclared artifact snapshot directory" + with pytest.raises(ManifestError, match=message): + construct_from_snapshot(s, monkeypatch, entry) + assert not s.cache.exists() + + +@pytest.mark.parametrize( + "filename", ["special_tokens_map.json", "generation_config.json", "model.safetensors", "chat_template.jinja"] +) +def test_snapshot_rejects_other_undeclared_loader_inputs(snapshot, monkeypatch, filename): + s = snapshot + (s.root / filename).write_bytes(b"synthetic undeclared sidecar") + with pytest.raises(ManifestError, match="Undeclared artifact snapshot file"): + construct_from_snapshot(s, monkeypatch, "client") + + +def test_snapshot_rejects_linked_directory_without_traversing_it(snapshot): + s = snapshot + external = s.root.parent / "external-directory" + external.mkdir() + create_symlink_or_skip(s.root / "linked-directory", external, directory=True) + with pytest.raises(ManifestError, match="links or reparse points"): + validate_artifact_snapshot(s.manifest, s.root, cache_dir=s.cache) + + +def test_snapshot_rejects_undeclared_empty_directory(snapshot): + s = snapshot + (s.root / "extra-directory").mkdir() + with pytest.raises(ManifestError, match="Undeclared artifact snapshot directory"): + validate_artifact_snapshot(s.manifest, s.root, cache_dir=s.cache) + + +def test_snapshot_accepts_declared_nested_regular_files(snapshot): + s = snapshot + nested = s.root / "tokenizer" / "nested" / "tokenizer.json" + nested.parent.mkdir(parents=True) + (s.root / "tokenizer.json").rename(nested) + manifest = replace( + s.manifest, + artifacts=tuple( + replace(artifact, path="tokenizer/nested/tokenizer.json") if artifact.path == "tokenizer.json" else artifact + for artifact in s.manifest.artifacts + ), + ) + assert validate_artifact_snapshot(manifest, s.root, cache_dir=s.cache) == s.root.resolve() + assert not s.cache.exists() + + +def test_real_server_constructor_uses_local_config(snapshot, monkeypatch): + s = snapshot + + def config(source, **kwargs): + assert source == s.root.resolve() + assert kwargs["local_files_only"] is True + assert kwargs["revision"] is None + raise ConstructionReached + + monkeypatch.setattr(server.AutoDistributedConfig, "from_pretrained", config) + with pytest.raises(ConstructionReached): + server.Server(**server_kwargs(s)) + assert set(p.name for p in s.root.iterdir()) == set(s.payloads) + + +@pytest.mark.parametrize("overlap", ["same", "cache-inside", "snapshot-inside"]) +def test_snapshot_and_writable_cache_must_be_disjoint(snapshot, overlap): + s = snapshot + cache = {"same": s.root, "cache-inside": s.root / "cache", "snapshot-inside": s.root.parent}[overlap] + with pytest.raises(ManifestError, match="disjoint"): + validate_artifact_snapshot(s.manifest, s.root, cache_dir=cache) + + +def test_snapshot_requires_manifest_and_explicit_cache(snapshot): + s = snapshot + with pytest.raises(ManifestError, match="exact model manifest"): + server.Server(**dict(server_kwargs(s), model_manifest=None)) + with pytest.raises(ManifestError, match="separate writable"): + validate_artifact_snapshot(s.manifest, s.root, cache_dir=None) + + +def test_snapshot_rejects_cache_file_before_model_construction(snapshot): + s = snapshot + s.cache.write_bytes(b"not a directory") + with pytest.raises(ManifestError, match="cache directory"): + server.Server(**server_kwargs(s)) + + +@pytest.mark.parametrize("field", ["converted_model_name_or_path", "revision"]) +def test_server_rejects_snapshot_loading_for_another_source_identity(snapshot, field): + s = snapshot + replacement = "other/repository" if field == "converted_model_name_or_path" else "b" * 40 + with pytest.raises(ManifestError): + server.Server(**dict(server_kwargs(s), **{field: replacement})) + + +def module_kwargs(s): + plan = select_manifest_block_artifacts( + s.manifest, + block_prefix="model.layers", + start_block=0, + end_block=2, + weight_map=json.loads(s.payloads["model.safetensors.index.json"])["weight_map"], + ) + return dict( + dht=None, + dht_prefix=s.manifest.dht_prefix, + converted_model_name_or_path=s.manifest.source.repository, + block_config=SimpleNamespace(block_prefix="model.layers"), + attn_cache_bytes=1, + server_info=SimpleNamespace(), + model_info=SimpleNamespace(), + block_indices=[0, 1], + min_batch_size=1, + max_batch_size=1, + max_chunk_size_bytes=1, + max_alloc_timeout=1, + torch_dtype=None, + cache_dir=str(s.cache.resolve()), + max_disk_space=None, + artifact_root=s.root, + device="cpu", + compression=None, + update_period=1, + expiration=None, + revision=s.manifest.source.revision, + token=False, + quant_type=None, + tensor_parallel_devices=(None,), + model_manifest=s.manifest, + expected_manifest_digest=s.manifest.digest_id, + expected_block_indices="0:2", + expected_artifact_bytes=plan.artifact_bytes, + expected_artifact_set_digest=plan.artifact_set_digest, + expected_cache_root=str(s.cache.resolve()), + ) + + +def test_real_module_constructor_preserves_root_and_all_five_claims(snapshot, monkeypatch): + s = snapshot + actual = server._scoped_manifest_artifact_verifier + + def scoped(manifest, **kwargs): + verifier = actual(manifest, **kwargs) + assert verifier.artifact_root == s.root.resolve() + assert verifier.cache_only is True + assert verifier.cache_dir == str(s.cache.resolve()) + assert verifier.allowed_paths == frozenset( + {"config.json", "model.safetensors.index.json", "shared.safetensors"} + ) + raise ConstructionReached + + monkeypatch.setattr(server, "_scoped_manifest_artifact_verifier", scoped) + with pytest.raises(ConstructionReached): + server.ModuleContainer.create(**module_kwargs(s)) + for field, value in [ + ("expected_manifest_digest", "sha256:" + "0" * 64), + ("expected_block_indices", "1:2"), + ("expected_artifact_bytes", 1), + ("expected_artifact_set_digest", "0" * 64), + ("expected_cache_root", str(s.cache.parent / "wrong")), + ]: + with pytest.raises(ManifestError): + server.ModuleContainer.create(**dict(module_kwargs(s), **{field: value})) + + +def test_real_block_file_loader_locks_only_writable_cache(snapshot, monkeypatch): + s = snapshot + verifier = server._scoped_manifest_artifact_verifier( + s.manifest, + repository=s.manifest.source.repository, + revision=s.manifest.source.revision, + token=False, + cache_dir=str(s.cache), + max_disk_space=None, + block_prefix="model.layers", + block_indices=[0, 1], + artifact_root=s.root, + ) + + def deserialize(path, **kwargs): + assert Path(path).parent == s.root + assert Path(path).read_bytes() == s.payloads["shared.safetensors"] + assert (s.cache / "blocks.lock").is_file() + return "synthetic state dict" + + monkeypatch.setattr(from_pretrained, "_load_state_dict_from_local_file", deserialize) + result = from_pretrained._load_state_dict_from_repo_file( + s.manifest.source.repository, "shared.safetensors", cache_dir=str(s.cache), artifact_verifier=verifier + ) + assert result == "synthetic state dict" + assert set(p.name for p in s.root.iterdir()) == set(s.payloads) + + +def fake_client_constructors(monkeypatch, s): + calls = [] + + def tokenizer(root, **kwargs): + assert root == s.root.resolve() and kwargs["local_files_only"] is True + calls.append("tokenizer") + return object() + + def model(repository, **kwargs): + verifier = kwargs["artifact_verifier"] + assert repository == s.manifest.source.repository + assert kwargs["revision"] == s.manifest.source.revision + assert kwargs["manifest_digest"] == s.manifest.digest + assert verifier.artifact_root == s.root.resolve() and verifier.cache_dir == str(s.cache) + assert verifier.cache_only is True + verifier.ensure_path("shared.safetensors", allowed_roles={"weight"}) + calls.append("model") + return SimpleNamespace( + transformer=SimpleNamespace( + h=SimpleNamespace( + sequence_manager=SimpleNamespace( + shutdown=lambda: calls.append("shutdown"), + dht=None, + _thread=SimpleNamespace(is_alive=lambda: False), + ) + ) + ) + ) + + monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", tokenizer) + monkeypatch.setattr(drift.AutoDistributedModelForCausalLM, "from_pretrained", model) + return calls + + +@pytest.mark.parametrize("service", [False, True]) +def test_actual_client_loader_and_text_service_use_snapshot(snapshot, monkeypatch, service): + s = snapshot + calls = fake_client_constructors(monkeypatch, s) + if service: + identity = SimpleNamespace(peer_id="synthetic-peer") + peer = TextPeerService( + identity, identity, s.manifest, initial_peers=[], cache_dir=str(s.cache), artifact_root=s.root + ) + peer._stop.set() # Exercise the actual loader without starting protocol/discovery. + peer._run() + assert peer.error is None + else: + runtime = make_manifest_loader(s.manifest, initial_peers=[], cache_dir=str(s.cache), artifact_root=s.root)() + runtime.close() + assert calls == ["tokenizer", "model", "shutdown"] + assert not s.cache.exists() + + +def test_server_forwards_snapshot_and_claims_into_actual_module_constructor(snapshot, monkeypatch): + s = snapshot + owner = server.Server.__new__(server.Server) + values = module_kwargs(s) + for name, value in values.items(): + setattr(owner, name, value) + for name in ( + "num_handlers", + "inference_max_length", + "stats_report_interval", + "request_timeout", + "session_timeout", + "step_timeout", + "prefetch_batches", + "sender_threads", + "admission_policy", + "protocol_identity", + "manifest_execution_profile", + "revocations", + "announcement_replay_guard", + "ready_timeout", + "health_state_path", + "paged_cache", + "page_size", + ): + setattr(owner, name, None) + owner.processing_budget = SimpleNamespace(percent=100, path=None) + actual = server._scoped_manifest_artifact_verifier + + def scoped(manifest, **kwargs): + verifier = actual(manifest, **kwargs) + assert verifier.artifact_root == s.root.resolve() and verifier.cache_only + assert verifier.cache_dir == values["expected_cache_root"] + raise ConstructionReached + + monkeypatch.setattr(server, "_scoped_manifest_artifact_verifier", scoped) + with pytest.raises(ConstructionReached): + owner._create_module_container([0, 1]) + + +def test_worker_cli_forwards_snapshot_and_managed_placement_claims(snapshot, monkeypatch): + s = snapshot + claims = module_kwargs(s) + argv = [ + s.manifest.source.repository, + "--new_swarm", + "--model_manifest", + str(s.path), + "--artifact_root", + str(s.root), + "--cache_dir", + str(s.cache.resolve()), + "--block_indices", + "0:2", + "--identity_path", + "unused.key", + "--increase_file_limit", + "0", + ] + fields = ( + "expected_manifest_digest", + "expected_block_indices", + "expected_artifact_bytes", + "expected_artifact_set_digest", + "expected_cache_root", + ) + for field in fields: + argv.extend(["--" + field, str(claims[field])]) + args = vars(run_server.build_parser(bound_worker=True).parse_args(argv)) + args.pop("config", None) + monkeypatch.setattr(run_server, "tie_child_processes_to_this_process", lambda: None) + monkeypatch.setattr(run_server, "log_version", lambda: None) + monkeypatch.setattr(run_server, "Server", lambda **kwargs: kwargs) + result = run_server.server_from_args(args) + assert result["artifact_root"] == str(s.root) + assert result["model_manifest"].digest == s.manifest.digest + for field in fields: + assert result[field] == claims[field] + + +def test_text_cli_forwards_root_and_separate_cache(snapshot, monkeypatch): + s = snapshot + monkeypatch.setattr( + sys, + "argv", + [ + "text-peer", + str(s.path), + "--cache_dir", + str(s.cache), + "--artifact_root", + str(s.root), + "--identity_path", + "unused.key", + "--initial_peers", + "unused", + ], + ) + monkeypatch.setattr(run_text_peer.NodeIdentity, "ensure", lambda path: object()) + monkeypatch.setattr(run_text_peer, "DHT", lambda **kwargs: object()) + + def service(*args, **kwargs): + assert kwargs["artifact_root"] == s.root.resolve() + assert kwargs["cache_dir"] == str(s.cache) + assert args[2].digest == s.manifest.digest + raise ConstructionReached + + monkeypatch.setattr(run_text_peer, "TextPeerService", service) + with pytest.raises(ConstructionReached): + run_text_peer.main() diff --git a/tests/test_model_manifest.py b/tests/test_model_manifest.py index de1e2e298..638b0a57e 100644 --- a/tests/test_model_manifest.py +++ b/tests/test_model_manifest.py @@ -855,6 +855,239 @@ def test_incremental_verifier_hashes_metadata_and_requested_shards(tmp_path): verifier.ensure_path("README.md") +def cache_only_test_manifest(payload: bytes = b"good") -> ModelManifest: + source = manifest_dict() + config = next(artifact for artifact in source["artifacts"] if artifact["role"] == "config") + config.update(sha256=hashlib.sha256(payload).hexdigest(), size=len(payload)) + return ModelManifest.from_dict(source) + + +def forbid_cache_only_transfer(monkeypatch, verifier): + def forbidden(*args, **kwargs): + pytest.fail("cache-only verification attempted a transfer or disk reservation") + + monkeypatch.setattr("requests.get", forbidden) + monkeypatch.setattr("drift.utils.hub_ranges.download_ranges", forbidden) + monkeypatch.setattr("drift.utils.disk_cache.free_disk_space_for", forbidden) + monkeypatch.setattr(verifier, "_resumable_hub_download", forbidden) + + +def test_cache_only_uses_local_hub_file_and_rehashes_on_repeated_access(tmp_path, monkeypatch): + manifest = cache_only_test_manifest() + cache_dir = tmp_path / "cache" + snapshot = tmp_path / "hub-snapshot" + snapshot.mkdir() + cached = snapshot / "config.json" + cached.write_bytes(b"good") + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=cache_dir, cache_only=True + ) + + def local_hub_file(repository, filename, **kwargs): + assert (repository, filename, kwargs["revision"]) == ( + manifest.source.repository, + "config.json", + manifest.source.revision, + ) + assert kwargs["local_files_only"] is True + return str(cached) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_hub_file) + forbid_cache_only_transfer(monkeypatch, verifier) + + assert verifier.ensure_path("config.json", allowed_roles={"config"}) == cached.absolute() + assert verifier.snapshot_root == snapshot.absolute() + assert not cache_dir.exists() + + original_stat = cached.stat() + cached.write_bytes(b"evil") # Same size and mtime must not reuse the earlier hash. + os.utime(cached, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns)) + with pytest.raises(ManifestError, match="declared SHA-256"): + verifier.ensure_path("config.json", allowed_roles={"config"}) + + +def test_cache_only_uses_verified_manifest_cache_final(tmp_path, monkeypatch): + from huggingface_hub.utils import LocalEntryNotFoundError + + manifest = cache_only_test_manifest() + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path, cache_only=True + ) + _, final, _ = verifier._resumable_paths(manifest.get_artifact("config.json")) + final.parent.mkdir(parents=True) + final.write_bytes(b"good") + + def local_miss(*args, **kwargs): + assert kwargs["local_files_only"] is True + raise LocalEntryNotFoundError("not cached in the Hub") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_miss) + forbid_cache_only_transfer(monkeypatch, verifier) + + assert verifier.ensure_path("config.json", allowed_roles={"config"}) == final.absolute() + assert verifier.snapshot_root == final.parent.absolute() + + +def test_cache_only_mixed_local_caches_share_one_snapshot_root(tmp_path, monkeypatch): + from huggingface_hub.utils import LocalEntryNotFoundError + + source = manifest_dict() + payloads = {"config.json": b"good", "tokenizer.json": b"tt"} + for artifact in source["artifacts"]: + if artifact["path"] in payloads: + payload = payloads[artifact["path"]] + artifact.update(sha256=hashlib.sha256(payload).hexdigest(), size=len(payload)) + manifest = ModelManifest.from_dict(source) + snapshot = tmp_path / "hub-snapshot" + snapshot.mkdir() + (snapshot / "config.json").write_bytes(payloads["config.json"]) + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path / "cache", cache_only=True + ) + _, custom_final, _ = verifier._resumable_paths(manifest.get_artifact("tokenizer.json")) + custom_final.parent.mkdir(parents=True) + custom_final.write_bytes(payloads["tokenizer.json"]) + + def local_hub_file(*args, **kwargs): + assert kwargs["local_files_only"] is True + if args[1] == "config.json": + return str(snapshot / "config.json") + raise LocalEntryNotFoundError("not cached in the Hub") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_hub_file) + forbid_cache_only_transfer(monkeypatch, verifier) + + assert verifier.ensure_path("config.json", allowed_roles={"config"}) == snapshot / "config.json" + assert verifier.ensure_path("tokenizer.json", allowed_roles={"tokenizer"}) == snapshot / "tokenizer.json" + assert verifier.snapshot_root == snapshot.absolute() + assert (snapshot / "tokenizer.json").read_bytes() == payloads["tokenizer.json"] + assert custom_final.read_bytes() == payloads["tokenizer.json"] + + +@pytest.mark.parametrize("partial_bytes", [None, b"go", b"good"]) +def test_cache_only_rejects_missing_final_even_with_partial(tmp_path, monkeypatch, partial_bytes): + from huggingface_hub.utils import LocalEntryNotFoundError + + manifest = cache_only_test_manifest() + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path, cache_only=True + ) + partial, final, _ = verifier._resumable_paths(manifest.get_artifact("config.json")) + if partial_bytes is not None: + partial.parent.mkdir(parents=True) + partial.write_bytes(partial_bytes) + + def local_miss(*args, **kwargs): + assert kwargs["local_files_only"] is True + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_miss) + forbid_cache_only_transfer(monkeypatch, verifier) + + with pytest.raises(ManifestError, match="not available in the local cache"): + verifier.ensure_path("config.json", allowed_roles={"config"}) + assert not final.exists() + if partial_bytes is not None: + assert partial.read_bytes() == partial_bytes + + +@pytest.mark.parametrize("cached_bytes", [b"go", b"evil"]) +def test_cache_only_rejects_truncated_or_corrupt_manifest_final(tmp_path, monkeypatch, cached_bytes): + from huggingface_hub.utils import LocalEntryNotFoundError + + manifest = cache_only_test_manifest() + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path, cache_only=True + ) + _, final, _ = verifier._resumable_paths(manifest.get_artifact("config.json")) + final.parent.mkdir(parents=True) + final.write_bytes(cached_bytes) + + def local_miss(*args, **kwargs): + assert kwargs["local_files_only"] is True + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_miss) + forbid_cache_only_transfer(monkeypatch, verifier) + + with pytest.raises(ManifestError, match="size .* expected|declared SHA-256"): + verifier.ensure_path("config.json", allowed_roles={"config"}) + assert final.read_bytes() == cached_bytes + + +def test_cache_only_explicit_artifact_root_never_falls_back_to_cache(tmp_path, monkeypatch): + manifest = cache_only_test_manifest() + root = tmp_path / "owner-snapshot" + root.mkdir() + verifier = ManifestArtifactVerifier( + manifest, + manifest.source.repository, + manifest.source.revision, + artifact_root=root, + cache_only=True, + ) + + def forbidden(*args, **kwargs): + pytest.fail("explicit artifact root attempted a cache lookup") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", forbidden) + forbid_cache_only_transfer(monkeypatch, verifier) + + with pytest.raises(ManifestError, match="Could not read artifact"): + verifier.ensure_path("config.json", allowed_roles={"config"}) + (root / "config.json").write_bytes(b"evil") + with pytest.raises(ManifestError, match="declared SHA-256"): + verifier.ensure_path("config.json", allowed_roles={"config"}) + + +def test_cache_only_rejects_direct_transfer_helper_call(tmp_path, monkeypatch): + manifest = cache_only_test_manifest() + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path, cache_only=True + ) + + def forbidden(*args, **kwargs): + pytest.fail("cache-only transfer helper attempted HTTP") + + monkeypatch.setattr("requests.get", forbidden) + monkeypatch.setattr("drift.utils.hub_ranges.download_ranges", forbidden) + + with pytest.raises(ManifestError, match="transfer is disabled in cache-only mode"): + verifier._resumable_hub_download(manifest.get_artifact("config.json")) + assert not (tmp_path / "manifest-artifacts").exists() + + +def test_default_verifier_still_downloads_after_local_cache_miss(tmp_path, monkeypatch): + from huggingface_hub.utils import LocalEntryNotFoundError + + manifest = cache_only_test_manifest() + verifier = ManifestArtifactVerifier( + manifest, manifest.source.repository, manifest.source.revision, cache_dir=tmp_path + ) + reservations = [] + transfers = [] + + def local_miss(*args, **kwargs): + assert kwargs["local_files_only"] is True + raise LocalEntryNotFoundError("not cached") + + def local_transfer(artifact, *, destination=None): + transfers.append((artifact.path, destination)) + _, final, _ = verifier._resumable_paths(artifact) + final.parent.mkdir(parents=True, exist_ok=True) + final.write_bytes(b"good") + return str(final) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", local_miss) + monkeypatch.setattr("drift.utils.disk_cache.free_disk_space_for", lambda size, **kwargs: reservations.append(size)) + monkeypatch.setattr(verifier, "_resumable_hub_download", local_transfer) + + result = verifier.ensure_path("config.json", allowed_roles={"config"}) + assert result.read_bytes() == b"good" + assert reservations == [4] + assert transfers == [("config.json", None)] + + def test_interrupted_download_can_resume_and_is_reverified(tmp_path, monkeypatch): """A real interrupted HTTP response resumes with Range and stays hidden until verified.""" import socket