[skycap] Pull a run's records and step index from W&B

examples/train_integrations/harbor_skycap/README.md

@@ -112,6 +112,30 @@ turns the index off. Logging runs on a background thread and never fails a
 step: W&B errors and timeouts are logged and counted, and at the end of
 training it waits up to two minutes for what is queued.
 
+### Pull a run
+
+`record_index pull` turns a run's artifact back into a record directory: each
+version's `step.json` from W&B, written as `index/<phase>/step-<N>.json`, and
+its record files, read straight from the mirror with your own credentials
+(through fsspec: `s3fs` for `s3://`, `gcsfs` for `gs://`; read access is
+enough). Give an alias to pull
+one version, or none to pull every version:
+
+```bash
+uv run --isolated --extra skyrl-train --extra harbor --extra skycap \
+  python -m examples.train_integrations.harbor_skycap.record_index \
+  pull my-team/my-project/skycap-records-train-<run id>:train-step-3 ./run-records
+```
+
+It fails open per record: a record whose files W&B can't fetch is reported as
+missing and left out whole, and the rest are pulled. Local-only records (no
+mirror) are in the index but have nothing to pull; the summary names where each
+is, `host:path` on the node that wrote it, for fetching it from there.
+A file already in the directory and identical is left alone, so pulling again,
+or pulling more versions into the same directory, only adds what is new. The
+command prints a summary (versions, records, missing, local-only) and exits
+non-zero only when it pulled nothing.
+
 ## How it fits
 
 | Piece | What it does |
@@ -120,7 +144,7 @@ training it waits up to two minutes for what is queued.
 | `servers.py` | The server pool: one Ray actor per server, each running a `skycap.CaptureService` on a port of its own. skycap builds how calls reach the model from the options; the integration supplies only its engine wire. |
 | `engine.py` | `SkyRLEngine`: skycap's vLLM wire on `/skyrl/v1/generate`, with packed routed experts and sampler support decoded by SkyRL's own `generate_wire`, and sessions released at `/finish_session`. |
 | `harbor_generator.py` | Per trial: create a trajectory, point the agent's `api_base` at it, run Harbor, and `finish` with the reward to get the samples. A retry gets a fresh trajectory. |
-| `record_index.py` | `SkycapRecordIndex`, a trainer callback: the per-step W&B index of the trajectories the generator logged in a `RecordLog`. Nothing in it is Harbor's. |
+| `record_index.py` | `SkycapRecordIndex`, a trainer callback: the per-step W&B index of the trajectories the generator logged in a `RecordLog`; and `pull`, which brings a run's records and index back from W&B. Nothing in it is Harbor's. |
 | `compose.py` | Samples to a step-wise `GeneratorOutput`: a trial's paths are contiguous under its `TrajectoryID`, the last one marked `is_last_step` and carrying the reward. |
 
 What's imposed on every call:

examples/train_integrations/harbor_skycap/record_index.py

