Node-local env staging: pack envs to single images, extract to local disk at spawn

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
 

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 <name>
+```
+
+Job scripts that spawn several calculators can pay the read once, up front:
+
+```bash
+rootstock stage uma-s-1p1 mace-mp-0-medium --cluster <name>
+```
+
+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

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 "

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 [<env> ...] [--all] [--root <path> | --cluster <name>]
+        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 <checkpoint-id> [...] [--root <path> | --cluster <name>]
+        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 <path>] [--checkpoints <id> ...] [--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/<env>-<sha>.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",

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:

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",

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'}")

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:

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