diff --git a/CLAUDE.md b/CLAUDE.md index 4fb1e92..5e3fe9d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co Rootstock runs MLIP (Machine Learning Interatomic Potential) calculators in isolated pre-built Python environments on HPC clusters, communicating via the i-PI protocol over Unix sockets. -Versioning is dynamic (git tags via uv-dynamic-versioning) — check `rootstock --version`. Manifest schema v6 (older schemas migrate in place on load; verification is per-cluster, and shared installs like sophia/polaris push one manifest per cluster); canonical-checkpoint-id API. +Versioning is dynamic (git tags via uv-dynamic-versioning) — check `rootstock --version`. Manifest schema v7 (older schemas migrate in place on load; verification is per-cluster, and shared installs like sophia/polaris push one manifest per cluster); canonical-checkpoint-id API. ## Commands diff --git a/docs/cluster-setup.md b/docs/cluster-setup.md index afc7f40..cbd520c 100644 --- a/docs/cluster-setup.md +++ b/docs/cluster-setup.md @@ -67,6 +67,32 @@ The cluster registry (`rootstock/clusters.py`) is only a name → install-path b Users don't need to set environment variables — `RootstockCalculator(cluster="perlmutter", ...)` resolves both automatically. +### Node-local staging (optional) + +On network filesystems with slow cold reads, worker startup can be dominated by faulting multi-GB libraries and weights in over the wire. When the compute nodes have local disk, an install can opt into **node-local staging**: `rootstock install` (and the `rootstock pack` backfill command) archives each built env into a single compressed image under `{root}/images/`, and worker spawns extract that image to local disk — one sequential read — then run entirely from the local copy. Checkpoint weights ride along via each checkpoint's recorded weight files. Nothing is required of end users; spawns fall back to the ordinary warm-up path whenever any piece is missing. + +Enable it by declaring where staged copies may live, in `{root}/layout.json`: + +```bash +rootstock init --stage-dir '/tmp' # or e.g. '$SLURM_TMPDIR' +``` + +(or add `"stage_dir": "/tmp"` to an existing `layout.json` by hand — like `cache_root`, `init` is the only command with a flag to *set* it; `install`/`sync`/`prune` rewrite the marker but always preserve an existing declaration). The value may contain environment variables, which expand on the node at spawn time; a value that doesn't expand (e.g. `$SLURM_TMPDIR` on a login node), doesn't exist, isn't writable, or is on the *same* filesystem as the install simply disables staging on that node. Users can override per run with `ROOTSTOCK_STAGE_DIR=...` or disable with `ROOTSTOCK_NO_STAGE=1`. + +Packing and extraction shell out to `tar` and `zstd`, so both must be on `PATH` (compute nodes usually have them; `module load zstd` otherwise). After enabling, backfill images for already-built envs — in a batch allocation if the login nodes cap CPU time: + +```bash +rootstock pack --cluster +``` + +Job scripts that spawn several calculators can pay the read once, up front: + +```bash +rootstock stage uma-s-1p1 mace-mp-0-medium --cluster +``` + +Staged copies are per-user, content-addressed by image, reused across spawns and jobs while the node directory survives, and evicted oldest-first when local disk runs short. + ### Trust model Using a shared install means trusting its maintainer. An environment's diff --git a/rootstock/batch.py b/rootstock/batch.py index 6fb50b7..b76b368 100644 --- a/rootstock/batch.py +++ b/rootstock/batch.py @@ -15,8 +15,9 @@ The prune planner takes the opposite half of the same diff — ``actual − desired`` — plus an internal-GC tier that needs no declared state at all -(``.build`` leftovers, orphaned interpreters, stale lockfiles, ``.trash``, -the uv cache). Same keep-going executor idiom, same single end-of-run +(``.build`` leftovers, orphaned interpreters, stale lockfiles, orphaned +staging images, ``.trash``, the uv cache). Same keep-going executor idiom, +same single end-of-run manifest refresh + push. Deletion orders are chosen so a crash mid-run strands only *unreferenced* state that a later run collects. """ @@ -677,7 +678,9 @@ class PruneEnvItem: class GCItem: """One piece of internal garbage; ``path`` is absolute.""" - kind: str # "build" | "trash" | "interpreter" | "lockfile" | "uv-cache" | "unattributed" + # "build" | "trash" | "interpreter" | "lockfile" | "image" | "uv-cache" + # | "unattributed" + kind: str path: str reason: str reclaim_bytes: int | None = None # None: unknown until executed (uv-cache) @@ -1118,6 +1121,32 @@ def collect(kind: str, entries: list[Path], reason_for: Callable[[Path], str]) - ] collect("lockfile", orphan_locks, lambda _: "no matching environment source") + images_dir = root / "images" + if images_dir.is_dir(): + # A staging archive is live iff some manifest env record points at + # it (#180) — retiring an env drops its record (refresh) but leaves + # the multi-GB archive; crashed packs leave .packing partials. The + # age guard in collect() protects in-flight packs and images whose + # manifest record is about to be written. + recorded_images = set() + for record in manifest.environments.values() if manifest else []: + image = record.image if isinstance(record.image, dict) else None + image_path = image.get("path") if image else None + if isinstance(image_path, str) and image_path: + recorded_images.add(Path(image_path).name) + orphan_images = [ + entry for entry in sorted(images_dir.iterdir()) if entry.name not in recorded_images + ] + collect( + "image", + orphan_images, + lambda entry: ( + "leftover from a crashed pack" + if ".packing." in entry.name + else "no manifest env records it" + ), + ) + for kind, count in sorted(age_skipped.items()): plan.notes.append( f"{kind}: left {count} entr{'y' if count == 1 else 'ies'} younger than " diff --git a/rootstock/cli.py b/rootstock/cli.py index ba7e500..897ba90 100644 --- a/rootstock/cli.py +++ b/rootstock/cli.py @@ -43,12 +43,26 @@ The subtractive half of sync: remove built envs with no registered source, checkpoint records (and their unshared weight files) no source declares, and internal garbage (.build/.trash leftovers, - orphaned interpreters, stale lockfiles, the uv cache). Plan-confirm - by default; batch jobs pass --yes. + orphaned interpreters, stale lockfiles, orphaned staging images, + the uv cache). Plan-confirm by default; batch jobs pass --yes. rootstock prune --cluster delta --dry-run rootstock prune ./environments/ --yes # retire everything not declared here rootstock prune --gc-only --yes # internal garbage only + rootstock pack [ ...] [--all] [--root | --cluster ] + Pack built envs into single-image archives ({root}/images/) for + node-local staging (#180). Bare `pack` covers every built env whose + image is missing or stale; `install` packs its own env automatically. + rootstock pack --cluster delta # backfill missing/stale + rootstock pack mace uma # exactly these + + rootstock stage [...] [--root | --cluster ] + Warm checkpoints once, ahead of worker spawns — for job prologues. + Extracts each hosting env's image (plus recorded weights) to the + configured node-local dir, or falls back to a page-cache prewarm + where staging isn't configured. + rootstock stage uma-s-1p1 mace-mp-0-medium --cluster delta + rootstock benchmark [--root ] [--checkpoints ...] [--devices cuda cpu] [--list] Measure i-PI IPC overhead: RootstockCalculator vs. the same calculator called directly inside its pre-built env. `--list` shows installed ids. @@ -93,11 +107,13 @@ cmd_manifest_push, cmd_manifest_show, cmd_new_env, + cmd_pack, cmd_prune, cmd_resolve, cmd_serve, cmd_setup_perms, cmd_smoke_test, + cmd_stage, cmd_status, cmd_sync, cmd_usage_compact, @@ -145,6 +161,16 @@ def main(): "layout.json and moving the weights." ), ) + init_parser.add_argument( + "--stage-dir", + default=None, + help=( + "Node-local directory worker spawns may stage packed env images " + "to (recorded in {root}/layout.json). May contain env vars that " + "expand on the compute node, e.g. '$SLURM_TMPDIR' or '/tmp'. " + "Pass an empty string to clear an existing declaration." + ), + ) init_parser.add_argument( "--skip-dirs", action="store_true", @@ -227,6 +253,15 @@ def main(): action="store_true", help="Skip the up-front shared-install permission check", ) + install_parser.add_argument( + "--no-pack", + action="store_true", + help=( + "Skip packing the staging image after the build (worker spawns " + "then warm this env via the prewarm path; pack later with " + "'rootstock pack')" + ), + ) install_parser.set_defaults(func=cmd_install) # add command @@ -432,7 +467,8 @@ def main(): "checkpoint records no source declares (plus their weight files, " "refcounted against surviving checkpoints), and internal garbage " "(.build/.trash leftovers, orphaned interpreters, stale " - "lockfiles, the uv cache). Prints the plan and asks for " + "lockfiles, orphaned staging images, the uv cache). Prints the " + "plan and asks for " "confirmation before deleting anything; idempotent and safe to " "re-run after a failure. Don't run while a sync is in flight." ), @@ -533,6 +569,77 @@ def main(): ) prune_parser.set_defaults(func=cmd_prune) + # pack command + pack_parser = subparsers.add_parser( + "pack", + help="Pack built envs into single-image archives for node-local staging", + description=( + "Pack built envs into {root}/images/-.tar.zst and " + "record the images in the manifest (#180). Worker spawns on " + "clusters with a configured staging dir extract these to " + "node-local disk instead of prewarming the shared tree. Bare " + "'pack' covers every built env whose image is missing or stale " + "(installs pack their own env automatically). Needs tar and " + "zstd on PATH; on login nodes with CPU-time caps, run inside a " + "batch allocation." + ), + ) + pack_parser.add_argument( + "envs", + nargs="*", + metavar="ENV", + help="Env name(s) to pack (default: every built env with a missing/stale image)", + ) + pack_parser.add_argument( + "--all", + action="store_true", + help="Repack every built env, fresh or not", + ) + pack_parser.add_argument( + "--root", + default=os.environ.get(ROOTSTOCK_ROOT_ENV), + help=f"Root directory (default: ${ROOTSTOCK_ROOT_ENV})", + ) + pack_parser.add_argument( + "--cluster", + help="Resolve the root from the cluster registry instead of --root", + ) + pack_parser.add_argument( + "--no-push", + action="store_true", + help="Don't push manifest to backend (useful during development)", + ) + pack_parser.set_defaults(func=cmd_pack) + + # stage command + stage_parser = subparsers.add_parser( + "stage", + help="Warm checkpoints once ahead of worker spawns (job prologues)", + description=( + "Stage each checkpoint's env image and recorded weights to the " + "configured node-local dir — or page-cache prewarm them where " + "staging isn't configured — so the shared-filesystem read " + "happens once, before any worker spawns. Intended for job " + "prologues; worker spawns do the same on their own at startup." + ), + ) + stage_parser.add_argument( + "checkpoints", + nargs="+", + metavar="CHECKPOINT", + help="Canonical checkpoint id(s) to warm", + ) + stage_parser.add_argument( + "--root", + default=os.environ.get(ROOTSTOCK_ROOT_ENV), + help=f"Root directory (default: ${ROOTSTOCK_ROOT_ENV})", + ) + stage_parser.add_argument( + "--cluster", + help="Resolve the root from the cluster registry instead of --root", + ) + stage_parser.set_defaults(func=cmd_stage) + # smoke-test command smoke_parser = subparsers.add_parser( "smoke-test", diff --git a/rootstock/clusters.py b/rootstock/clusters.py index 9322381..f33903d 100644 --- a/rootstock/clusters.py +++ b/rootstock/clusters.py @@ -23,10 +23,18 @@ class Cluster: `cache_root` defaults to `root`. Override it only when the right filesystem for code/venvs differs from the right filesystem for the model-weight cache. + + `stage_dir` names a node-local directory worker spawns may stage packed + env images to (#180); like `cache_root`, it is only a legacy fallback — + the install's own `stage_dir` declaration in `{root}/layout.json` wins. + The value is a raw string because it may contain environment variables + (`$SLURM_TMPDIR`) that expand on the compute node at spawn time. None + (the default) disables staging via the registry. """ root: Path cache_root: Path | None = None + stage_dir: str | None = None @property def resolved_cache_root(self) -> Path: diff --git a/rootstock/commands/__init__.py b/rootstock/commands/__init__.py index 0d30293..1c95a95 100644 --- a/rootstock/commands/__init__.py +++ b/rootstock/commands/__init__.py @@ -7,11 +7,13 @@ from .init import cmd_init from .install import cmd_install from .manifest import cmd_manifest_init, cmd_manifest_push, cmd_manifest_show +from .pack import cmd_pack from .prune import cmd_prune from .resolve import cmd_resolve from .serve import cmd_serve from .setup_perms import cmd_setup_perms from .smoke_test import cmd_smoke_test +from .stage import cmd_stage from .status import cmd_list, cmd_status from .sync import cmd_sync from .usage import cmd_usage_compact, cmd_usage_push, cmd_usage_report @@ -27,11 +29,13 @@ "cmd_manifest_push", "cmd_manifest_show", "cmd_new_env", + "cmd_pack", "cmd_prune", "cmd_resolve", "cmd_serve", "cmd_setup_perms", "cmd_smoke_test", + "cmd_stage", "cmd_status", "cmd_sync", "cmd_usage_compact", diff --git a/rootstock/commands/init.py b/rootstock/commands/init.py index 6532b2c..f79c065 100644 --- a/rootstock/commands/init.py +++ b/rootstock/commands/init.py @@ -202,9 +202,12 @@ def cmd_init(args) -> int: print(f" Exists: {dir_path}") try: - # Declare the cache root so the install is self-describing — - # readers prefer this over the baked-in cluster registry. - write_layout_marker(root, cache_root=cache_root) + # Declare the cache root (and any node-local staging base) so + # the install is self-describing — readers prefer this over the + # baked-in cluster registry. + write_layout_marker( + root, cache_root=cache_root, stage_dir=getattr(args, "stage_dir", None) + ) print(f" Created: {root / 'layout.json'}") except PermissionError: print(f" Skipped (no permission): {root / 'layout.json'}") diff --git a/rootstock/commands/install.py b/rootstock/commands/install.py index 3fa3ab2..b76eb77 100644 --- a/rootstock/commands/install.py +++ b/rootstock/commands/install.py @@ -26,6 +26,7 @@ def _install_one(root: Path, source: str, args) -> int: upgrade=args.upgrade, verbose=args.verbose, push=not args.no_push, + pack=not getattr(args, "no_pack", False), progress=print, ) except OperationError as exc: diff --git a/rootstock/commands/pack.py b/rootstock/commands/pack.py new file mode 100644 index 0000000..f81e969 --- /dev/null +++ b/rootstock/commands/pack.py @@ -0,0 +1,61 @@ +"""``rootstock pack`` — pack built envs into single-image archives (#180). + +Thin argparse adapter over :func:`rootstock.operations.pack_environments`. +``install`` packs each env it builds; this is the backfill for envs built +before packing existed (or whose install-time pack failed — no zstd on +PATH, say). On clusters with login-node CPU-time caps (Delta), run it in a +batch allocation: zstd across a multi-GB env is exactly the kind of burst +those caps kill. +""" + +from __future__ import annotations + +import os +import sys + +from ..layout import ensure_layout_compatible +from ..manifest import ManifestError +from ..operations import OperationError, pack_environments +from .common import resolve_root + + +def cmd_pack(args) -> int: + """ + Pack staging images and record them in the manifest. + + Exit codes: + 0: Every requested image packed (or nothing needed packing) + 1: One or more packs failed + """ + # Images land on the shared install; same umask stance as install/sync. + os.umask(0o002) + + root = resolve_root(args) + try: + ensure_layout_compatible(root) + except RuntimeError as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + + if args.all and args.envs: + print("Error: name envs or pass --all, not both.", file=sys.stderr) + return 2 + + env_names = args.envs or None + if args.all: + from ..environment import list_built_environments + + env_names = [name for name, _ in list_built_environments(root)] + if not env_names: + print(f"No built envs at {root}.", file=sys.stderr) + return 1 + + try: + packed = pack_environments(root, env_names, push=not args.no_push, progress=print) + except (OperationError, ManifestError) as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + + if packed: + print(f"\nPacked {len(packed)} image(s) into {root / 'images'}.") + return 0 diff --git a/rootstock/commands/stage.py b/rootstock/commands/stage.py new file mode 100644 index 0000000..086f448 --- /dev/null +++ b/rootstock/commands/stage.py @@ -0,0 +1,92 @@ +"""``rootstock stage`` — warm checkpoints once, ahead of worker spawns. + +The job-prologue command the prewarm ladder called for (#179, subsumed by +node-local staging #180): an sbatch/PBS script stages (or, where staging +isn't configured, page-cache prewarms) the checkpoints it is about to run, +paying the shared-filesystem read exactly once before any worker spawns. +On compute nodes the staged copy / warmth then persists for the job:: + + rootstock stage uma-s-1p1 mace-mp-0-medium --cluster delta + +Spawns do all of this on their own; the command exists for explicit intent — +warm N checkpoints up front rather than serially on first use. +""" + +from __future__ import annotations + +import sys + +from .. import prewarm +from ..environment import ( + CheckpointNotFoundError, + get_checkpoint_prewarm_paths, + resolve_checkpoint, +) +from ..layout import resolve_cache_root +from ..stage import resolve_stage_base, stage_env, stage_weights +from .common import resolve_root + + +def cmd_stage(args) -> int: + """ + Stage (or prewarm) each checkpoint's env and weights. + + Exit codes: + 0: Every checkpoint staged or prewarmed + 1: One or more checkpoints failed to resolve or warm + """ + root = resolve_root(args) + cluster = getattr(args, "cluster", None) + cache_root = resolve_cache_root(root) + base = resolve_stage_base(root) + if base is None: + print( + "Node-local staging is not configured here " + "(no ROOTSTOCK_STAGE_DIR / layout.json stage_dir); " + "falling back to a page-cache prewarm pass." + ) + + failures = 0 + for checkpoint_id in args.checkpoints: + try: + resolved = resolve_checkpoint(root, checkpoint_id, cluster) + except CheckpointNotFoundError as exc: + print(f"Error: {exc}", file=sys.stderr) + failures += 1 + continue + env_name = resolved.env_name + env_dir = root / "envs" / env_name + if not (env_dir / "bin" / "python").exists(): + print(f"Error: env '{env_name}' is not built at {env_dir}.", file=sys.stderr) + failures += 1 + continue + + print(f"{checkpoint_id} (env {env_name}):") + staged_root = stage_env(root, env_name, base) if base is not None else None + # Weights mirror only alongside a staged env — that is the only copy + # spawns will point a worker's caches at (stage_for_spawn returns + # None outright when the env can't stage). + staged_weights = None + if base is not None and staged_root is not None and not resolved.is_custom: + staged_weights = stage_weights(root, cache_root, env_name, checkpoint_id, base) + + # Whatever didn't stage gets the classic sequential warm instead — + # through prewarm_from_spec, so the prologue keeps the cgroup + # working-set warning and ROOTSTOCK_NO_PREWARM semantics. + spec: dict = {} + if staged_root is None: + spec["env_dir"] = str(env_dir) + if staged_weights is None and not resolved.is_custom: + try: + paths, tier = get_checkpoint_prewarm_paths( + root, env_name, checkpoint_id, cache_root + ) + except Exception: + paths, tier = [], None + spec["prewarm_paths"] = paths + if tier is not None: + spec["prewarm_weights_tier"] = tier + if spec.get("env_dir") or spec.get("prewarm_paths"): + prewarm.prewarm_from_spec(spec, log=sys.stdout, label="[Stage]") + + return 1 if failures else 0 diff --git a/rootstock/layout.py b/rootstock/layout.py index 7c3483c..dd07f9b 100644 --- a/rootstock/layout.py +++ b/rootstock/layout.py @@ -64,6 +64,20 @@ def read_declared_cache_root(root: Path) -> Path | None: return Path(declared) if isinstance(declared, str) and declared else None +def read_declared_stage_dir(root: Path) -> str | None: + """Return the node-local staging base this install declares, if any. + + The declaration lives in {root}/layout.json for the same reason + ``cache_root`` does: a registry baked into pinned clients goes stale, + while the install's own declaration travels with the install. The value + is kept as the raw string — it may contain environment variables + (``$SLURM_TMPDIR``, ``$TMPDIR``) that must expand on the *node* doing + the staging, not wherever the maintainer ran init. + """ + declared = _read_marker(root).get("stage_dir") + return declared if isinstance(declared, str) and declared else None + + def resolve_cache_root(root: Path, explicit: Path | str | None = None) -> Path: """Resolve the model-weight cache root for an install root. @@ -108,15 +122,21 @@ def ensure_layout_compatible(root: Path) -> None: ) -def write_layout_marker(root: Path, cache_root: Path | str | None = None) -> None: - """Record the layout version — and the install's cache root — in - {root}/layout.json. +def write_layout_marker( + root: Path, + cache_root: Path | str | None = None, + stage_dir: str | None = None, +) -> None: + """Record the layout version — and the install's cache root and staging + base — in {root}/layout.json. Called from maintainer commands that write the root anyway (install, init). ``cache_root`` records where this install keeps its model-weight - cache; when omitted, an existing declaration is preserved. No-op when - nothing would change, so repeated installs don't churn the file. Atomic - write, mode honoring the process umask — same recipe as save_manifest. + cache; ``stage_dir`` records the node-local base worker spawns may stage + envs to (#180). When omitted, existing declarations are preserved; pass + ``stage_dir=""`` to remove one. No-op when nothing would change, so + repeated installs don't churn the file. Atomic write, mode honoring the + process umask — same recipe as save_manifest. """ from . import __version__ from .manifest import now_iso @@ -127,8 +147,15 @@ def write_layout_marker(root: Path, cache_root: Path | str | None = None) -> Non existing = read_declared_cache_root(root) declared = str(existing) if existing is not None else None + declared_stage = stage_dir if stage_dir is not None else read_declared_stage_dir(root) + declared_stage = declared_stage or None # "" clears the declaration + current = _read_marker(root) - if current.get("layout_version") == LAYOUT_VERSION and current.get("cache_root") == declared: + if ( + current.get("layout_version") == LAYOUT_VERSION + and current.get("cache_root") == declared + and current.get("stage_dir") == declared_stage + ): return data = { @@ -138,6 +165,8 @@ def write_layout_marker(root: Path, cache_root: Path | str | None = None) -> Non } if declared is not None: data["cache_root"] = declared + if declared_stage is not None: + data["stage_dir"] = declared_stage root.mkdir(parents=True, exist_ok=True) fd, temp_path = tempfile.mkstemp(dir=root, suffix=".json") diff --git a/rootstock/manifest.py b/rootstock/manifest.py index 50772dd..b3501b3 100644 --- a/rootstock/manifest.py +++ b/rootstock/manifest.py @@ -28,7 +28,7 @@ from .config import UserConfig from .exceptions import RootstockError -SCHEMA_VERSION = 6 +SCHEMA_VERSION = 7 # manifest_lock defaults. The lock is only held across a load → mutate → save # cycle (plus the env refresh, which shells out to `uv pip list` per env), so @@ -172,6 +172,19 @@ def _migrate_v5_to_v6(data: dict) -> tuple[dict, str | None]: return data, note +def _migrate_v6_to_v7(data: dict) -> tuple[dict, str | None]: + """v7 added the optional per-env ``image`` record (packed env archive, + #180). + + ``image`` defaults to absent, so this is a pure version bump — the bump + matters for the same reason v5's did: without it a v6 client's + from_dict/asdict round-trip would silently drop the record on its next + save. + """ + data["schema_version"] = 7 + return data, None + + # One entry per historical schema version, upgrading one step. A schema bump # without a migration here strands every deployed manifest of that vintage — # add the migration in the same change as the bump. @@ -181,6 +194,7 @@ def _migrate_v5_to_v6(data: dict) -> tuple[dict, str | None]: 3: _migrate_v3_to_v4, 4: _migrate_v4_to_v5, 5: _migrate_v5_to_v6, + 6: _migrate_v6_to_v7, } @@ -345,6 +359,13 @@ class EnvironmentInfo: # refresh time. None = serves every cluster the install does; a list # restricts it (a cluster-specific variant on a shared install, #208). clusters: list[str] | None = None + # Packed single-image archive of this env (#180): a dict of {"path" + # (relative to the root), "sha256" (of the archive — the staging + # identity), "format", "compressed_bytes", "uncompressed_bytes", + # "packed_at"}. None = never packed. The image is current only while + # ``packed_at >= built_at`` (see :func:`image_is_current`) — a rebuild + # without a repack must never serve a stale image. + image: dict | None = None # Field order is the JSON key order (asdict): lock_hash stays ahead of # checkpoints to match the layout pushed manifests have always had. checkpoints: dict[str, CheckpointInfo] = field(default_factory=dict) @@ -372,6 +393,7 @@ def from_dict(cls, data: dict) -> EnvironmentInfo: checkpoints=checkpoints, lock_hash=data.get("lock_hash"), clusters=data.get("clusters"), + image=data.get("image"), ) @@ -384,6 +406,19 @@ def is_verified(env: EnvironmentInfo, ckpt: CheckpointInfo, cluster: str) -> boo return verified_at > env.built_at # ISO 8601 sorts lexically +def image_is_current(env: EnvironmentInfo) -> bool: + """True if the env's packed image record describes the current build. + + ``packed_at >= built_at`` (ISO 8601 sorts lexically; ``>=`` because an + install stamps both in the same refresh): an env rebuilt without a repack + reads as image-less rather than serving a stale archive. + """ + if not isinstance(env.image, dict): + return False + packed_at = env.image.get("packed_at") + return isinstance(packed_at, str) and packed_at >= env.built_at + + @dataclass class Manifest: """Root manifest for a rootstock installation. diff --git a/rootstock/operations.py b/rootstock/operations.py index 73def46..3bd1572 100644 --- a/rootstock/operations.py +++ b/rootstock/operations.py @@ -36,6 +36,7 @@ from .environment import ( declares_setup_from_path, find_env_for_checkpoint, + list_built_environments, parse_checkpoints_dict, parse_clusters_list, parse_custom_checkpoint_ids, @@ -55,6 +56,7 @@ now_iso, save_manifest, ) +from .pack import PackError, pack_environment, pack_environment_best_effort from .pep723 import ( get_dependencies, get_requires_python, @@ -162,7 +164,10 @@ def parse_setup_kwargs(kwarg_specs: list[str] | None) -> dict[str, object]: def refresh_manifest_environments( - manifest: Manifest, root: Path, built_env: str | None = None + manifest: Manifest, + root: Path, + built_env: str | None = None, + packed_images: dict[str, dict] | None = None, ) -> Manifest: """ Update manifest with current environment state. @@ -176,6 +181,13 @@ def refresh_manifest_environments( env the manifest has never seen gets the env directory's mtime as a best-effort estimate — never `now`, which would fake freshness into the `verified_at > built_at` staleness comparison. + + `packed_images` carries freshly packed image records (env name -> record + from :func:`rootstock.pack.pack_environment`); each is stamped with + ``packed_at`` here, in the same refresh that stamps ``built_at``, so a + just-installed env's image always satisfies the ``packed_at >= built_at`` + currency check. Other envs keep their existing image record — currency is + judged at read time, never by dropping history. """ from . import __version__ from .install_state import read_install_state @@ -223,6 +235,11 @@ def refresh_manifest_environments( else: built_at = built_at_estimate(env.path) + if packed_images and env_name in packed_images: + image = {**packed_images[env_name], "packed_at": now_iso()} + else: + image = env.record.image if env.record else None + manifest.environments[env_name] = EnvironmentInfo( built_at=built_at, source_hash=env.source_hash, @@ -232,6 +249,7 @@ def refresh_manifest_environments( checkpoints=checkpoints, lock_hash=env.lock_hash, clusters=env_clusters, + image=image, ) # The filesystem is the truth for what's installed: a record whose env @@ -252,6 +270,7 @@ def update_and_push_manifest( quiet: bool = False, push: bool = True, built_env: str | None = None, + packed_images: dict[str, dict] | None = None, ) -> bool: """ Update manifest with current state and optionally push to backend. @@ -267,6 +286,8 @@ def update_and_push_manifest( push: Whether to push to backend (default True) built_env: Env name that was (re)built by the calling command, if any; its built_at is stamped to now + packed_images: Freshly packed image records to stamp into the + refreshed manifest (see refresh_manifest_environments) Returns: True if push succeeded or was skipped (no API key), False on error @@ -302,7 +323,9 @@ def update_and_push_manifest( manifest = create_manifest(root, clusters, config) # Refresh environment info from current state - manifest = refresh_manifest_environments(manifest, root, built_env=built_env) + manifest = refresh_manifest_environments( + manifest, root, built_env=built_env, packed_images=packed_images + ) # Save locally save_manifest(manifest, root) @@ -581,6 +604,7 @@ def install_environment( upgrade: bool = False, verbose: bool = False, push: bool = True, + pack: bool = True, progress: Progress | None = None, ) -> InstallResult: """ @@ -706,6 +730,7 @@ def install_environment( verbose=verbose, push=push, upgrade=upgrade, + pack=pack, progress=progress, ) finally: @@ -726,6 +751,7 @@ def _build_and_swap( verbose: bool, push: bool, upgrade: bool, + pack: bool, progress: Progress | None, ) -> bool: """Build the venv into build_dir, then atomically swap it into env_target. @@ -932,15 +958,106 @@ def _build_and_swap( _say(progress, "7. Pre-compiling bytecode...") _precompile_environment(env_python, env_target) + # Pack the single-image archive for node-local staging (#180) inside the + # same install transaction that produced the env, so image and env can + # only ever describe the same build. Best-effort: a failed pack degrades + # spawns to the prewarm path, never fails the install. + packed_images = None + if pack: + _say(progress, "8. Packing staging image...") + record = pack_environment_best_effort( + root, env_name, progress=None if progress is None else (lambda m: _say(progress, m)) + ) + if record is not None: + packed_images = {env_name: record} + _say(progress, f"\nBuilt environment: {env_target}") # Update manifest. built_env stamps this env's built_at to now — the one # moment the true build time is known. - update_and_push_manifest(root, quiet=progress is None, push=push, built_env=env_name) + update_and_push_manifest( + root, + quiet=progress is None, + push=push, + built_env=env_name, + packed_images=packed_images, + ) return locked +def pack_environments( + root: Path, + env_names: list[str] | None = None, + *, + push: bool = True, + progress: Progress | None = None, +) -> dict[str, dict]: + """Pack staging images for built envs and record them in the manifest. + + The backfill half of image packing (#180): ``install`` packs its own env, + this covers everything built before packing existed (or whose pack was + skipped/failed). ``env_names=None`` packs every built env whose recorded + image is missing or stale; naming envs packs exactly those, fresh or not. + + Returns the new image records by env name. Raises OperationError when a + named env isn't built, and — after recording whatever did pack in the + manifest — when any requested pack failed. The manifest update must not + be skippable by a late failure: earlier packs in the same run already + deleted the superseded archives their old records point at. + """ + root = Path(root) + built = {name for name, _ in list_built_environments(root)} + + if env_names: + missing = sorted(set(env_names) - built) + if missing: + raise OperationError( + f"not built at {root}: {', '.join(missing)}. " + f"Built envs: {', '.join(sorted(built)) or '(none)'}" + ) + targets = list(env_names) + else: + manifest = load_manifest(root) + recorded = manifest.environments if manifest else {} + targets = sorted( + name + for name in built + if name not in recorded or not _image_usable(root, recorded[name]) + ) + if not targets: + _say(progress, "All built envs already have current images.") + return {} + + packed: dict[str, dict] = {} + failures: list[str] = [] + for env_name in targets: + _say(progress, f"Packing {env_name}...") + try: + packed[env_name] = pack_environment(root, env_name, progress=progress) + except PackError as exc: + failures.append(f"{env_name}: {exc}") + + if packed: + update_and_push_manifest(root, quiet=progress is None, push=push, packed_images=packed) + if failures: + raise OperationError("some envs could not be packed:\n " + "\n ".join(failures)) + return packed + + +def _image_usable(root: Path, env: EnvironmentInfo) -> bool: + """Current by timestamp AND the archive actually on disk. The sweep's + repack filter must see through a purged images/ dir — a current-looking + record with no file behind it would otherwise report "nothing to pack" + forever while every spawn falls back to prewarm.""" + from .manifest import image_is_current + + if not image_is_current(env): + return False + path = env.image.get("path") if isinstance(env.image, dict) else None + return isinstance(path, str) and bool(path) and (Path(root) / path).is_file() + + # ----------------------------------------------------------------------------- # Checkpoint add (download + verify) # ----------------------------------------------------------------------------- diff --git a/rootstock/pack.py b/rootstock/pack.py new file mode 100644 index 0000000..6d7f24e --- /dev/null +++ b/rootstock/pack.py @@ -0,0 +1,259 @@ +"""Pack a built env into a single compressed image for node-local staging. + +Cold worker starts on network filesystems are bounded by per-file metadata +RPCs and mmap fault storms (#167). A single ``tar.zst`` of the env tree +makes the transfer one sequential stream, and extracting it to node-local +disk (see :mod:`rootstock.stage`, #180) makes every subsequent ``stat``, +import, and mmap local. The image is an accelerator artifact derived from +the env — the Lustre tree stays the source of truth, and images are +regenerable at will. + +Each image contains the env directory *and* the interpreter directory its +venv symlinks resolve through, both as root-relative paths:: + + envs//... + .python/cpython-3.11.15-linux-x86_64-gnu/... + +so extraction reproduces the layout the venv expects. Weights are NOT +packed: they change on every ``rootstock add`` while envs change only on +``install``, so bundling them would stale the image constantly. They stage +separately from the manifest's per-checkpoint ``weight_files`` records. + +Packing shells out to ``tar`` and ``zstd`` (probed up front) rather than +using Python archive modules: the same two tools are all extraction needs, +so a client can stage on any node where they exist, and zstd's multi-thread +compressor is far faster than anything in-process. +""" + +from __future__ import annotations + +import hashlib +import os +import re +import shutil +import subprocess +import sys +import time +from pathlib import Path + +from .exceptions import RootstockError + +IMAGE_FORMAT = "tar.zst" +IMAGES_DIRNAME = "images" + +# zstd level 3: within a few percent of higher levels on .so-heavy trees but +# several times faster to compress — and decompression speed (the spawn-path +# cost) is essentially level-independent. +_ZSTD_LEVEL = 3 + +# A .packing partial older than this is a crashed pack's leftover even if +# some process holds its recorded pid (pid reuse); far beyond any real pack. +_PACK_STALE_SECONDS = 6 * 3600.0 + + +def _pid_alive(pid: int) -> bool: + """Shared by the pack-side partial sweep and the staging module (which + imports it from here — stage imports pack, so the helper can't live + there without a cycle).""" + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except (PermissionError, OSError): + pass + return True + + +class PackError(RootstockError, RuntimeError): + """Packing an env image failed. Messages are user-presentable.""" + + +def pack_tools_missing() -> str | None: + """Name the missing archive tool(s), or None when packing/staging can run.""" + missing = [tool for tool in ("tar", "zstd") if shutil.which(tool) is None] + if missing: + return " and ".join(missing) + return None + + +def env_interpreter_dir(root: Path, env_name: str) -> Path: + """The ``{root}/.python/`` directory this env's venv runs on. + + Resolved through ``envs//bin/python`` — the one link that must hold + for the env to work at all. Raises PackError when it resolves outside the + root's ``.python/`` (an env built against a foreign interpreter can't be + made self-contained by this packer). + """ + python_dir = (root / ".python").resolve() + real_python = (root / "envs" / env_name / "bin" / "python").resolve() + try: + relative = real_python.relative_to(python_dir) + except ValueError: + raise PackError( + f"env '{env_name}' runs on an interpreter outside {root / '.python'} " + f"({real_python}) — cannot pack a self-contained image." + ) from None + return root / ".python" / relative.parts[0] + + +def _tree_bytes(*trees: Path) -> int: + """Total apparent size of every regular file under ``trees`` (symlinks not + followed). Runs right after a build, so the stats are metadata-cache warm.""" + total = 0 + for tree in trees: + for dirpath, _dirnames, filenames in os.walk(tree): + for filename in filenames: + try: + total += os.lstat(os.path.join(dirpath, filename)).st_size + except OSError: + continue + return total + + +def pack_environment(root: Path | str, env_name: str, progress=None) -> dict: + """Pack one built env into ``{root}/images/-.tar.zst``. + + Streams ``tar | zstd`` straight into the image file while hashing, so + nothing is read back to compute the identity. The finished archive is + renamed into place atomically and superseded images of the same env are + removed. Returns the manifest ``image`` record (without ``packed_at``, + which the manifest refresh stamps alongside ``built_at`` so the + ``packed_at >= built_at`` currency check can't lose the race with its + own install). + + Raises PackError when the env is not built, tools are missing, or the + archive pipeline fails. + """ + root = Path(root) + env_dir = root / "envs" / env_name + if not (env_dir / "bin" / "python").exists(): + raise PackError(f"env '{env_name}' is not built at {env_dir} — nothing to pack.") + + missing = pack_tools_missing() + if missing: + raise PackError(f"packing needs {missing} on PATH (on clusters, try `module load zstd`).") + + interp_dir = env_interpreter_dir(root, env_name) + members = [ + str(env_dir.relative_to(root)), + str(interp_dir.relative_to(root)), + ] + uncompressed = _tree_bytes(env_dir, interp_dir) + + images_dir = root / IMAGES_DIRNAME + images_dir.mkdir(parents=True, exist_ok=True) + partial = images_dir / f".{env_name}.packing.{os.getpid()}" + + if progress is not None: + progress(f" Packing {env_name}: {', '.join(members)} ({uncompressed / 1e9:.1f} GB)") + + digest = hashlib.sha256() + compressed = 0 + tar_proc = zstd_proc = None + try: + with open(partial, "wb") as out: + tar_proc = subprocess.Popen( + ["tar", "-cf", "-", "-C", str(root), *members], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + zstd_proc = subprocess.Popen( + ["zstd", f"-{_ZSTD_LEVEL}", "-T0", "-q", "-c"], + stdin=tar_proc.stdout, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + assert tar_proc.stdout is not None and zstd_proc.stdout is not None + tar_proc.stdout.close() # so tar sees EPIPE if zstd dies + while True: + chunk = zstd_proc.stdout.read(1 << 20) + if not chunk: + break + digest.update(chunk) + compressed += len(chunk) + out.write(chunk) + + zstd_err = zstd_proc.communicate()[1] + tar_err = tar_proc.communicate()[1] + if tar_proc.returncode != 0: + raise PackError(f"tar failed packing '{env_name}': {tar_err.decode(errors='replace')}") + if zstd_proc.returncode != 0: + raise PackError( + f"zstd failed packing '{env_name}': {zstd_err.decode(errors='replace')}" + ) + + sha256 = digest.hexdigest() + image_name = f"{env_name}-{sha256[:12]}.{IMAGE_FORMAT}" + try: + partial.rename(images_dir / image_name) + except OSError as exc: + # Most plausibly a concurrent pack of the same env swept our + # partial; surface it as the domain error so batch callers + # (pack_environments) report it instead of crashing. + raise PackError( + f"could not move the finished image for '{env_name}' into " + f"place (concurrent pack?): {exc}" + ) from exc + except Exception: + for proc in (tar_proc, zstd_proc): + if proc is not None and proc.poll() is None: + proc.kill() + proc.wait() + partial.unlink(missing_ok=True) + raise + + # Superseded images of this env are dead weight on the shared + # filesystem; the new archive is the only one the manifest will point + # at. Exact-match on the <12-hex-sha> suffix — a bare `{env_name}-*` + # glob would also swallow a dash-extended sibling env's images + # (packing 'ani' must never delete 'ani-tuned-.tar.zst'). + stale_image = re.compile(rf"^{re.escape(env_name)}-[0-9a-f]{{12}}\.{re.escape(IMAGE_FORMAT)}$") + for stale in images_dir.iterdir(): + if stale.name != image_name and stale_image.match(stale.name): + stale.unlink(missing_ok=True) + # Our own partial was renamed into place above; remaining .packing.* + # entries are crashed packs' leftovers — unless their recorded pid is + # still alive and the file is fresh (a concurrent pack of this env, + # e.g. install's auto-pack racing a batch `rootstock pack`). + for stale in images_dir.glob(f".{env_name}.packing.*"): + try: + pid = int(stale.name.rsplit(".", 1)[-1]) + except ValueError: + pid = None + try: + age = time.time() - stale.stat().st_mtime + except OSError: + continue # gone already + if pid is not None and _pid_alive(pid) and age < _PACK_STALE_SECONDS: + continue + stale.unlink(missing_ok=True) + + if progress is not None: + progress( + f" Packed {image_name}: {compressed / 1e9:.1f} GB (from {uncompressed / 1e9:.1f} GB)" + ) + + return { + "path": f"{IMAGES_DIRNAME}/{image_name}", + "sha256": sha256, + "format": IMAGE_FORMAT, + "compressed_bytes": compressed, + "uncompressed_bytes": uncompressed, + } + + +def pack_environment_best_effort(root: Path, env_name: str, progress=None) -> dict | None: + """Pack, degrading to a warning: the image is an accelerator, so a failed + pack (no zstd on PATH, say) must never fail the install that triggered it. + The manifest's currency check then reports no usable image and spawns fall + back to the prewarm path.""" + try: + return pack_environment(root, env_name, progress=progress) + except Exception as exc: # noqa: BLE001 - packing is strictly optional + print( + f"Warning: could not pack an image for '{env_name}' " + f"({exc}); worker spawns will use the prewarm path instead. " + f"Retry later with `rootstock pack {env_name}`.", + file=sys.stderr, + ) + return None diff --git a/rootstock/prewarm.py b/rootstock/prewarm.py index caf8698..e801185 100644 --- a/rootstock/prewarm.py +++ b/rootstock/prewarm.py @@ -145,7 +145,7 @@ def prewarm_files(paths, max_workers: int | None = None) -> tuple[int, int]: return _read_sized(_stat_files(paths), max_workers) -def _fmt_bytes(n: int) -> str: +def _fmt_bytes(n: float) -> str: return f"{n / 1e9:.1f} GB" if n >= 1e9 else f"{n / 1e6:.0f} MB" @@ -194,7 +194,7 @@ def _cgroup_memory_limit( return min(limits) if limits else None -def prewarm_from_spec(spec: dict, log=None) -> None: +def prewarm_from_spec(spec: dict, log=None, label: str = "[Worker]") -> None: """Warm the page cache for a worker spawn spec; never raises. The one-line summary goes to ``log`` (default stderr, so it lands in @@ -203,7 +203,8 @@ def prewarm_from_spec(spec: dict, log=None) -> None: replaces an hours-long stall with a visible, bounded read. The weights portion is reported separately, tagged with the tier that resolved it (#178) — field data for retiring the heuristic once manifest records - are universal. + are universal. ``label`` prefixes each line: the default is the worker + wrapper's; the ``rootstock stage`` prologue fallback passes its own. """ if os.environ.get("ROOTSTOCK_NO_PREWARM"): return @@ -227,7 +228,7 @@ def prewarm_from_spec(spec: dict, log=None) -> None: limit = _cgroup_memory_limit() if limit is not None and expected > limit: print( - f"[Worker] Warning: expected cold working set " + f"{label} Warning: expected cold working set " f"{_fmt_bytes(expected)} exceeds this job's memory limit " f"{_fmt_bytes(limit)}; warmed pages will be evicted before " f"the worker reads them — request more memory for the job", @@ -239,7 +240,7 @@ def prewarm_from_spec(spec: dict, log=None) -> None: elapsed = time.monotonic() - began summary = ( - f"[Worker] Prewarmed page cache: {n_files} files, " + f"{label} Prewarmed page cache: {n_files} files, " f"{_fmt_bytes(n_bytes)} in {elapsed:.1f}s" ) tier = spec.get("prewarm_weights_tier") @@ -257,6 +258,6 @@ def prewarm_from_spec(spec: dict, log=None) -> None: print(summary, file=log, flush=True) except Exception as exc: # noqa: BLE001 - never take the worker down try: - print(f"[Worker] Prewarm skipped: {type(exc).__name__}: {exc}", file=log, flush=True) + print(f"{label} Prewarm skipped: {type(exc).__name__}: {exc}", file=log, flush=True) except Exception: pass diff --git a/rootstock/spawn.py b/rootstock/spawn.py index 144b2ea..a47f64f 100644 --- a/rootstock/spawn.py +++ b/rootstock/spawn.py @@ -152,6 +152,13 @@ def spawn_in_env( Stage ``wrapper_source`` + a JSON sidecar for ``payload`` and yield the command that runs them with the env's Python. + When node-local staging is configured and the env has a current packed + image (#180), worker spawns run from a node-local extraction instead of + the shared tree — ``env_dir``, the interpreter, and (when the + checkpoint's weights could be overlaid) the cache env vars all point at + local disk, and the page-cache prewarm is skipped as redundant. Any + missing piece degrades back to the shared tree + prewarm. + Args: root: Rootstock install root. env_name: Name of the pre-built environment. @@ -180,15 +187,37 @@ def spawn_in_env( env_python = get_env_python(root, env_name) env_dir = root / "envs" / env_name + # Node-local staging (#180): extract the env's packed image (and overlay + # the checkpoint's recorded weights) to local disk, and run the worker + # from the copy. Best-effort — None means every piece below behaves + # exactly as before staging existed. Download spawns never stage: their + # job is to *write* the shared cache. + staged = None + if wrapper_source != DOWNLOAD_WRAPPER: + from .stage import stage_for_spawn + + staged = stage_for_spawn(root, env_name, payload, cache_root) + if staged is not None: + env_dir = staged.env_dir + env_python = env_dir / "bin" / "python" + + # Everything is node-local and the worker's prewarm will be disabled — + # except when a user-supplied (:custom) weights file still lives on the + # shared filesystem and needs the warm. + fully_local = ( + staged is not None and staged.cache_base is not None and not payload.get("checkpoint_path") + ) + # Fill the weight-prewarm hint here rather than in each caller: every # spawn (calculator, verify, serve, add) benefits, and this is the same # choke point that already owns env_dir and the cache env vars. # Best-effort by contract — a failed lookup must never fail the spawn. - # Download spawns get the record tier only: their weight record is - # written after the download, so the heuristic would cold-read whole - # family cache dirs (typically on a login node) for weights the - # download may be about to (re)write. - if payload.get("checkpoint") and "prewarm_paths" not in payload: + # Skipped when fully staged (the worker won't prewarm at all). Download + # spawns get the record tier only: their weight record is written after + # the download, so the heuristic would cold-read whole family cache dirs + # (typically on a login node) for weights the download may be about to + # (re)write. + if payload.get("checkpoint") and "prewarm_paths" not in payload and not fully_local: try: paths, tier = get_checkpoint_prewarm_paths( root, @@ -204,7 +233,17 @@ def spawn_in_env( payload = {**payload, "prewarm_paths": paths, "prewarm_weights_tier": tier} env = os.environ.copy() - env.update(get_model_cache_env(root, cache_root)) + if staged is not None and staged.cache_base is not None: + env.update(get_model_cache_env(root, staged.cache_base)) + # With env and weights both node-local, the page-cache prewarm is + # pure overhead (see fully_local above for the :custom exception). + if fully_local: + env["ROOTSTOCK_NO_PREWARM"] = "1" + else: + # Staged env without staged weights keeps the prewarm: the env-tree + # portion re-reads from local disk at memory-ish speed, and the + # weight portion still streams the shared cache ahead of the mmaps. + env.update(get_model_cache_env(root, cache_root)) if offline: env["HF_HUB_OFFLINE"] = "1" diff --git a/rootstock/stage.py b/rootstock/stage.py new file mode 100644 index 0000000..c9a5cbf --- /dev/null +++ b/rootstock/stage.py @@ -0,0 +1,856 @@ +"""Stage packed env images (and checkpoint weights) to node-local disk. + +The structural fix for cold starts on network filesystems (#180): the spawn +path downloads each env as **one sequential read of a compressed image** +(see :mod:`rootstock.pack`) and extracts it to node-local disk, after which +imports, ``stat`` calls, and weight mmaps never touch the network again — +the two costs page-cache prewarming can't remove (metadata RPCs, and warmth +eviction under memory pressure) simply disappear. + +Graceful degradation is load-bearing: any missing piece — no staging dir +configured, no packed image, image stale relative to the env build, missing +tools, insufficient disk, a lost lock race that times out, any unexpected +error — falls back to the existing prewarm path. A cluster where nobody +configures anything behaves exactly as before this module existed. + +Where the staging base comes from (first declaration wins, then must +validate — exist, be writable, and live on a *different* filesystem than +the install root — or staging is disabled): + +1. ``ROOTSTOCK_STAGE_DIR`` (user override / experiments / A-B testing), +2. ``stage_dir`` in ``{root}/layout.json`` (the install's own declaration; + may contain env vars like ``$SLURM_TMPDIR``, expanded here on the node), +3. ``Cluster.stage_dir`` in the registry (legacy fallback). + +``ROOTSTOCK_NO_STAGE=1`` force-disables, mirroring ``ROOTSTOCK_NO_PREWARM``. + +Layout under the resolved base (per-user — staged trees are exec'd, so they +are never shared between users):: + + {base}/rootstock/{user}/ + ├── envs-by-hash/ + │ ├── {archive sha256}/ # one extracted image (content-addressed: + │ │ ├── envs// # reuse across spawns/jobs is free, and a + │ │ └── .python// # rebuilt env lands in a new dir) + │ ├── {sha256}.lock # one extractor per node; others wait + │ └── {sha256}.partial. # in-flight extraction (swept when dead) + └── cache-mirror/ # weight overlay, shared across checkpoints + +The weight mirror is an **overlay, not a copy**: each checkpoint's +manifest-recorded ``weight_files`` (#177) are materialized locally, and +everything else in the shared cache appears as a symlink back to it. The +fallthrough matters — the records only cover mmap-visible files, so small +side files (HF ``config.json`` and friends) must still resolve, just at +ordinary read speed. Worker caches (``HOME``/``XDG_CACHE_HOME``/``HF_HOME``) +are repointed at the mirror only when the overlay succeeds. +""" + +from __future__ import annotations + +import getpass +import hashlib +import json +import os +import shutil +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + +from . import __version__ +from .pack import IMAGE_FORMAT, _pid_alive, pack_tools_missing +from .prewarm import _fmt_bytes + +STAGE_DIR_ENV = "ROOTSTOCK_STAGE_DIR" +NO_STAGE_ENV = "ROOTSTOCK_NO_STAGE" + +# A lockfile older than this is a crashed extractor's leftover (holds are +# minutes even on a congested night); a waiter that hasn't seen the winner +# finish by the same deadline gives up and falls back to prewarm. +_LOCK_STALE_SECONDS = 900.0 +_WAIT_POLL_SECONDS = 1.0 + +# Extraction preflight headroom over the recorded uncompressed size (tar +# rounding, filesystem overhead, the in-flight compressed stream). +_SPACE_HEADROOM = 1.2 + +# Never evict a staged env younger than this: its job is plausibly still +# running from it. Older entries are fair game oldest-first. +_EVICT_MIN_AGE_SECONDS = 6 * 3600.0 + + +@dataclass +class StagedSpawn: + """What a successful staging pass hands back to ``spawn_in_env``.""" + + env_dir: Path # staged replacement for {root}/envs/ + cache_base: Path | None # weight-mirror base to point caches at, if staged + + +def _log(message: str) -> None: + """Staging progress goes to the *client's* stderr (unlike the prewarm, + which runs inside the worker): it is visible in job logs while the worker + is still being spawned, and bracket lines before/after each phase keep a + stall distinguishable from a hang.""" + print(f"[Rootstock] {message}", file=sys.stderr, flush=True) + + +def _same_filesystem(a: Path, b: Path) -> bool: + return os.stat(a).st_dev == os.stat(b).st_dev + + +def resolve_stage_base(root: Path) -> Path | None: + """Resolve and validate the node-local staging base for an install. + + Returns None — staging disabled — when nothing is declared or the first + declaration found fails validation. Deliberately does not fall through + to later declarations: a declared-but-broken dir is a configuration + problem to surface (via debug-level fallback behavior), not to paper + over with a stale registry entry. + """ + if os.environ.get(NO_STAGE_ENV): + return None + + raw = os.environ.get(STAGE_DIR_ENV) + if not raw: + from .layout import read_declared_stage_dir + + raw = read_declared_stage_dir(root) + if not raw: + from .clusters import get_cluster, get_cluster_for_root + + cluster_name = get_cluster_for_root(root) + if cluster_name is not None: + raw = get_cluster(cluster_name).stage_dir + if not raw: + return None + + expanded = os.path.expandvars(os.path.expanduser(raw)) + if "$" in expanded: + # An env var in the declaration didn't expand here (e.g. + # $SLURM_TMPDIR outside a job) — that node has no staging dir. + return None + base = Path(expanded) + try: + if not base.is_dir() or not os.access(base, os.W_OK | os.X_OK): + return None + # A "node-local" path that is actually the same filesystem as the + # install would double every cost this module exists to remove. + if _same_filesystem(base, Path(root)): + return None + except OSError: + return None + return base + + +def _read_manifest_env(root: Path, env_name: str) -> dict | None: + """Raw ``manifest.json`` read of one env record — same contract as the + prewarm-path lookup (#178): this runs on the spawn path as an + optimization hint, so it must never print migration notes, take locks, + or refuse a newer schema. Anything unreadable means "no record".""" + try: + with open(Path(root) / "manifest.json") as f: + data = json.load(f) + env = data.get("environments", {}).get(env_name) + except (OSError, ValueError, AttributeError): + return None + return env if isinstance(env, dict) else None + + +def _current_image_record(root: Path, env_record: dict | None) -> dict | None: + """The env record's image entry, iff it describes the current build and + the archive file is present. Mirrors + :func:`rootstock.manifest.image_is_current` on the raw dict.""" + if env_record is None: + return None + image = env_record.get("image") + built_at = env_record.get("built_at") + if not (isinstance(image, dict) and isinstance(built_at, str)): + return None + packed_at = image.get("packed_at") + if not (isinstance(packed_at, str) and packed_at >= built_at): + return None + if not ( + isinstance(image.get("path"), str) + and isinstance(image.get("sha256"), str) + and image.get("format") == IMAGE_FORMAT + and isinstance(image.get("uncompressed_bytes"), int) + ): + return None + if not (Path(root) / image["path"]).is_file(): + return None + return image + + +def _recorded_weight_entries(env_record: dict | None, checkpoint: str) -> list | None: + """The checkpoint's ``weight_files`` list from a pre-read raw env record + (same semantics as ``environment._recorded_weight_files``, which reads + the manifest itself for the prewarm path).""" + if not isinstance(env_record, dict): + return None + record = (env_record.get("checkpoints") or {}).get(checkpoint) or {} + weight_files = record.get("weight_files") + return weight_files if isinstance(weight_files, list) else None + + +def _user_stage_root(base: Path) -> Path: + """``{base}/rootstock/{user}``, created 0700: staged trees are exec'd, so + they must not be writable — or trustingly reusable — across users. + + The shared ``{base}/rootstock`` intermediate is made sticky-1777 (the + /tmp recipe, same as the usage spool): whoever stages first must not + lock everyone else's leaf-mkdir out via their umask. chmod is + best-effort — on an already-existing dir owned by someone else it + fails, and the mkdir below then either works or disables staging for + this user with a visible log line (via the callers' guards). + """ + shared = base / "rootstock" + shared.mkdir(exist_ok=True) + try: + os.chmod(shared, 0o1777) + except OSError: + pass + user_root = shared / getpass.getuser() + user_root.mkdir(exist_ok=True) + os.chmod(user_root, 0o700) + return user_root + + +def _sweep_partials(envs_root: Path) -> None: + """Remove extraction leftovers whose owning process is gone (SIGKILL + hygiene — same "leak into node-local temp, auto-cleaned" posture as the + spawn tmp dir).""" + for partial in envs_root.glob("*.partial.*"): + try: + pid = int(partial.name.rsplit(".", 1)[-1]) + except ValueError: + continue + try: + stale_by_age = time.time() - partial.stat().st_mtime > _LOCK_STALE_SECONDS + except OSError: + continue # renamed/removed mid-scan by its extractor + if not _pid_alive(pid) or stale_by_age: + shutil.rmtree(partial, ignore_errors=True) + + +def _mark_in_use(staged_root: Path) -> None: + """Record this process as a user of a staged env, in + ``{staged}/.users/``. The recording process is the *client* (the + calculator / server process), which lives as long as any worker it + spawns — so pid-aliveness of these files is the eviction shield for + long MD runs. Dead pidfiles are cleaned during eviction scans; nothing + removes them on exit, and nothing needs to.""" + users = staged_root / ".users" + try: + users.mkdir(exist_ok=True) + (users / str(os.getpid())).touch() + except OSError: + pass # marking is best-effort; min-age still shields young dirs + + +def _in_use(staged_root: Path) -> bool: + """Whether any recorded user of a staged env is still alive (pruning + dead pidfiles as a side effect).""" + live = False + try: + pidfiles = list((staged_root / ".users").iterdir()) + except OSError: + return False + for pidfile in pidfiles: + if pidfile.name.isdigit() and _pid_alive(int(pidfile.name)): + live = True + else: + pidfile.unlink(missing_ok=True) + return live + + +def _evict_lru(envs_root: Path, keep: Path, bytes_needed: int) -> None: + """Free space by removing the oldest staged envs — never ``keep``, + never anything younger than the min age, and never a dir some live + client process is registered against (mtime alone can't shield a + multi-day MD run on a persistent /tmp). Best-effort: rechecks free + space after each removal.""" + try: + entries = sorted( + (d for d in envs_root.iterdir() if d.is_dir() and d != keep), + key=lambda d: d.stat().st_mtime, + ) + except OSError: + return + now = time.time() + for entry in entries: + if shutil.disk_usage(envs_root).free >= bytes_needed: + return + try: + if now - entry.stat().st_mtime < _EVICT_MIN_AGE_SECONDS: + break # sorted oldest-first: everything after is younger + except OSError: + continue + if _in_use(entry): + continue # a live client is running workers out of it + _log(f"Stage evicting {entry.name} (LRU) to make room") + shutil.rmtree(entry, ignore_errors=True) + + +class _StageLock: + """O_EXCL lockfile granting one extractor per archive per node. Losing + the race is not an error — the loser waits for the winner's atomic + rename instead of duplicating a multi-GB extraction (four contending + prewarm streams are exactly the pathology staging replaces).""" + + def __init__(self, path: Path): + self.path = path + self.acquired = False + + def try_acquire(self) -> bool: + for _ in range(2): # second try after clearing a stale lock + try: + fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + if self._is_stale(): + self.path.unlink(missing_ok=True) + continue + return False + with os.fdopen(fd, "w") as f: + f.write(str(os.getpid())) + self.acquired = True + return True + return False + + def _is_stale(self) -> bool: + try: + content = self.path.read_text().strip() + age = time.time() - self.path.stat().st_mtime + except OSError: + return False # vanished — holder just released it + if age > _LOCK_STALE_SECONDS: + return True + return content.isdigit() and not _pid_alive(int(content)) + + def release(self) -> None: + if self.acquired: + self.path.unlink(missing_ok=True) + self.acquired = False + + +class _FixupError(RuntimeError): + """A staged tree could not be pointed at its local interpreter. + + Deterministic for a given (archive, client version) — unlike transient + extraction failures — so stage_env caches it per sha and stops paying + the multi-GB extract-and-discard on every subsequent spawn. + """ + + +def _remap_into_stage(value: str, root: Path, staged_root: Path) -> str | None: + """Rewrite an absolute path under the shared ``root`` to its staged + equivalent; None when it doesn't point under the root. + + Tries the root as spelled and as resolved, then — because uv bakes the + *install-time* spelling into symlink targets and pyvenv.cfg, which on + multi-alias mounts (/eagle vs /lus/eagle at ALCF) matches neither — + resolves the value itself and retries against the resolved root. + """ + for prefix in {str(root), str(Path(root).resolve())}: + if value == prefix or value.startswith(prefix.rstrip("/") + "/"): + return str(staged_root) + value[len(prefix.rstrip("/")) :] + if os.path.isabs(value): + try: + real = str(Path(value).resolve()) + except OSError: + return None + root_real = str(Path(root).resolve()).rstrip("/") + if real == root_real or real.startswith(root_real + "/"): + return str(staged_root) + real[len(root_real) :] + return None + + +def _fixup_staged_env(tree: Path, final_root: Path, root: Path, env_name: str) -> None: + """Point the extracted venv (still in the ``tree`` partial dir) at the + interpreter it will have once the tree is renamed to ``final_root``. + + The archive reproduces root-relative layout, but two things inside a + venv name the shared root absolutely and would quietly put the + interpreter and stdlib back on the network filesystem: + + - the ``bin/python*`` symlinks (uv writes absolute targets), and + - ``pyvenv.cfg``'s ``home =`` line, which is where the stdlib is + resolved from at startup. + + Targets are written in ``final_root`` terms — dangling until the atomic + rename, which is exactly the publication point waiters key off — and + verified through ``tree``. Raises RuntimeError when ``bin/python`` can't + be made local: a staged env that still executes from Lustre is worse + than falling back. + """ + env_dir = tree / "envs" / env_name + for entry in (env_dir / "bin").iterdir(): + if not entry.is_symlink(): + continue + target = os.readlink(entry) + if not os.path.isabs(target): + continue + remapped = _remap_into_stage(target, root, final_root) + if remapped is None: + if entry.name.startswith("python"): + raise _FixupError(f"staged {entry.name} links outside the install root ({target})") + continue + entry.unlink() + os.symlink(remapped, entry) + + pyvenv = env_dir / "pyvenv.cfg" + if pyvenv.is_file(): + lines = [] + for line in pyvenv.read_text().splitlines(): + key, sep, value = line.partition("=") + remapped = _remap_into_stage(value.strip(), root, final_root) if sep else None + lines.append(f"{key.rstrip()} {sep} {remapped}" if remapped else line) + pyvenv.write_text("\n".join(lines) + "\n") + + python = env_dir / "bin" / "python" + target = os.readlink(python) if python.is_symlink() else None + if target is None or not (tree / Path(target).relative_to(final_root)).exists(): + raise _FixupError("staged bin/python does not resolve after fixup") + + +def stage_env(root: Path, env_name: str, base: Path, env_record: dict | None = None) -> Path | None: + """Materialize the env's packed image under ``base``; return the staged + root (containing ``envs/`` and ``.python/``) or None to fall back. + + Content-addressed by archive sha256: a warm dir is reused with zero + reads, and a rebuilt env (new archive, new sha) lands beside the old one, + which ages out via LRU eviction. ``env_record`` is the pre-read raw + manifest env record, for callers that already have it. Never raises — + the CLI calls this bare, and even the preamble (mkdir on a shared /tmp, + disk_usage, lock files) can fail on a hostile node. + """ + try: + return _stage_env(root, env_name, base, env_record) + except Exception as exc: # noqa: BLE001 - staging must never fail the caller + try: + _log( + f"Stage skipped ({env_name}): {type(exc).__name__}: {exc}; falling back to prewarm" + ) + except Exception: + pass + return None + + +def _stage_env( + root: Path, env_name: str, base: Path, env_record: dict | None = None +) -> Path | None: + if env_record is None: + env_record = _read_manifest_env(root, env_name) + image = _current_image_record(root, env_record) + if image is None: + return None + missing_tools = pack_tools_missing() + if missing_tools: + _log(f"Stage skipped ({env_name}): {missing_tools} not on PATH; falling back to prewarm") + return None + + envs_root = _user_stage_root(base) / "envs-by-hash" + envs_root.mkdir(parents=True, exist_ok=True) + final = envs_root / image["sha256"] + marker = final / "envs" / env_name / "bin" / "python" + + if marker.exists(): + final.touch() # LRU freshness + _mark_in_use(final) + _log(f"Stage reused (warm): {env_name} at {final}") + return final + + # A recorded deterministic fixup failure for this archive + client + # version: don't repeat a multi-GB extract-and-discard on every spawn. + failed_note = envs_root / f"{image['sha256']}.failed" + try: + failed_version, failed_reason = failed_note.read_text().splitlines()[:2] + except (OSError, ValueError): + failed_version = None + failed_reason = "" + if failed_version == __version__: + _log( + f"Stage skipped ({env_name}): previously failed on this node " + f"({failed_reason}); falling back to prewarm" + ) + return None + failed_note.unlink(missing_ok=True) # other-version note: retry below + + _sweep_partials(envs_root) + + needed = int(image["uncompressed_bytes"] * _SPACE_HEADROOM) + if shutil.disk_usage(envs_root).free < needed: + _evict_lru(envs_root, keep=final, bytes_needed=needed) + if shutil.disk_usage(envs_root).free < needed: + _log( + f"Stage skipped ({env_name}): needs {_fmt_bytes(needed)} free " + f"at {envs_root}; falling back to prewarm" + ) + return None + + lock = _StageLock(envs_root / f"{image['sha256']}.lock") + if not lock.try_acquire(): + # Another spawn on this node is extracting the same archive; one + # image stream then N local reuses beats N contending streams. + _log(f"Staging of {env_name} in progress by another process; waiting") + deadline = time.monotonic() + _LOCK_STALE_SECONDS + while time.monotonic() < deadline: + if marker.exists(): + _mark_in_use(final) + _log(f"Stage reused (warm): {env_name} at {final}") + return final + if lock.try_acquire(): + break # winner died; take over below + time.sleep(_WAIT_POLL_SECONDS) + if not lock.acquired: + _log(f"Stage skipped ({env_name}): timed out waiting; falling back to prewarm") + return None + + partial = envs_root / f"{image['sha256']}.partial.{os.getpid()}" + image_path = Path(root) / image["path"] + try: + if marker.exists(): # completed while we raced for the lock + _mark_in_use(final) + return final + _log( + f"Staging {env_name} ({_fmt_bytes(image.get('compressed_bytes', 0))} " + f"compressed) to {final}" + ) + began = time.monotonic() + partial.mkdir() + zstd = subprocess.Popen( + ["zstd", "-dc", str(image_path)], stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + tar = subprocess.Popen( + ["tar", "-xf", "-", "-C", str(partial)], + stdin=zstd.stdout, + stderr=subprocess.PIPE, + ) + assert zstd.stdout is not None + zstd.stdout.close() + tar_err = tar.communicate()[1] + zstd_err = zstd.communicate()[1] + if zstd.returncode != 0 or tar.returncode != 0: + err = (zstd_err or tar_err).decode(errors="replace").strip() + raise RuntimeError(f"extraction failed: {err}") + + _fixup_staged_env(partial, final, Path(root), env_name) + + try: + partial.rename(final) + except OSError: + if not marker.exists(): # a real failure, not a lost race + raise + _mark_in_use(final) + _log(f"Staged {env_name} in {time.monotonic() - began:.1f}s") + return final + except Exception as exc: # noqa: BLE001 - staging must never fail the spawn + if isinstance(exc, _FixupError): + # Deterministic for this archive + client version; note it so + # later spawns skip straight to prewarm instead of re-paying + # the extraction. A client upgrade invalidates the note. + try: + failed_note.write_text(f"{__version__}\n{exc}\n") + except OSError: + pass + _log(f"Stage skipped ({env_name}): {exc}; falling back to prewarm") + shutil.rmtree(partial, ignore_errors=True) + return None + finally: + lock.release() + + +# ----------------------------------------------------------------------------- +# Weight overlay +# ----------------------------------------------------------------------------- + + +def _hub_sibling_dirs(needed: list[PurePosixPath]) -> set[PurePosixPath]: + """HuggingFace-hub special case: recorded weights are blob files, but the + worker opens them through ``snapshots//`` relative symlinks + (``../../blobs/``). If the repo's ``snapshots``/``refs`` dirs were + left as whole-directory symlinks into the shared tree, those relative + links would resolve to the *shared* blobs and the local copies would + never be read — so the sibling dirs must be recreated in the mirror.""" + extras: set[PurePosixPath] = set() + for rel in needed: + parts = rel.parts + if "blobs" in parts[:-1]: + repo = PurePosixPath(*parts[: parts.index("blobs")]) + extras.update({repo / "snapshots", repo / "refs"}) + return extras + + +def _copy_file_atomic(src: Path, dest: Path, src_stat: os.stat_result | None = None) -> None: + """Copy via tmp + rename. When ``src_stat`` is given, the source's mtime + is preserved on the copy — that is the mirror's staleness signal: weight + captures record only {path, size}, so a same-size in-place overwrite on + the shared cache (a retrained ``:custom``-adjacent file, a torch-hub + refresh) would otherwise serve stale bytes from a warm mirror forever.""" + tmp = dest.with_name(dest.name + f".copying.{os.getpid()}") + try: + shutil.copyfile(src, tmp) + if src_stat is not None: + os.utime(tmp, ns=(src_stat.st_atime_ns, src_stat.st_mtime_ns)) + tmp.rename(dest) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def _mirror_current(mirror: Path, rel: str, src_stat: os.stat_result) -> bool: + """Whether the mirror's copy of ``rel`` matches the shared source (size + and preserved mtime). A fallthrough *symlink* is never current — its + stat would trivially match the very shared file it points at, and + treating it as a copy would leave a later-recorded weight file on the + shared filesystem with the worker's prewarm switched off (worse than no + staging). ``_copy_file_atomic``'s rename-over replaces the link.""" + dest = mirror / rel + if os.path.islink(dest): + return False + try: + st = os.stat(dest) + except OSError: + return False + return st.st_size == src_stat.st_size and st.st_mtime_ns == src_stat.st_mtime_ns + + +def _weights_digest(sources: list[tuple[str, os.stat_result]]) -> str: + """Identity of one checkpoint's overlay: the recorded paths plus each + shared source's (size, mtime). Any source change — or a different + record — produces a new digest and forces a re-overlay.""" + entries = sorted((rel, st.st_size, st.st_mtime_ns) for rel, st in sources) + return hashlib.sha256(json.dumps(entries).encode()).hexdigest() + + +def _overlay_tree(shared: Path, mirror: Path, rel_dir: PurePosixPath) -> None: + """Deep-copy one small subtree's *structure*: dirs recreated, symlink + entries copied verbatim (their relative targets then resolve inside the + mirror), small regular files copied. Used only for the hub sibling dirs, + which hold pointers, not weights.""" + src_dir = shared / rel_dir + dest_dir = mirror / rel_dir + if dest_dir.is_symlink(): + dest_dir.unlink() + dest_dir.mkdir(parents=True, exist_ok=True) + try: + entries = list(src_dir.iterdir()) + except OSError: + return + for entry in entries: + dest = dest_dir / entry.name + if entry.is_symlink(): + if dest.is_symlink() or dest.exists(): + continue + os.symlink(os.readlink(entry), dest) + elif entry.is_dir(): + _overlay_tree(shared, mirror, rel_dir / entry.name) + else: + if not dest.exists(): + _copy_file_atomic(entry, dest) + + +def _overlay_recorded( + shared_base: Path, mirror: Path, sources: list[tuple[str, os.stat_result]] +) -> int: + """Build/refresh the weight overlay; returns bytes copied this pass. + + Ancestor directories of recorded files are materialized as real dirs; + at each level, siblings not on any recorded path become symlinks back + into the shared cache (the fallthrough for unrecorded side files). An + existing real dir/file in the mirror is never downgraded to a symlink — + it may be another checkpoint's materialized copy. + """ + needed = {PurePosixPath(rel): st for rel, st in sources} + + materialize: set[PurePosixPath] = {PurePosixPath("cache"), PurePosixPath("home")} + for rel in needed: + materialize.update(p for p in rel.parents if p.parts) + hub_siblings = _hub_sibling_dirs(list(needed)) + + copied = 0 + for rel_dir in sorted(materialize, key=lambda p: len(p.parts)): + src_dir = shared_base / rel_dir + dest_dir = mirror / rel_dir + if dest_dir.is_symlink(): + dest_dir.unlink() + dest_dir.mkdir(parents=True, exist_ok=True) + try: + entries = list(src_dir.iterdir()) + except OSError: + continue # shared side absent (e.g. bare home/): an empty local dir is right + for entry in entries: + rel_child = rel_dir / entry.name + dest = mirror / rel_child + if rel_child in materialize or rel_child in hub_siblings: + continue # handled by its own pass + if rel_child in needed: + src_stat = needed[rel_child] + try: + if not _mirror_current(mirror, str(rel_child), src_stat): + _copy_file_atomic(entry, dest, src_stat) + copied += src_stat.st_size + except OSError: + raise RuntimeError(f"could not copy recorded weight file {rel_child}") + elif not (dest.is_symlink() or dest.exists()): + os.symlink(entry, dest) + + for rel_dir in sorted(hub_siblings, key=lambda p: len(p.parts)): + _overlay_tree(shared_base, mirror, rel_dir) + return copied + + +def stage_weights( + root: Path, + cache_root: Path | None, + env_name: str, + checkpoint: str, + base: Path, + env_record: dict | None = None, +) -> Path | None: + """Overlay the checkpoint's recorded weight files into the node-local + cache mirror; return the mirror base or None to leave the worker's + caches on the shared filesystem. Never raises (see stage_env). + + Only the manifest record tier is trusted here (unlike the prewarm's + heuristic tier): redirecting a worker's caches at a mirror is only safe + when we know exactly which files it will mmap. Every recorded file must + exist in the shared cache — a purged file means the record is stale and + the whole overlay is skipped, self-healing on the next add/verify pass. + + A per-checkpoint completion marker records the digest of the last + finished overlay (recorded paths + shared sizes/mtimes). Matching it is + the lock-free warm path: concurrent same-checkpoint spawns on one node — + the committee-demo case — read the marker and return without ever + touching the mirror lock or re-walking the shared cache. + """ + try: + return _stage_weights(root, cache_root, env_name, checkpoint, base, env_record) + except Exception as exc: # noqa: BLE001 - staging must never fail the caller + try: + _log(f"Weights not staged ({checkpoint}): {type(exc).__name__}: {exc}") + except Exception: + pass + return None + + +def _stage_weights( + root: Path, + cache_root: Path | None, + env_name: str, + checkpoint: str, + base: Path, + env_record: dict | None = None, +) -> Path | None: + if env_record is None: + env_record = _read_manifest_env(root, env_name) + recorded = _recorded_weight_entries(env_record, checkpoint) + if not recorded: + return None + rels: list[str] = [] + for entry in recorded: + if not (isinstance(entry, dict) and isinstance(entry.get("path"), str) and entry["path"]): + return None + rels.append(entry["path"]) + + shared_base = Path(cache_root) if cache_root is not None else Path(root) + sources: list[tuple[str, os.stat_result]] = [] + for rel in rels: + try: + st = os.stat(shared_base / rel) + except OSError: + return None # purged / stale record — self-heals on the next verify + sources.append((rel, st)) + + user_root = _user_stage_root(base) + mirror = user_root / "cache-mirror" + marker = user_root / f"cache-mirror.{checkpoint}.ok" + digest = _weights_digest(sources) + + try: + if marker.read_text().strip() == digest: + return mirror # completed overlay, sources unchanged since + except OSError: + pass + + missing = sum(st.st_size for rel, st in sources if not _mirror_current(mirror, rel, st)) + if missing and shutil.disk_usage(user_root).free < missing * 1.1: + _log( + f"Weights not staged ({checkpoint}): needs {_fmt_bytes(missing * 1.1)} free at {mirror}" + ) + return None + + # One overlay mutator per node: concurrent spawns touching one family's + # dirs would race on the symlink/materialize transitions. + lock = _StageLock(user_root / "cache-mirror.lock") + deadline = time.monotonic() + _LOCK_STALE_SECONDS + while not lock.try_acquire(): + if time.monotonic() > deadline: + return None + time.sleep(_WAIT_POLL_SECONDS) + try: + mirror.mkdir(parents=True, exist_ok=True) + began = time.monotonic() + copied = _overlay_recorded(shared_base, mirror, sources) + # Marker only after the whole overlay (copies, fallthrough symlinks, + # hub siblings) finished — a crash mid-overlay leaves no marker, so + # the next spawn takes the lock and completes it. + marker_tmp = marker.with_name(marker.name + f".{os.getpid()}") + marker_tmp.write_text(digest) + marker_tmp.rename(marker) + if copied: + _log( + f"Staged weights for {checkpoint}: {_fmt_bytes(copied)} " + f"in {time.monotonic() - began:.1f}s" + ) + else: + _log(f"Weights reused (warm) for {checkpoint}") + return mirror + except Exception as exc: # noqa: BLE001 - overlay failure must not fail the spawn + _log(f"Weights not staged ({checkpoint}): {exc}") + return None + finally: + lock.release() + + +# ----------------------------------------------------------------------------- +# The spawn seam +# ----------------------------------------------------------------------------- + + +def stage_for_spawn( + root: Path, + env_name: str, + payload: dict, + cache_root: Path | None = None, +) -> StagedSpawn | None: + """Best-effort staging pass for one worker spawn; never raises. + + The env stages for every worker spawn. Weights stage — and the worker's + caches are repointed — only when the payload names a checkpoint and does + *not* request weight capture: capture runs record files relative to the + shared cache root, and add/verify passes must observe (and write) the + shared cache, not a node-local mirror. + """ + try: + base = resolve_stage_base(Path(root)) + if base is None: + return None + env_record = _read_manifest_env(root, env_name) + staged_root = stage_env(Path(root), env_name, base, env_record=env_record) + if staged_root is None: + return None + + cache_base: Path | None = None + if payload.get("checkpoint") and "weights_capture" not in payload: + cache_base = stage_weights( + root, cache_root, env_name, payload["checkpoint"], base, env_record=env_record + ) + return StagedSpawn(env_dir=staged_root / "envs" / env_name, cache_base=cache_base) + except Exception as exc: # noqa: BLE001 - staging must never fail the spawn + try: + _log(f"Stage skipped ({env_name}): {type(exc).__name__}: {exc}") + except Exception: + pass + return None diff --git a/tests/commands/test_prune_plan.py b/tests/commands/test_prune_plan.py index 349230e..1920ebb 100644 --- a/tests/commands/test_prune_plan.py +++ b/tests/commands/test_prune_plan.py @@ -455,6 +455,36 @@ def test_orphaned_interpreters_are_collected_live_ones_kept(tmp_path): assert by_name[stranded.name].reason == "stranded staging copy" +def test_orphaned_staging_images_are_collected_recorded_ones_kept(tmp_path): + source = env_source("mace-mp-0-medium") + register(tmp_path, "mace", source) + build(tmp_path, "mace", source) + images = tmp_path / "images" + images.mkdir() + recorded_image = images / "mace-abc123def456.tar.zst" + recorded_image.write_bytes(b"z") + age(recorded_image) + orphan = images / "retired-0123456789ab.tar.zst" # env pruned in a past run + orphan.write_bytes(b"z") + age(orphan) + crashed = images / ".uma.packing.4242" + crashed.write_bytes(b"z") + age(crashed) + fresh_partial = images / ".orb.packing.9999" # in-flight pack: age guard + fresh_partial.write_bytes(b"z") + + env = record(tmp_path, "mace", checkpoints={"mace-mp-0-medium": fetched()}) + env.image = {"path": "images/mace-abc123def456.tar.zst", "packed_at": NEWER} + save(tmp_path, {"mace": env}) + + plan = plan_prune(tmp_path, cache_root=tmp_path) + + by_name = {Path(i.path).name: i for i in plan.gc if i.kind == "image"} + assert set(by_name) == {orphan.name, crashed.name} + assert by_name[crashed.name].reason == "leftover from a crashed pack" + assert by_name[orphan.name].reason == "no manifest env records it" + + def test_orphaned_lockfile_is_collected(tmp_path): register(tmp_path, "mace", env_source("mace-mp-0-medium")) build(tmp_path, "mace", env_source("mace-mp-0-medium")) diff --git a/tests/manifest/test_image_record.py b/tests/manifest/test_image_record.py new file mode 100644 index 0000000..ad5d9cf --- /dev/null +++ b/tests/manifest/test_image_record.py @@ -0,0 +1,105 @@ +"""The per-env packed-image record (#180): round-trip, the packed_at >= +built_at currency check, and how the manifest refresh preserves/stamps it.""" + +from __future__ import annotations + +from pathlib import Path + +from rootstock.manifest import ( + SCHEMA_VERSION, + EnvironmentInfo, + Maintainer, + Manifest, + image_is_current, +) +from rootstock.operations import refresh_manifest_environments + +ENV_SOURCE = ( + "# /// script\n" + '# requires-python = ">=3.10"\n' + '# dependencies = ["six>=1.0"]\n' + "# ///\n" + "CHECKPOINTS = {}\n" +) + +IMAGE = { + "path": "images/demo-abc123def456.tar.zst", + "sha256": "abc123def456" + "0" * 52, + "format": "tar.zst", + "compressed_bytes": 100, + "uncompressed_bytes": 250, +} + + +def _env(built_at: str, image: dict | None) -> EnvironmentInfo: + return EnvironmentInfo( + built_at=built_at, + source_hash=None, + source="", + python_requires=">=3.10", + dependencies={}, + image=image, + ) + + +def test_image_round_trips_through_dataclasses(): + env = _env("2026-09-01T00:00:00Z", {**IMAGE, "packed_at": "2026-09-01T00:00:00Z"}) + assert EnvironmentInfo.from_dict(env.to_dict()).image == env.image + + +def test_image_currency(): + built = "2026-09-01T00:00:00Z" + assert image_is_current(_env(built, {**IMAGE, "packed_at": built})) # same refresh + assert image_is_current(_env(built, {**IMAGE, "packed_at": "2026-09-02T00:00:00Z"})) + # Rebuilt after the pack: the record survives but must not read as usable. + assert not image_is_current(_env("2026-09-03T00:00:00Z", {**IMAGE, "packed_at": built})) + assert not image_is_current(_env(built, None)) + assert not image_is_current(_env(built, IMAGE)) # no packed_at at all + + +def _make_built_env(root: Path, name: str) -> None: + env_dir = root / "envs" / name + (env_dir / "bin").mkdir(parents=True) + (env_dir / "bin" / "python").touch() + (env_dir / "env_source.py").write_text(ENV_SOURCE) + + +def _manifest(root: Path, environments=None) -> Manifest: + return Manifest( + schema_version=SCHEMA_VERSION, + clusters=["test"], + root=str(root), + maintainer=Maintainer(name="a", email="a@b.c"), + rootstock_version="0.0.0", + python_version="3.10", + last_updated="2026-01-01T00:00:00Z", + environments=environments or {}, + ) + + +def test_refresh_preserves_existing_image_record(tmp_path: Path, monkeypatch): + monkeypatch.setattr("rootstock.operations.get_installed_versions", lambda *a, **k: {}) + _make_built_env(tmp_path, "demo") + existing = _manifest( + tmp_path, + {"demo": _env("2026-09-01T00:00:00Z", {**IMAGE, "packed_at": "2026-09-01T00:00:00Z"})}, + ) + + refreshed = refresh_manifest_environments(existing, tmp_path) + + assert refreshed.environments["demo"].image == {**IMAGE, "packed_at": "2026-09-01T00:00:00Z"} + + +def test_refresh_stamps_freshly_packed_image_current(tmp_path: Path, monkeypatch): + monkeypatch.setattr("rootstock.operations.get_installed_versions", lambda *a, **k: {}) + _make_built_env(tmp_path, "demo") + + refreshed = refresh_manifest_environments( + _manifest(tmp_path), tmp_path, built_env="demo", packed_images={"demo": IMAGE} + ) + + env = refreshed.environments["demo"] + assert env.image is not None and "packed_at" in env.image + # Stamped in the same refresh as built_at: the currency check must hold + # for an install's own pack, or every fresh install would fall back. + assert image_is_current(env) diff --git a/tests/manifest/test_manifest.py b/tests/manifest/test_manifest.py index 3d2cd55..e59970c 100644 --- a/tests/manifest/test_manifest.py +++ b/tests/manifest/test_manifest.py @@ -41,7 +41,7 @@ def _make_manifest(envs: dict[str, EnvironmentInfo] | None = None) -> Manifest: def test_schema_version_constant(): - assert SCHEMA_VERSION == 6 + assert SCHEMA_VERSION == 7 def test_checkpoint_info_round_trip(): diff --git a/tests/manifest/test_migrations.py b/tests/manifest/test_migrations.py index 26d7dfc..25098ae 100644 --- a/tests/manifest/test_migrations.py +++ b/tests/manifest/test_migrations.py @@ -84,6 +84,7 @@ def test_v2_without_checkpoints_migrates_quietly(): "migrated manifest schema v3 -> v4", "migrated manifest schema v4 -> v5", "migrated manifest schema v5 -> v6", + "migrated manifest schema v6 -> v7", ] @@ -114,6 +115,7 @@ def test_v3_drops_dead_status_fields(): "migrated manifest schema v3 -> v4", "migrated manifest schema v4 -> v5", "migrated manifest schema v5 -> v6", + "migrated manifest schema v6 -> v7", ] assert "mace" in Manifest.from_dict(migrated).environments @@ -134,6 +136,7 @@ def test_v4_bumps_cleanly_with_checkpoints_intact(): assert notes == [ "migrated manifest schema v4 -> v5", "migrated manifest schema v5 -> v6", + "migrated manifest schema v6 -> v7", ] ckpt = Manifest.from_dict(migrated).environments["mace"].checkpoints["mace-mp-0-medium"] assert ckpt.fetched_at == "2026-01-02T00:00:00Z" @@ -154,7 +157,7 @@ def test_v1_chain_migrates_to_current(): assert migrated["clusters"] == ["test"] # v1->v2 mints empty CheckpointInfo dicts; v2->v3 then drops them assert migrated["environments"]["mace"]["checkpoints"] == {} - assert len(notes) == 5 + assert len(notes) == 6 assert Manifest.from_dict(migrated).environments["mace"].source_hash == "sha256:abc" diff --git a/tests/manifest/test_shared_install.py b/tests/manifest/test_shared_install.py index ef907fc..4aa3f58 100644 --- a/tests/manifest/test_shared_install.py +++ b/tests/manifest/test_shared_install.py @@ -118,7 +118,10 @@ def test_v5_unregistered_root_migrates_to_single_cluster(): migrated, notes = migrate_manifest_data(_v5(cluster="della", root="/nowhere")) assert migrated["clusters"] == ["della"] # No sibling seeding — nothing extra to announce. - assert notes == ["migrated manifest schema v5 -> v6"] + assert notes == [ + "migrated manifest schema v5 -> v6", + "migrated manifest schema v6 -> v7", + ] def test_migrated_v5_round_trips_through_dataclasses(): diff --git a/tests/stage/conftest.py b/tests/stage/conftest.py new file mode 100644 index 0000000..9e76e42 --- /dev/null +++ b/tests/stage/conftest.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from stagelib import build_install_root + + +@pytest.fixture +def install_root(tmp_path: Path) -> Path: + return build_install_root(tmp_path) diff --git a/tests/stage/stagelib.py b/tests/stage/stagelib.py new file mode 100644 index 0000000..e05c7dc --- /dev/null +++ b/tests/stage/stagelib.py @@ -0,0 +1,66 @@ +"""Shared helpers for the pack/stage tests: a miniature install root with +one built env whose venv symlinks and pyvenv.cfg point at the root's +.python/, the way uv-built envs do on a real install.""" + +from __future__ import annotations + +import json +import os +import shutil +from pathlib import Path + +import pytest + +requires_archive_tools = pytest.mark.skipif( + shutil.which("tar") is None or shutil.which("zstd") is None, + reason="tar/zstd not on PATH", +) + +ENV_NAME = "demo" +INTERP = "cpython-3.11.99-test" + + +def build_install_root(tmp_path: Path) -> Path: + root = tmp_path / "root" + interp_bin = root / ".python" / INTERP / "bin" + interp_bin.mkdir(parents=True) + (interp_bin / "python3.11").write_text("#!/bin/sh\n") + (interp_bin / "python3.11").chmod(0o755) + + env_bin = root / "envs" / ENV_NAME / "bin" + env_bin.mkdir(parents=True) + os.symlink(interp_bin / "python3.11", env_bin / "python") + os.symlink("python", env_bin / "python3") + + env_dir = root / "envs" / ENV_NAME + (env_dir / "pyvenv.cfg").write_text( + f"home = {interp_bin}\nversion_info = 3.11.99\nrelocatable = true\n" + ) + site = env_dir / "lib" / "python3.11" / "site-packages" + site.mkdir(parents=True) + (site / "libdemo.so").write_bytes(b"\x7fELF" + b"x" * 1000) + (env_dir / "env_source.py").write_text("CHECKPOINTS = {}\n") + return root + + +def write_manifest_env( + root: Path, + env_name: str = ENV_NAME, + built_at: str = "2026-09-01T00:00:00Z", + image: dict | None = None, + checkpoints: dict | None = None, +) -> None: + (root / "manifest.json").write_text( + json.dumps( + { + "schema_version": 7, + "environments": { + env_name: { + "built_at": built_at, + "image": image, + "checkpoints": checkpoints or {}, + } + }, + } + ) + ) diff --git a/tests/stage/test_pack.py b/tests/stage/test_pack.py new file mode 100644 index 0000000..729a90b --- /dev/null +++ b/tests/stage/test_pack.py @@ -0,0 +1,137 @@ +"""Packing envs into single-image archives (#180).""" + +from __future__ import annotations + +import hashlib +import subprocess +from pathlib import Path + +import pytest +from stagelib import ENV_NAME, INTERP, requires_archive_tools + +from rootstock.pack import ( + PackError, + env_interpreter_dir, + pack_environment, + pack_environment_best_effort, +) + + +def test_interpreter_dir_resolves_through_venv_symlink(install_root: Path): + assert env_interpreter_dir(install_root, ENV_NAME) == install_root / ".python" / INTERP + + +def test_interpreter_outside_root_refuses(install_root: Path, tmp_path: Path): + foreign = tmp_path / "foreign-python" + foreign.write_text("#!/bin/sh\n") + python_link = install_root / "envs" / ENV_NAME / "bin" / "python" + python_link.unlink() + python_link.symlink_to(foreign) + with pytest.raises(PackError, match="outside"): + env_interpreter_dir(install_root, ENV_NAME) + + +def test_pack_unbuilt_env_refuses(install_root: Path): + with pytest.raises(PackError, match="not built"): + pack_environment(install_root, "nope") + + +@requires_archive_tools +def test_pack_produces_verifiable_archive(install_root: Path): + record = pack_environment(install_root, ENV_NAME) + + image = install_root / record["path"] + assert image.is_file() + assert record["format"] == "tar.zst" + assert record["sha256"] == hashlib.sha256(image.read_bytes()).hexdigest() + assert image.name == f"{ENV_NAME}-{record['sha256'][:12]}.tar.zst" + assert record["compressed_bytes"] == image.stat().st_size + assert record["uncompressed_bytes"] > 0 + # "packed_at" is stamped by the manifest refresh, not the packer. + assert "packed_at" not in record + + # The archive holds root-relative paths for the env AND its interpreter. + listing = subprocess.run( + f"zstd -dc {image} | tar -tf -", shell=True, capture_output=True, text=True, check=True + ).stdout + assert f"envs/{ENV_NAME}/pyvenv.cfg" in listing + assert f".python/{INTERP}/bin/python3.11" in listing + + +@requires_archive_tools +def test_repack_removes_superseded_images(install_root: Path): + first = pack_environment(install_root, ENV_NAME) + # Change the env so the archive bytes (and sha) differ. + site = install_root / "envs" / ENV_NAME / "lib" / "python3.11" / "site-packages" + (site / "extra.so").write_bytes(b"y" * 2048) + second = pack_environment(install_root, ENV_NAME) + + assert first["sha256"] != second["sha256"] + images = sorted(p.name for p in (install_root / "images").iterdir()) + assert images == [Path(second["path"]).name] + + +def test_best_effort_pack_degrades_to_warning(install_root: Path, capsys): + assert pack_environment_best_effort(install_root, "nope") is None + assert "rootstock pack" in capsys.readouterr().err + + +@requires_archive_tools +def test_pack_spares_live_concurrent_partials(install_root: Path): + """install's auto-pack racing a batch `rootstock pack` of the same env: + the winner must not sweep the loser's in-flight partial.""" + import subprocess as sp + + images = install_root / "images" + images.mkdir() + live = images / ".demo.packing.1" # pid 1 is always alive + live.write_bytes(b"x") + reaped = sp.Popen(["true"]) + reaped.wait() + dead = images / f".demo.packing.{reaped.pid}" + dead.write_bytes(b"x") + + pack_environment(install_root, ENV_NAME) + + assert live.exists() + assert not dead.exists() + + +@requires_archive_tools +def test_repack_spares_dash_extended_sibling_env_images(install_root: Path): + """Cleanup matches `-<12 hex>.tar.zst` exactly: packing 'demo' must + never delete 'demo-tuned-.tar.zst'.""" + images = install_root / "images" + images.mkdir() + sibling = images / "demo-tuned-0123456789ab.tar.zst" + sibling.write_bytes(b"z") + old_own = images / "demo-ba9876543210.tar.zst" + old_own.write_bytes(b"z") + + pack_environment(install_root, ENV_NAME) + + assert sibling.exists() + assert not old_own.exists() + + +def test_image_usable_requires_the_archive_on_disk(tmp_path: Path): + """The pack sweep's filter must see through a purged images/ dir — a + current-looking record with no file behind it means repack, not 'all + current'.""" + from rootstock.manifest import EnvironmentInfo + from rootstock.operations import _image_usable + + image = {"path": "images/demo-abc123def456.tar.zst", "packed_at": "2026-09-01T00:00:01Z"} + env = EnvironmentInfo( + built_at="2026-09-01T00:00:00Z", + source_hash=None, + source="", + python_requires=">=3.11", + dependencies={}, + image=image, + ) + assert not _image_usable(tmp_path, env) # record fine, file purged + target = tmp_path / image["path"] + target.parent.mkdir(parents=True) + target.write_bytes(b"z") + assert _image_usable(tmp_path, env) diff --git a/tests/stage/test_resolve_stage_base.py b/tests/stage/test_resolve_stage_base.py new file mode 100644 index 0000000..f1d162d --- /dev/null +++ b/tests/stage/test_resolve_stage_base.py @@ -0,0 +1,101 @@ +"""The staging-base resolution chain: env var > layout.json > registry, +with validation that fails closed (disabled) rather than staging badly.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +import rootstock.stage as stage +from rootstock.clusters import Cluster +from rootstock.stage import NO_STAGE_ENV, STAGE_DIR_ENV, resolve_stage_base + + +@pytest.fixture +def different_fs(monkeypatch): + """tmp_path trees share one device; pretend base and root differ.""" + monkeypatch.setattr(stage, "_same_filesystem", lambda a, b: False) + + +@pytest.fixture +def root(tmp_path: Path) -> Path: + root = tmp_path / "root" + root.mkdir() + return root + + +def _declare_layout(root: Path, stage_dir: str) -> None: + (root / "layout.json").write_text(json.dumps({"layout_version": 1, "stage_dir": stage_dir})) + + +def test_env_var_wins(monkeypatch, root: Path, tmp_path: Path, different_fs): + env_base = tmp_path / "from-env" + env_base.mkdir() + layout_base = tmp_path / "from-layout" + layout_base.mkdir() + _declare_layout(root, str(layout_base)) + monkeypatch.setenv(STAGE_DIR_ENV, str(env_base)) + assert resolve_stage_base(root) == env_base + + +def test_layout_declaration(monkeypatch, root: Path, tmp_path: Path, different_fs): + monkeypatch.delenv(STAGE_DIR_ENV, raising=False) + base = tmp_path / "local" + base.mkdir() + _declare_layout(root, str(base)) + assert resolve_stage_base(root) == base + + +def test_registry_fallback(monkeypatch, root: Path, tmp_path: Path, different_fs): + monkeypatch.delenv(STAGE_DIR_ENV, raising=False) + base = tmp_path / "registry-local" + base.mkdir() + monkeypatch.setattr( + "rootstock.clusters.CLUSTER_REGISTRY", + {"testcluster": Cluster(root=root, stage_dir=str(base))}, + ) + assert resolve_stage_base(root) == base + + +def test_nothing_declared_disables(monkeypatch, root: Path, different_fs): + monkeypatch.delenv(STAGE_DIR_ENV, raising=False) + assert resolve_stage_base(root) is None + + +def test_no_stage_env_force_disables(monkeypatch, root: Path, tmp_path: Path, different_fs): + base = tmp_path / "local" + base.mkdir() + monkeypatch.setenv(STAGE_DIR_ENV, str(base)) + monkeypatch.setenv(NO_STAGE_ENV, "1") + assert resolve_stage_base(root) is None + + +def test_unexpanded_env_var_disables(monkeypatch, root: Path, different_fs): + # $SLURM_TMPDIR outside a job: the declaration is fine, this node isn't. + monkeypatch.delenv("SLURM_TMPDIR", raising=False) + monkeypatch.setenv(STAGE_DIR_ENV, "$SLURM_TMPDIR/stage") + assert resolve_stage_base(root) is None + + +def test_env_var_expansion(monkeypatch, root: Path, tmp_path: Path, different_fs): + base = tmp_path / "jobtmp" + base.mkdir() + monkeypatch.setenv("SLURM_TMPDIR", str(base)) + monkeypatch.setenv(STAGE_DIR_ENV, "$SLURM_TMPDIR") + assert resolve_stage_base(root) == base + + +def test_missing_dir_disables(monkeypatch, root: Path, tmp_path: Path, different_fs): + monkeypatch.setenv(STAGE_DIR_ENV, str(tmp_path / "nope")) + assert resolve_stage_base(root) is None + + +def test_same_filesystem_disables(monkeypatch, root: Path, tmp_path: Path): + # No _same_filesystem patch here: base and root genuinely share a device, + # and staging onto the filesystem we're escaping must be refused. + base = tmp_path / "samefs" + base.mkdir() + monkeypatch.setenv(STAGE_DIR_ENV, str(base)) + assert resolve_stage_base(root) is None diff --git a/tests/stage/test_spawn_integration.py b/tests/stage/test_spawn_integration.py new file mode 100644 index 0000000..be25975 --- /dev/null +++ b/tests/stage/test_spawn_integration.py @@ -0,0 +1,116 @@ +"""How staging plugs into spawn_in_env: the sidecar, interpreter, cache env +vars, and prewarm toggle all follow the staged copy — and downloads and +capture runs never see a mirror.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from rootstock.spawn import DOWNLOAD_WRAPPER, WORKER_WRAPPER, spawn_in_env +from rootstock.stage import StagedSpawn, stage_for_spawn + + +@pytest.fixture +def root(tmp_path: Path) -> Path: + root = tmp_path / "root" + bin_dir = root / "envs" / "demo" / "bin" + bin_dir.mkdir(parents=True) + (bin_dir / "python").write_text("#!/bin/sh\n") + return root + + +@pytest.fixture +def staged(tmp_path: Path) -> StagedSpawn: + staged_env = tmp_path / "local" / "sha" / "envs" / "demo" + (staged_env / "bin").mkdir(parents=True) + (staged_env / "bin" / "python").write_text("#!/bin/sh\n") + mirror = tmp_path / "local" / "cache-mirror" + mirror.mkdir(parents=True) + return StagedSpawn(env_dir=staged_env, cache_base=mirror) + + +def _sidecar(spec) -> dict: + return json.loads(Path(spec.cmd[2]).read_text()) + + +def test_staged_spawn_is_fully_local(monkeypatch, root: Path, staged: StagedSpawn): + monkeypatch.setattr("rootstock.stage.stage_for_spawn", lambda *a, **k: staged) + payload = {"checkpoint": "demo-ckpt", "device": "cpu", "setup_kwargs": {}} + with spawn_in_env(root, "demo", WORKER_WRAPPER, payload) as spec: + assert spec.cmd[0] == str(staged.env_dir / "bin" / "python") + assert spec.cwd == str(staged.env_dir) + assert _sidecar(spec)["env_dir"] == str(staged.env_dir) + # caches point at the mirror, prewarm is redundant + assert spec.env["HOME"] == str(staged.cache_base / "home") + assert spec.env["XDG_CACHE_HOME"] == str(staged.cache_base / "cache") + assert spec.env["ROOTSTOCK_NO_PREWARM"] == "1" + + +def test_staged_env_without_weights_keeps_prewarm(monkeypatch, root: Path, staged: StagedSpawn): + staged.cache_base = None + monkeypatch.setattr("rootstock.stage.stage_for_spawn", lambda *a, **k: staged) + payload = {"checkpoint": "demo-ckpt", "device": "cpu", "setup_kwargs": {}} + with spawn_in_env(root, "demo", WORKER_WRAPPER, payload) as spec: + assert spec.cmd[0] == str(staged.env_dir / "bin" / "python") + # caches stay on the shared filesystem, and the prewarm still runs + # (streaming the shared weights ahead of the mmaps). + assert spec.env["HOME"] == str(root / "home") + assert "ROOTSTOCK_NO_PREWARM" not in spec.env + + +def test_custom_weights_keep_prewarm_even_when_staged(monkeypatch, root: Path, staged: StagedSpawn): + monkeypatch.setattr("rootstock.stage.stage_for_spawn", lambda *a, **k: staged) + payload = { + "checkpoint": "demo:custom", + "checkpoint_path": "/scratch/me/ft.pt", + "device": "cpu", + "setup_kwargs": {}, + } + with spawn_in_env(root, "demo", WORKER_WRAPPER, payload) as spec: + # The user's weights file still lives on the shared filesystem; the + # prewarm must stay on to stream it. + assert "ROOTSTOCK_NO_PREWARM" not in spec.env + + +def test_download_spawns_never_stage(monkeypatch, root: Path): + def boom(*a, **k): # pragma: no cover - the assertion is that it's unused + raise AssertionError("download spawns must not stage") + + monkeypatch.setattr("rootstock.stage.stage_for_spawn", boom) + payload = {"checkpoint": "demo-ckpt", "device": "cpu", "setup_kwargs": {}} + with spawn_in_env(root, "demo", DOWNLOAD_WRAPPER, payload) as spec: + assert spec.cmd[0] == str(root / "envs" / "demo" / "bin" / "python") + + +def test_unstaged_spawn_unchanged(monkeypatch, root: Path): + monkeypatch.setattr("rootstock.stage.stage_for_spawn", lambda *a, **k: None) + payload = {"checkpoint": "demo-ckpt", "device": "cpu", "setup_kwargs": {}} + with spawn_in_env(root, "demo", WORKER_WRAPPER, payload) as spec: + assert spec.cmd[0] == str(root / "envs" / "demo" / "bin" / "python") + assert spec.env["HOME"] == str(root / "home") + assert "ROOTSTOCK_NO_PREWARM" not in spec.env + + +def test_capture_spawns_stage_env_but_not_weights(monkeypatch, tmp_path: Path): + """Verify/add runs (weights_capture in the payload) may stage the env, + but must observe the shared cache — never a mirror.""" + staged_root = tmp_path / "sha" + + def fake_stage_weights(*a, **k): # pragma: no cover + raise AssertionError("capture spawns must not stage weights") + + monkeypatch.setattr("rootstock.stage.resolve_stage_base", lambda root: tmp_path) + monkeypatch.setattr("rootstock.stage.stage_env", lambda *a, **k: staged_root) + monkeypatch.setattr("rootstock.stage.stage_weights", fake_stage_weights) + + payload = { + "checkpoint": "demo-ckpt", + "weights_capture": {"result_path": "/tmp/x", "cache_root": "/shared"}, + } + staged = stage_for_spawn(tmp_path / "root", "demo", payload) + assert staged is not None + assert staged.env_dir == staged_root / "envs" / "demo" + assert staged.cache_base is None diff --git a/tests/stage/test_stage_env.py b/tests/stage/test_stage_env.py new file mode 100644 index 0000000..8e1f731 --- /dev/null +++ b/tests/stage/test_stage_env.py @@ -0,0 +1,172 @@ +"""Extracting a packed env image to a node-local dir (#180): the content- +addressed round trip, the venv fixups, warm reuse, and every fall-back-to- +prewarm edge (stale image, missing archive, no record).""" + +from __future__ import annotations + +import os +from pathlib import Path + +from stagelib import ENV_NAME, INTERP, requires_archive_tools, write_manifest_env + +from rootstock.pack import pack_environment +from rootstock.stage import stage_env + +pytestmark = requires_archive_tools + + +def _pack_and_record(install_root: Path, packed_at="2026-09-01T00:00:01Z", **kwargs) -> dict: + record = pack_environment(install_root, ENV_NAME) + record["packed_at"] = packed_at + write_manifest_env(install_root, image=record, **kwargs) + return record + + +def test_stage_extracts_and_localizes(install_root: Path, tmp_path: Path): + record = _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + + staged_root = stage_env(install_root, ENV_NAME, base) + + assert staged_root is not None + assert staged_root.name == record["sha256"] + env_dir = staged_root / "envs" / ENV_NAME + + # The venv now runs on the *staged* interpreter, not the shared one: + python = env_dir / "bin" / "python" + assert os.readlink(python) == str(staged_root / ".python" / INTERP / "bin" / "python3.11") + assert python.exists() + # relative sibling symlinks survive untouched + assert os.readlink(env_dir / "bin" / "python3") == "python" + # pyvenv.cfg's home (where the stdlib is resolved from) is local too + assert ( + f"home = {staged_root / '.python' / INTERP / 'bin'}" in (env_dir / "pyvenv.cfg").read_text() + ) + # payload files came through + assert (env_dir / "lib" / "python3.11" / "site-packages" / "libdemo.so").is_file() + # no lock/partial litter + leftovers = [p.name for p in staged_root.parent.iterdir() if p.name != record["sha256"]] + assert leftovers == [] + + +def test_second_stage_reuses_warm_copy(install_root: Path, tmp_path: Path, capsys): + _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + + first = stage_env(install_root, ENV_NAME, base) + second = stage_env(install_root, ENV_NAME, base) + + assert first == second + assert "Stage reused (warm)" in capsys.readouterr().err + + +def test_rebuilt_env_stages_beside_old_copy(install_root: Path, tmp_path: Path): + _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + first = stage_env(install_root, ENV_NAME, base) + + # Rebuild: env content changes, repack, manifest re-records. + site = install_root / "envs" / ENV_NAME / "lib" / "python3.11" / "site-packages" + (site / "v2.so").write_bytes(b"z" * 4096) + _pack_and_record( + install_root, built_at="2026-09-02T00:00:00Z", packed_at="2026-09-02T00:00:01Z" + ) + second = stage_env(install_root, ENV_NAME, base) + + assert first is not None and second is not None + assert first != second # content-addressed: new build, new dir + assert (second / "envs" / ENV_NAME / "lib" / "python3.11" / "site-packages" / "v2.so").exists() + + +def test_stale_image_falls_back(install_root: Path, tmp_path: Path): + # Env rebuilt after the pack: built_at is newer than packed_at. + _pack_and_record( + install_root, built_at="2026-09-03T00:00:00Z", packed_at="2026-09-01T00:00:01Z" + ) + base = tmp_path / "local" + base.mkdir() + assert stage_env(install_root, ENV_NAME, base) is None + + +def test_missing_archive_falls_back(install_root: Path, tmp_path: Path): + record = _pack_and_record(install_root) + (install_root / record["path"]).unlink() + base = tmp_path / "local" + base.mkdir() + assert stage_env(install_root, ENV_NAME, base) is None + + +def test_no_image_record_falls_back(install_root: Path, tmp_path: Path): + write_manifest_env(install_root, image=None) + base = tmp_path / "local" + base.mkdir() + assert stage_env(install_root, ENV_NAME, base) is None + + +def test_corrupt_archive_falls_back_and_cleans_up(install_root: Path, tmp_path: Path, capsys): + record = _pack_and_record(install_root) + (install_root / record["path"]).write_bytes(b"not a zstd stream") + base = tmp_path / "local" + base.mkdir() + + assert stage_env(install_root, ENV_NAME, base) is None + assert "falling back to prewarm" in capsys.readouterr().err + import getpass + + envs_root = base / "rootstock" / getpass.getuser() / "envs-by-hash" + assert not any(envs_root.glob("*.partial.*")) + assert not any(envs_root.glob("*.lock")) + + +def test_insufficient_space_falls_back(install_root: Path, tmp_path: Path, monkeypatch, capsys): + _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + + import shutil as _shutil + + usage = _shutil.disk_usage(base) + monkeypatch.setattr("rootstock.stage.shutil.disk_usage", lambda p: usage._replace(free=10)) + assert stage_env(install_root, ENV_NAME, base) is None + assert "free at" in capsys.readouterr().err + + +def test_fixup_failure_is_cached_per_node(install_root: Path, tmp_path: Path, monkeypatch, capsys): + """A deterministic fixup failure (e.g. targets outside every known root + spelling) must not re-pay the multi-GB extract-and-discard on every + spawn — it is noted per archive and skipped until the client changes.""" + import rootstock.stage as stage_module + + _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + calls = {"n": 0} + + def boom(*a, **k): + calls["n"] += 1 + raise stage_module._FixupError("no local interpreter mapping") + + monkeypatch.setattr(stage_module, "_fixup_staged_env", boom) + + assert stage_module.stage_env(install_root, ENV_NAME, base) is None + assert stage_module.stage_env(install_root, ENV_NAME, base) is None + + assert calls["n"] == 1 # the second call skipped extraction entirely + assert "previously failed" in capsys.readouterr().err + + +def test_transient_extraction_failure_is_not_cached(install_root: Path, tmp_path: Path, capsys): + """A corrupt read is not a deterministic failure: fixing the archive + (repack) must let the very next spawn stage again.""" + record = _pack_and_record(install_root) + base = tmp_path / "local" + base.mkdir() + image = install_root / record["path"] + good = image.read_bytes() + image.write_bytes(b"corrupt") + assert stage_env(install_root, ENV_NAME, base) is None + image.write_bytes(good) + assert stage_env(install_root, ENV_NAME, base) is not None diff --git a/tests/stage/test_stage_hygiene.py b/tests/stage/test_stage_hygiene.py new file mode 100644 index 0000000..f67c003 --- /dev/null +++ b/tests/stage/test_stage_hygiene.py @@ -0,0 +1,91 @@ +"""Node-local staging hygiene: multi-user directory permissions, eviction's +live-worker shield, and remapping through install-time mount-alias +spellings.""" + +from __future__ import annotations + +import os +import shutil +import time +from pathlib import Path + +from rootstock.stage import ( + _evict_lru, + _mark_in_use, + _remap_into_stage, + _user_stage_root, +) + +SEVEN_HOURS_AGO = time.time() - 7 * 3600 + + +def _age(path: Path) -> None: + os.utime(path, (SEVEN_HOURS_AGO, SEVEN_HOURS_AGO)) + + +def test_shared_intermediate_is_sticky_world_writable(tmp_path: Path): + """{base}/rootstock is created by whoever stages first; under their + umask it must still let every other user create a leaf — the /tmp + recipe (sticky 1777), with the per-user leaf locked to 0700.""" + old_umask = os.umask(0o077) + try: + user_root = _user_stage_root(tmp_path) + finally: + os.umask(old_umask) + shared = tmp_path / "rootstock" + assert shared.stat().st_mode & 0o7777 == 0o1777 + assert user_root.stat().st_mode & 0o7777 == 0o700 + + +def test_eviction_spares_envs_with_live_users(tmp_path: Path): + """Dir mtime alone can't shield a multi-day MD run: a staged env with a + live registered client pid must survive eviction however old it is.""" + envs_root = tmp_path / "envs-by-hash" + envs_root.mkdir() + keep = envs_root / "keep" + keep.mkdir() + + in_use = envs_root / "sha-in-use" + in_use.mkdir() + _mark_in_use(in_use) # registers our own (alive) pid + _age(in_use) + + import subprocess + + reaped = subprocess.Popen(["true"]) + reaped.wait() + dead = envs_root / "sha-dead" + (dead / ".users").mkdir(parents=True) + (dead / ".users" / str(reaped.pid)).touch() + _age(dead) + + # An unreachable target forces the scan over every candidate. + unreachable = shutil.disk_usage(envs_root).free + 10**15 + _evict_lru(envs_root, keep=keep, bytes_needed=unreachable) + + assert in_use.exists() # live client: shielded + assert not dead.exists() # dead pidfile: evicted + assert keep.exists() + + +def test_remap_resolves_install_time_alias_spelling(tmp_path: Path): + """uv bakes the install-time path spelling into symlink targets and + pyvenv.cfg; on multi-alias mounts (/eagle vs /lus/eagle) that spelling + matches neither the spawn-time root nor its resolution — the value + itself must be resolved and retried.""" + real_root = tmp_path / "lus" / "eagle" / "rootstock" + (real_root / ".python").mkdir(parents=True) + (tmp_path / "eagle").symlink_to(tmp_path / "lus" / "eagle") + alias_target = str(tmp_path / "eagle" / "rootstock" / ".python" / "cp311" / "bin" / "python") + staged = tmp_path / "staged" + + remapped = _remap_into_stage(alias_target, real_root, staged) + + assert remapped == str(staged / ".python" / "cp311" / "bin" / "python") + + +def test_remap_leaves_foreign_paths_alone(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + assert _remap_into_stage("/usr/bin/python3", root, tmp_path / "staged") is None + assert _remap_into_stage("3.11.99", root, tmp_path / "staged") is None diff --git a/tests/stage/test_weights_overlay.py b/tests/stage/test_weights_overlay.py new file mode 100644 index 0000000..05e7531 --- /dev/null +++ b/tests/stage/test_weights_overlay.py @@ -0,0 +1,228 @@ +"""The node-local weight mirror (#180): recorded files materialize locally, +everything else falls through to the shared cache via symlinks — including +the HuggingFace-hub snapshot indirection, whose relative symlinks must +resolve to the *local* blob copies. Currency is size+mtime, and a completed +overlay leaves a per-checkpoint marker that later spawns re-enter lock-free. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +from stagelib import ENV_NAME, write_manifest_env + +import rootstock.stage as stage_module +from rootstock.stage import stage_weights + +CKPT = "demo-checkpoint" + + +def _record(root: Path, entries: list[dict]) -> None: + write_manifest_env(root, checkpoints={CKPT: {"weight_files": entries}}) + + +def _write(base: Path, rel: str, data: bytes) -> Path: + path = base / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return path + + +def _bump_mtime(path: Path) -> None: + st = path.stat() + os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns + 10**9)) + + +def test_recorded_files_copy_and_siblings_symlink(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + weights = _write(root, "cache/demo/weights.pt", b"w" * 512) + config = _write(root, "cache/demo/config.json", b"{}") + _write(root, "cache/other-family/big.bin", b"b" * 128) + _record(root, [{"path": "cache/demo/weights.pt", "size": 512}]) + base = tmp_path / "local" + base.mkdir() + + mirror = stage_weights(root, None, ENV_NAME, CKPT, base) + + assert mirror is not None + staged = mirror / "cache" / "demo" / "weights.pt" + assert staged.is_file() and not staged.is_symlink() + assert staged.read_bytes() == weights.read_bytes() + # The copy carries the source's mtime — that is the staleness signal. + assert staged.stat().st_mtime_ns == weights.stat().st_mtime_ns + # Unrecorded sibling falls through to the shared copy… + linked_config = mirror / "cache" / "demo" / "config.json" + assert linked_config.is_symlink() and linked_config.resolve() == config.resolve() + # …and untouched families are one whole-directory symlink, not a walk. + other = mirror / "cache" / "other-family" + assert other.is_symlink() + # The worker's HOME must exist locally even with no recorded files there. + assert (mirror / "home").is_dir() and not (mirror / "home").is_symlink() + + +def test_second_pass_is_lock_free_and_copies_nothing(tmp_path: Path, monkeypatch): + root = tmp_path / "root" + root.mkdir() + _write(root, "cache/demo/weights.pt", b"w" * 512) + _record(root, [{"path": "cache/demo/weights.pt", "size": 512}]) + base = tmp_path / "local" + base.mkdir() + + first = stage_weights(root, None, ENV_NAME, CKPT, base) + assert first is not None + inode = (first / "cache" / "demo" / "weights.pt").stat().st_ino + + # The completed-overlay marker means the warm path never takes the + # mirror lock (the committee-demo serialization finding). + def no_lock(*a, **k): # pragma: no cover - the assertion is that it's unused + raise AssertionError("warm mirror must not take the lock") + + monkeypatch.setattr(stage_module._StageLock, "try_acquire", no_lock) + second = stage_weights(root, None, ENV_NAME, CKPT, base) + + assert second == first + # No re-copy: the atomic copy replaces inodes, so an unchanged inode + # proves nothing was rewritten. + assert (second / "cache" / "demo" / "weights.pt").stat().st_ino == inode + + +def test_same_size_source_update_invalidates_mirror(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + src = _write(root, "cache/demo/weights.pt", b"v1" * 256) + _record(root, [{"path": "cache/demo/weights.pt", "size": 512}]) + base = tmp_path / "local" + base.mkdir() + + assert stage_weights(root, None, ENV_NAME, CKPT, base) is not None + + # In-place retrain: same size, different bytes, newer mtime. Size-only + # currency would serve stale forces from the warm mirror forever. + src.write_bytes(b"v2" * 256) + _bump_mtime(src) + + mirror = stage_weights(root, None, ENV_NAME, CKPT, base) + + assert mirror is not None + assert (mirror / "cache" / "demo" / "weights.pt").read_bytes() == b"v2" * 256 + + +def test_hub_snapshot_symlinks_resolve_to_local_blobs(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + repo = "cache/huggingface/hub/models--facebook--UMA" + blob = _write(root, f"{repo}/blobs/abc123", b"weights" * 100) + snap_dir = root / repo / "snapshots" / "rev1" + snap_dir.mkdir(parents=True) + os.symlink("../../blobs/abc123", snap_dir / "model.safetensors") + _write(root, f"{repo}/refs/main", b"rev1") + _record(root, [{"path": f"{repo}/blobs/abc123", "size": blob.stat().st_size}]) + base = tmp_path / "local" + base.mkdir() + + mirror = stage_weights(root, None, ENV_NAME, CKPT, base) + + assert mirror is not None + local_blob = mirror / repo / "blobs" / "abc123" + assert local_blob.is_file() and not local_blob.is_symlink() + # The snapshot's relative symlink must land on the LOCAL blob — if + # snapshots/ were a whole-directory symlink into the shared tree, the + # worker would silently mmap the shared copy and staging would be moot. + snapshot = mirror / repo / "snapshots" / "rev1" / "model.safetensors" + assert snapshot.is_symlink() + assert snapshot.resolve() == local_blob.resolve() + assert (mirror / repo / "refs" / "main").read_bytes() == b"rev1" + + +def test_no_record_skips_overlay(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + _record(root, []) + base = tmp_path / "local" + base.mkdir() + assert stage_weights(root, None, ENV_NAME, CKPT, base) is None + + +def test_purged_recorded_file_skips_overlay(tmp_path: Path): + # A record whose file was scratch-swept is stale; redirecting caches at + # a mirror that can't materialize it would break the worker offline. + root = tmp_path / "root" + root.mkdir() + _record(root, [{"path": "cache/demo/gone.pt", "size": 10}]) + base = tmp_path / "local" + base.mkdir() + assert stage_weights(root, None, ENV_NAME, CKPT, base) is None + + +def test_split_cache_root(tmp_path: Path): + root = tmp_path / "root" + root.mkdir() + cache_root = tmp_path / "scratch-cache" + _write(cache_root, "cache/demo/weights.pt", b"w" * 64) + _record(root, [{"path": "cache/demo/weights.pt", "size": 64}]) + base = tmp_path / "local" + base.mkdir() + + mirror = stage_weights(root, cache_root, ENV_NAME, CKPT, base) + + assert mirror is not None + assert (mirror / "cache" / "demo" / "weights.pt").read_bytes() == b"w" * 64 + + +def test_warm_mirror_with_full_disk_stays_staged(tmp_path: Path, monkeypatch): + """The free-space gate counts only bytes that still need copying — a + warm mirror on a full disk must not get demoted to the shared path.""" + root = tmp_path / "root" + root.mkdir() + _write(root, "cache/demo/weights.pt", b"w" * 512) + _record(root, [{"path": "cache/demo/weights.pt", "size": 512}]) + base = tmp_path / "local" + base.mkdir() + first = stage_weights(root, None, ENV_NAME, CKPT, base) + assert first is not None + + # Wipe the marker so the pass re-evaluates copies (not the fast path), + # then report a full disk: nothing needs copying, so it must succeed. + for marker in (base / "rootstock").rglob("cache-mirror.*.ok"): + marker.unlink() + import shutil as _shutil + + usage = _shutil.disk_usage(base) + monkeypatch.setattr("rootstock.stage.shutil.disk_usage", lambda p: usage._replace(free=0)) + + assert stage_weights(root, None, ENV_NAME, CKPT, base) == first + + +def test_fallthrough_symlink_upgrades_to_copy_when_recorded(tmp_path: Path): + """A file symlinked earlier as an unrecorded sibling must become a real + copy once a checkpoint records it — a symlink's stat matches the shared + source trivially, and treating it as current would leave the worker + cold-mmapping the shared file with its prewarm switched off.""" + root = tmp_path / "root" + root.mkdir() + _write(root, "cache/demo/a.pt", b"a" * 128) + _write(root, "cache/demo/b.pt", b"b" * 128) + base = tmp_path / "local" + base.mkdir() + + # First checkpoint records only a.pt: b.pt gets the fallthrough symlink. + _record(root, [{"path": "cache/demo/a.pt", "size": 128}]) + mirror = stage_weights(root, None, ENV_NAME, CKPT, base) + assert mirror is not None + assert (mirror / "cache" / "demo" / "b.pt").is_symlink() + + # A second checkpoint records b.pt: the symlink must become a copy. + write_manifest_env( + root, + checkpoints={ + CKPT: {"weight_files": [{"path": "cache/demo/a.pt", "size": 128}]}, + "other-ckpt": {"weight_files": [{"path": "cache/demo/b.pt", "size": 128}]}, + }, + ) + mirror2 = stage_weights(root, None, ENV_NAME, "other-ckpt", base) + assert mirror2 == mirror + staged_b = mirror / "cache" / "demo" / "b.pt" + assert staged_b.is_file() and not staged_b.is_symlink() + assert staged_b.read_bytes() == b"b" * 128