[skycap] Land the W&B record index and pull on main (#2351, #2436); pull writes a directory per phase

examples/train_integrations/harbor_skycap/README.md

@@ -89,6 +89,54 @@ Terminus-2 keeps the server's own URL. A quick tunnel takes at most 200 calls in
 flight, and a call whose reply hasn't started after about 125 s fails, so use
 `external_host` for many agents or long replies.
 
+## Records and W&B
+
+Each skycap server writes its trajectories to `skycap.record_dir` on its own
+node. `skycap.record_mirror=s3://bucket/prefix` (any fsspec URL; install `s3fs`
+or `gcsfs`) has every server also copy them there, in the background.
+`skycap.record_mirror_config` takes the mirror's options, e.g.
+`'+skycap.record_mirror_config={exclude: [experts, sampling_mask]}'` to keep
+routed experts and sampling masks out of the remote copy.
+
+With `trainer.logger=wandb`, each step is indexed as a version of the artifact
+`skycap-records-train-<run id>`, aliased `train-step-N` and `latest`. It holds
+a `step.json`, the step's run index as [`run_index.md`](run_index.md)
+specifies it:
+`format_version`, `run`, `phase` and `step`, and a row per trajectory opened in
+the step (every attempt, with `trained`, `superseded` and the record's `path`,
+`mirror` and `files`). For mirrored records it also holds a W&B reference
+`records/<name>` to each of the record's files in the mirror (no bytes are
+copied). A local-only record is in the index only. `skycap.wandb.phases=[train,eval]`
+indexes eval into `skycap-records-eval-<run id>`; `skycap.wandb.enabled=false`
+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 record directories, one
+per phase: `<out_dir>/train/` and `<out_dir>/eval/`. Each holds that phase's
+record files, read straight from the mirror with your own credentials (through
+fsspec: `s3fs` for `s3://`, `gcsfs` for `gs://`; read access is enough), and
+each step's `step.json` from W&B as `index/step-<N>.json`. Each is a plain
+skycap record directory, readable by anything that reads one. 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 |
@@ -97,6 +145,7 @@ flight, and a call whose reply hasn't started after about 125 s fails, so use
 | `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`; 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:
@@ -121,6 +170,7 @@ counted in `generate/skycap/num_missing_route_trajectories`.
 
 - **Sampler support** (`enable_return_sample_support_set`) is passed through,
   padded to `top_k`.
+- **No W&B record index under the fully-async trainer.** It fires no callbacks.
 
 ## Tests
 

examples/train_integrations/harbor_skycap/entrypoints/main_harbor_skycap.py

@@ -16,7 +16,7 @@
 import os
 import sys
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import ray
 import yaml
@@ -34,6 +34,7 @@
     _deep_merge,
 )
 from ..harbor_generator import HarborSkycapGenerator
+from ..record_index import RecordLog, SkycapRecordIndex
 from ..servers import SkycapServers, exposure_for, start_servers
 
 
@@ -53,6 +54,19 @@ class ExposureConfig:
     ``attempts`` for ``cloudflare``, or a custom class's own."""
 
 
+@dataclass
+class SkycapWandbConfig:
+    enabled: bool = True
+    """Index each step's skycap records in W&B when ``trainer.logger`` is wandb.
+
+    One version per step of the artifact ``skycap-records-<phase>-<run id>``, aliased ``<phase>-step-N`` and
+    ``latest``, holding a ``step.json`` (every attempt, trained or superseded, and where its record is) and,
+    for records in ``record_mirror``, a reference to each record file. No record bytes are uploaded.
+    Logging runs off the step and fails open."""
+    phases: List[str] = field(default_factory=lambda: ["train"])
+    """The training phases to index: ``train``, ``eval``. Each gets its own artifact."""
+
+
 @dataclass
 class SkycapConfig:
     num_servers: int = 1
@@ -64,6 +78,17 @@ class SkycapConfig:
     record_dir: Optional[str] = None
     """Where ended trajectories are written. Defaults to ``{trainer.export_path}/skycap``; each server writes
     on its own node, so point it at a shared filesystem to have one directory for the run."""