@@ -33,13 +33,17 @@
 left. Every loss is logged and counted in ``stats()``.
 """
 
+import argparse
+import hashlib
 import json
+import os
 import queue
+import sys
 import re
 import tempfile
 import threading
 import time
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
 
@@ -461,3 +465,234 @@ def _transient(error: BaseException, wandb: Any) -> bool:
     if isinstance(error, (ConnectionError, TimeoutError)):
         return True
     return isinstance(error, OSError) and not isinstance(error, _PERMANENT)
+
+
+# -- pulling a run back -----------------------------------------------------------
+@dataclass
+class PullSummary:
+    """What ``pull`` brought back, and what it couldn't."""
+
+    #: ``<artifact>:<version>`` per version whose ``step.json`` was written.
+    versions: List[str] = field(default_factory=list)
+    #: Versions that couldn't be read at all, with why.
+    failed_versions: List[str] = field(default_factory=list)
+    #: Records whose files are all in ``out_dir`` now (already there and identical included).
+    records: int = 0
+    #: Files moved in; files left alone because an identical one was already there.
+    files: int = 0
+    unchanged: int = 0
+    #: ``<version>/<id>: why`` per mirrored record whose files couldn't all be fetched. Skipped, not partial.
+    missing: List[str] = field(default_factory=list)
+    #: ``<version>/<id>`` per row whose record has no mirror: indexed, with nothing in W&B to pull.
+    local_only: List[str] = field(default_factory=list)
+
+    @property
+    def pulled_anything(self) -> bool:
+        return bool(self.versions or self.records)
+
+    def report(self) -> str:
+        lines = [
+            f"versions: {len(self.versions)} pulled, {len(self.failed_versions)} failed",
+            f"records: {self.records} pulled ({self.files} files moved in, {self.unchanged} already there), "
+            f"{len(self.missing)} missing, {len(self.local_only)} local-only (indexed, nothing to pull)",
+        ]
+        lines += [f"  failed version {item}" for item in self.failed_versions]
+        lines += [f"  missing {item}" for item in self.missing]
+        lines += [f"  local-only {item}" for item in self.local_only]
+        return "\n".join(lines)
+
+
+def parse_ref(ref: str) -> Tuple[str, Optional[str]]:
+    """``entity/project/artifact[:alias]`` to ``(entity/project/artifact, alias or None)``."""
+    path, _, alias = ref.rpartition(":") if ":" in ref.rsplit("/", 1)[-1] else (ref, "", "")
+    if path.count("/") != 2 or not all(path.split("/")):
+        raise ValueError(f"expected <entity>/<project>/<artifact>[:<alias>], got {ref!r}")
+    return path, alias or None
+
+
+def pull(ref: str, out_dir: Any, *, api: Any = None) -> PullSummary:
+    """Pull a run's records and step index from W&B into ``out_dir``, a record directory.
+
+    With an alias (``...:train-step-3``, ``...:v7``), that version; without one, every version. Each
+    version's records go to ``out_dir`` and its ``step.json`` to ``out_dir/index/<phase>/step-<N>.json``.
+    A record whose files can't all be fetched is reported and skipped; the rest go on.
+
+    Args:
+        ref: ``<entity>/<project>/<artifact>[:<alias>]``.
+        out_dir: the record directory to fill; created if missing.
+        api: a ``wandb.Api``; one is made when None.
+    """
+    path, alias = parse_ref(ref)
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    if api is None:
+        import wandb
+
+        api = wandb.Api()
+    summary = PullSummary()
+    try:
+        versions = [api.artifact(f"{path}:{alias}")] if alias else list(api.artifacts(ARTIFACT_TYPE, path))
+    except Exception as error:  # noqa: BLE001 - reported; nothing else to pull
+        summary.failed_versions.append(f"{ref}: {type(error).__name__}: {error}")
+        return summary
+    for artifact in versions:
+        label = f"{path.rsplit('/', 1)[-1]}:{getattr(artifact, 'version', '?')}"
+        try:
+            _pull_version(artifact, label, out, summary)
+        except Exception as error:  # noqa: BLE001 - one version must not stop the others
+            summary.failed_versions.append(f"{label}: {type(error).__name__}: {error}")
+    return summary
+
+
+def _pull_version(artifact: Any, label: str, out: Path, summary: PullSummary) -> None:
+    # A scratch directory inside out_dir, so each file moves in with an atomic rename.
+    with tempfile.TemporaryDirectory(prefix=".pull-", dir=out) as tmp:
+        root = Path(tmp)
+        # Only the index comes from W&B. The records come from the mirror, read directly (below): W&B's own
+        # download of an s3:// reference needs s3:ListBucketVersions, which a reader of the bucket may lack.
+        artifact.get_entry("step.json").download(root=str(root))
+        step = json.loads((root / "step.json").read_bytes())
+        if not isinstance(step, dict) or not isinstance(step.get("rows"), list):
+            raise ValueError("step.json is not a run index (no rows)")
+        phase, number = step["phase"], int(step["step"])
+        if phase not in PHASES:
+            raise ValueError(f"step.json has an unknown phase {phase!r}")
+        for row in step["rows"]:
+            _fetch_record(row, root / "records")
+            _move_record(row, root / "records", out, f"{label}/{row.get('id')}", summary)
+        target = out / "index" / phase / f"step-{number}.json"
+        target.parent.mkdir(parents=True, exist_ok=True)
+        _move_in(root / "step.json", target)
+    summary.versions.append(label)
+
+
+def _fetch_record(row: Dict[str, Any], dest: Path) -> None:
+    """A mirrored record's files, read from the mirror into ``dest`` with the caller's own credentials.
+
+    A file that can't be read is left out, and ``_move_record`` reports the record missing, naming it.
+    """
+    try:
+        names = _record_files(row)
+    except ValueError:
+        return  # reported by _move_record
+    if not names:
+        return
+    import fsspec
+
+    mirror_dir = row["record"]["mirror"].rsplit("/", 1)[0]
+    dest.mkdir(parents=True, exist_ok=True)
+    for name in names:
+        uri = f"{mirror_dir}/{name}"
+        try:
+            fs, path = fsspec.core.url_to_fs(uri)
+            fs.get_file(path, str(dest / name))
+        except Exception as error:  # noqa: BLE001 - this record is reported missing
+            logger.warning(f"skycap record pull: fetching {uri} failed: {type(error).__name__}: {error}")
+
+
+def _record_files(row: Dict[str, Any]) -> List[str]:
+    """A mirrored row's file names, as ``record_references`` referenced them; [] for a local-only row.
+
+    Raises ValueError when a name isn't a plain file name: the names come from the downloaded ``step.json``,
+    and one like ``../x`` would move a file out of the directory.
+    """
+    record = row.get("record")
+    if not record or not record.get("mirror"):
+        return []
+    names = list(record.get("files") or [record["mirror"].rsplit("/", 1)[-1]])
+    unsafe = [name for name in names if not _plain_name(name)]
+    if unsafe:
+        raise ValueError(f"file names that aren't plain file names: {unsafe!r}")
+    return names
+
+
+def _plain_name(name: Any) -> bool:
+    return (
+        isinstance(name, str)
+        and name not in ("", ".", "..")
+        and "/" not in name
+        and "\\" not in name
+        and "\0" not in name
+    )
+
+
+def _move_record(row: Dict[str, Any], downloaded: Path, out: Path, label: str, summary: PullSummary) -> None:
+    record = row.get("record")
+    if not record:
+        return
+    try:
+        names = _record_files(row)
+    except ValueError as error:
+        summary.missing.append(f"{label}: not pulled: {error}")
+        return
+    if not names:
+        from skycap import RecordLocation
+
+        # Not in the mirror: say where it is, `host:path`, for whoever wants to fetch it from that node.
+        summary.local_only.append(f"{label} at {RecordLocation.from_json(record).local}")
+        return
+    absent = [name for name in names if not (downloaded / name).is_file()]
+    if absent:
+        summary.missing.append(f"{label}: could not fetch {', '.join(absent)}")
+        return
+    # Sidecars before the document, as skycap writes them, so a document in out_dir has its sidecars.
+    placed: List[Path] = []
+    moved = unchanged = 0
+    try:
+        for name in names:
+            target = out / name
+            existed = target.exists()
+            if _move_in(downloaded / name, target):
+                moved += 1
+                if not existed:
+                    placed.append(target)
+            else:
+                unchanged += 1
+    except OSError as error:
+        # One record that can't move in (say, a directory where a file goes) mustn't stop the others:
+        # take back what it newly placed, report it, and go on.
+        for path in placed:
+            path.unlink(missing_ok=True)
+        summary.missing.append(f"{label}: could not move into {out}: {type(error).__name__}: {error}")
+        return
+    summary.files += moved
+    summary.unchanged += unchanged
+    summary.records += 1
+
+
+def _move_in(source: Path, target: Path) -> bool:
+    """Rename ``source`` onto ``target`` unless an identical file is there. Returns whether it moved."""
+    if target.is_file() and target.stat().st_size == source.stat().st_size and _digest(target) == _digest(source):
+        return False
+    # The scratch directory is inside out_dir, so this is an atomic rename on one filesystem.
+    os.replace(source, target)
+    return True
+
+
+def _digest(path: Path) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for chunk in iter(lambda: handle.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def main(argv: Optional[List[str]] = None, *, api: Any = None) -> int:
+    parser = argparse.ArgumentParser(prog="python -m examples.train_integrations.harbor_skycap.record_index")
+    commands = parser.add_subparsers(dest="command", required=True)
+    pull_parser = commands.add_parser(
+        "pull", help="pull a run's records and step index from W&B into a record directory"
+    )
+    pull_parser.add_argument("ref", help="<entity>/<project>/<artifact>[:<alias>]; every version without an alias")
+    pull_parser.add_argument("out_dir", help="the record directory to fill")
+    args = parser.parse_args(argv)
+    try:
+        summary = pull(args.ref, args.out_dir, api=api)
+    except ValueError as error:
+        parser.error(str(error))
+    print(summary.report())
+    return 0 if summary.pulled_anything else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())