+    record_mirror: Optional[str] = None
+    """Where every server also copies its records, as an fsspec URL (``s3://bucket/prefix``).
+
+    The copy is made in the background after the record is written to ``record_dir``, and fails open: a slow
+    or failing store never fails a rollout. Needs the store's fsspec implementation (``s3fs``, ``gcsfs``)."""
+    record_mirror_config: Dict[str, Any] = field(default_factory=dict)
+    """The mirror's options (``skycap.mirror.RecordMirror``), e.g. ``{exclude: [experts, sampling_mask]}`` to
+    leave sidecars out of the remote copy, or ``timeout``, ``attempts``, ``queue_size``, ``storage_options``.
+    Needs ``record_mirror``."""
+    wandb: SkycapWandbConfig = field(default_factory=SkycapWandbConfig)
+    """The per-step index of the records in W&B."""
     ttl: float = 3600.0
     """Seconds an open trajectory may be idle before skycap writes it as abandoned and releases it."""
     renderer_pool_size: int = 8
@@ -120,6 +145,8 @@ def start_skycap(cfg: Any, engine_url: str) -> SkycapServers:
         "require_api_key": _require_api_key(cfg),
         # A custom rule is imported by each server, under the name the generator finishes with.
         "path_rules": {} if train_paths in BUILTIN_RULES else {train_paths: train_paths},
+        "record_mirror": cfg.skycap.record_mirror,
+        "record_mirror_config": dict(cfg.skycap.record_mirror_config or {}) or None,
     }
     return start_servers(
         settings,
@@ -155,19 +182,29 @@ def _exposure(cfg: Any) -> Optional[str]:
 class HarborSkycapExp(HarborExp):
     skycap: Optional[SkycapServers] = None
     generator: Optional[HarborSkycapGenerator] = None
+    records: Optional[RecordLog] = None
 
     def get_generator(self, cfg, tokenizer, inference_engine_client):
         if self.skycap is None:
             self.skycap = start_skycap(cfg, inference_engine_client.get_endpoint_url())
+        if self.records is None and cfg.skycap.wandb.enabled and cfg.trainer.logger == "wandb":
+            self.records = RecordLog()
         self.generator = HarborSkycapGenerator(
             generator_cfg=cfg.generator,
             harbor_cfg=cfg.harbor_trial_config,
             capture_urls=self.skycap.urls,
             inference_engine_client=inference_engine_client,
             train_paths=cfg.skycap.train_paths,
+            records=self.records,
         )
         return self.generator
 
+    def get_trainer(self, *args, **kwargs):
+        trainer = super().get_trainer(*args, **kwargs)
+        if self.records is not None:
+            trainer.add_callback(SkycapRecordIndex(self.records, self.cfg.skycap.wandb.phases))
+        return trainer
+
     def run(self):
         try:
             super().run()

examples/train_integrations/harbor_skycap/harbor_generator.py

@@ -49,6 +49,7 @@
 
 from ..harbor.trial_metrics import TrialAttempts, trial_metrics
 from .compose import TrialOutcome, compose, split
+from .record_index import RecordLog
 
 litellm.suppress_debug_info = True
 
@@ -67,6 +68,7 @@ def __init__(
         capture_urls: List[str],
         inference_engine_client: Any = None,
         train_paths: str = "all",
+        records: Optional[RecordLog] = None,
     ) -> None:
         """
         Args:
@@ -76,6 +78,7 @@ def __init__(
             inference_engine_client: read for its ``weight_version``, which keys the prefix-cache salt.
             train_paths: the skycap path rule every trajectory is finished with: ``all``, ``final``, or a custom
                 rule's ``"pkg.module:function"``, which the servers must have been started with.
+            records: where every trajectory opened is logged, per phase, for the W&B record index.
         """
         # Imported here too, so a bad import path fails at startup rather than at the first finish.
         load_rule(train_paths)
@@ -95,6 +98,7 @@ def __init__(
         self.capture_urls = list(capture_urls)
         self.pool = CapturePool(self.capture_urls)
         self.inference_engine_client = inference_engine_client
+        self.records = records
         served = generator_cfg.inference_engine.served_model_name
         if served is None or "/" in served:
             raise ValueError("generator.inference_engine.served_model_name must be set, without '/'")
@@ -135,6 +139,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             raise ValueError(f"Prompt count ({len(prompts)}) doesn't match trajectory_ids ({len(trajectory_ids)})")
         metadata = input_batch.get("batch_metadata")
         step = getattr(metadata, "global_step", None)
+        phase = getattr(metadata, "training_phase", "train")
         cache_salt = self._cache_salt()
 
         outcomes: List[Optional[TrialOutcome]] = [None] * len(prompts)
@@ -148,7 +153,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         )
 
         async def worker(index: int, prompt: ConversationType, trajectory_id: TrajectoryID) -> None:
-            outcomes[index] = await self._trial(prompt, trajectory_id, cache_salt, step, attempts[index])
+            outcomes[index] = await self._trial(prompt, trajectory_id, cache_salt, step, phase, attempts[index])
             progress.update(1)
 
         try:
@@ -174,6 +179,7 @@ async def _trial(
         trajectory_id: TrajectoryID,
         cache_salt: Optional[str],
         step: Optional[int],
+        phase: str,
         attempts: TrialAttempts,
     ) -> TrialOutcome:
         """One rollout, retried on unknown errors. Never raises: one failure must not cancel the batch.
@@ -187,7 +193,7 @@ async def _trial(
             prefix = f"Trajectory {trajectory_id} attempt {attempt + 1}/{MAX_NUM_RETRIES_PER_TRIAL}"
             attempts.start()
             try:
-                outcome = await self._attempt(prompt, trajectory_id, cache_salt, step, attempt, attempts)
+                outcome = await self._attempt(prompt, trajectory_id, cache_salt, step, phase, attempt, attempts)
             except PathRuleError as error:
                 logger.error(f"{prefix}: path rule {self.train_paths!r} failed, not retrying: {error}")
                 attempts.fail(error)
@@ -214,6 +220,7 @@ async def _attempt(
         trajectory_id: TrajectoryID,
         cache_salt: Optional[str],
         step: Optional[int],
+        phase: str,
         attempt: int,
         attempts: TrialAttempts,
     ) -> TrialOutcome:
@@ -225,25 +232,31 @@ async def _attempt(
             "step": step,
             "attempt": attempt,
         }
-        async with self.pool.trajectory(meta, paths=self.train_paths) as trajectory:
-            config = self._trial_config(prompt, self._agent_url(trajectory), cache_salt, trajectory.api_key)
-            async with self._rate_limiter:
-                results = await (await Trial.create(TrialConfig.model_validate(config))).run()
+        trajectory = None
+        try:
+            async with self.pool.trajectory(meta, paths=self.train_paths) as trajectory:
+                config = self._trial_config(prompt, self._agent_url(trajectory), cache_salt, trajectory.api_key)
+                async with self._rate_limiter:
+                    results = await (await Trial.create(TrialConfig.model_validate(config))).run()
 
-            attempts.record(results)
-            exception = results.exception_info.exception_type if results.exception_info else None
-            if exception == "AgentTimeoutError":
-                # Masked, not retried, as the sibling does.
-                reward, stop_reason = 0.0, "agent_timeout"
-            elif exception == "ContextLengthExceededError":
-                # Trains with reward 0, as the sibling does.
-                reward, stop_reason = 0.0, "context_length"
-            elif not results.verifier_result:
-                reward, stop_reason = 0.0, "error"
-                logger.warning(f"Trajectory {trajectory_id} has no verifier result: {results.exception_info}")
-            else:
-                reward, stop_reason = float(results.verifier_result.rewards["reward"]), "complete"
-            finished = await trajectory.finish({"reward": reward, "stop_reason": stop_reason})
+                attempts.record(results)
+                exception = results.exception_info.exception_type if results.exception_info else None
+                if exception == "AgentTimeoutError":
+                    # Masked, not retried, as the sibling does.
+                    reward, stop_reason = 0.0, "agent_timeout"
+                elif exception == "ContextLengthExceededError":
+                    # Trains with reward 0, as the sibling does.
+                    reward, stop_reason = 0.0, "context_length"
+                elif not results.verifier_result:
+                    reward, stop_reason = 0.0, "error"
+                    logger.warning(f"Trajectory {trajectory_id} has no verifier result: {results.exception_info}")
+                else:
+                    reward, stop_reason = float(results.verifier_result.rewards["reward"]), "complete"
+                finished = await trajectory.finish({"reward": reward, "stop_reason": stop_reason})
+        finally:
+            # Every attempt that got a trajectory, including one that raised: the pool finished it.
+            if self.records is not None and trajectory is not None:
+                self.records.add(phase, trajectory_id, attempt, trajectory)
 
         if finished.status != "finished":
             # The trajectory failed inside skycap (e.g. an unattributable prompt): its samples may miss a turn.