diff --git a/examples/train_integrations/harbor_skycap/README.md b/examples/train_integrations/harbor_skycap/README.md index 22f0f6246a..ef963a80ae 100644 --- a/examples/train_integrations/harbor_skycap/README.md +++ b/examples/train_integrations/harbor_skycap/README.md @@ -27,7 +27,9 @@ uv run --isolated --extra fsdp --extra harbor --extra skycap \ The rest of the configuration is the sibling's: `harbor_trial_config` holds Harbor's `TrialConfig`, with defaults from `../harbor/harbor_trial_config/default.yaml`. `skycap.*` sets the record directory (default `{trainer.export_path}/skycap`), -the idle TTL, the port and the renderer pool size. +the idle TTL, the port and the renderer pool size. With `trainer.logger=wandb`, each step's +documents go to W&B, without sidecars, as a version of `skycap-records-train-`, aliased +`step-N` and `latest`. `skycap.wandb.phases` adds `eval`; `skycap.wandb.enabled=false` turns it off. ## How it fits @@ -37,6 +39,7 @@ the idle TTL, the port and the renderer pool size. | `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. | +| `artifacts.py` | `SkycapUploads`, a trainer callback: after each step it fetches every attempt's document from its server and uploads them on a background thread, with a `step.json` that marks each one trained or dropped. An upload failure is logged and does not fail the step. | | `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: @@ -56,6 +59,7 @@ Masking is the sibling's: `enable_return_routed_experts=true`, and the trainer change is a follow-up. - **Sampler support** (`enable_return_sample_support_set`) is passed through, padded to `top_k`. +- **No W&B uploads under the fully-async trainer.** It fires no callbacks. - **One skycap server per run.** The generator takes a list of URLs and spreads trajectories over them, for when servers are launched separately. diff --git a/examples/train_integrations/harbor_skycap/artifacts.py b/examples/train_integrations/harbor_skycap/artifacts.py new file mode 100644 index 0000000000..4b11e83152 --- /dev/null +++ b/examples/train_integrations/harbor_skycap/artifacts.py @@ -0,0 +1,162 @@ +"""Each step's skycap documents, uploaded to W&B as one version of a ``skycap-records--`` artifact.""" + +import json +import re +import tempfile +import time +import urllib.request +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional + +import orjson +import zstandard +from loguru import logger + +from skyrl.train.utils.callbacks import CallbackInput, TrainingCallback, TrainingControl + +ARTIFACT_TYPE = "skycap-records" +FETCH_TIMEOUT = 60.0 +UPLOAD_ATTEMPTS = 3 +UPLOAD_TIMEOUT = 600.0 +RETRY_DELAY = 10.0 + + +@dataclass +class Created: + """One skycap trajectory the generator opened: a trial attempt.""" + + server: str + id: str + instance_id: str + repetition_id: int + attempt: int + + +class RecordLog: + """The trajectories created since the last upload, per training phase.""" + + def __init__(self) -> None: + self._created: Dict[str, List[Created]] = {} + + def add(self, phase: str, created: Created) -> None: + self._created.setdefault(phase, []).append(created) + + def take(self, phase: str) -> List[Created]: + return self._created.pop(phase, []) + + +def artifact_name(phase: str, run_id: str) -> str: + return f"{ARTIFACT_TYPE}-{phase}-{re.sub(r'[^a-zA-Z0-9_.-]', '-', run_id)}" + + +def fetch_document(created: Created) -> Optional[dict]: + """The stored document, from the server that wrote it; None unless it is on that server's disk.""" + url = f"{created.server}/trajectories/{created.id}" + try: + with urllib.request.urlopen(url, timeout=FETCH_TIMEOUT) as response: + document = json.load(response) + except Exception as error: # noqa: BLE001 - counted as missing + logger.warning(f"skycap artifact: fetching {url} failed: {type(error).__name__}: {error}") + return None + # A document still in memory (its write failed) has no format_version and no sidecars manifest. + return document if "format_version" in document else None + + +def upload( + wandb: Any, run_id: str, phase: str, step: int, created: List[Created], trained: Optional[set] +) -> Optional[str]: + """Upload the documents and a ``step.json`` index, aliased ``step-N`` and ``latest``. Returns the artifact name.""" + last_attempt: Dict[tuple, int] = {} + for entry in created: + key = (entry.instance_id, entry.repetition_id) + last_attempt[key] = max(last_attempt.get(key, -1), entry.attempt) + index = [] + with tempfile.TemporaryDirectory() as tmp: + artifact = wandb.Artifact(name=artifact_name(phase, run_id), type=ARTIFACT_TYPE) + for entry in created: + document = fetch_document(entry) + key = (entry.instance_id, entry.repetition_id) + row = asdict(entry) + del row["server"] + row["uploaded"] = document is not None + # A superseded attempt never trains; without dynamic sampling every final attempt does. + row["trained"] = entry.attempt == last_attempt[key] and (trained is None or key in trained) + index.append(row) + if document is None: + continue + path = Path(tmp) / f"{entry.id}.json.zst" + path.write_bytes(zstandard.ZstdCompressor().compress(orjson.dumps(document))) + artifact.add_file(str(path), name=path.name) + uploaded = sum(row["uploaded"] for row in index) + missing = len(index) - uploaded + if missing: + logger.warning(f"skycap artifact: {missing} of {len(index)} {phase} documents of step {step} are missing") + if not uploaded: + return None + (Path(tmp) / "step.json").write_bytes(orjson.dumps(index)) + artifact.add_file(str(Path(tmp) / "step.json"), name="step.json") + artifact.metadata.update( + { + "global_step": step, + "training_phase": phase, + "run_id": run_id, + "contents": "documents", + "num_uploaded": uploaded, + "num_missing": missing, + "num_trained": sum(row["trained"] for row in index), + } + ) + wandb.log_artifact(artifact, aliases=[f"step-{step}", "latest"]).wait(timeout=UPLOAD_TIMEOUT) + logger.info(f"skycap artifact {artifact.name}:step-{step}: {uploaded} documents") + return artifact.name + + +def upload_with_retries(*args: Any) -> Optional[str]: + """``upload``, retried from the fetch on any exception, up to ``UPLOAD_ATTEMPTS`` times.""" + for attempt in range(1, UPLOAD_ATTEMPTS + 1): + try: + return upload(*args) + except Exception as error: # noqa: BLE001 - retried, then raised + if attempt == UPLOAD_ATTEMPTS: + raise + logger.warning( + f"skycap artifact: attempt {attempt}/{UPLOAD_ATTEMPTS} failed: {type(error).__name__}: {error}" + ) + time.sleep(RETRY_DELAY * attempt) + + +class SkycapUploads(TrainingCallback): + """Uploads the train records after each step and the eval records after each eval pass, off the step.""" + + def __init__(self, records: RecordLog, phases: List[str]) -> None: + self.records = records + self.phases = set(phases) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="skycap-artifacts") + + def _submit(self, trainer: Any, phase: str, step: int, trained: Optional[set]) -> None: + created = self.records.take(phase) + tracker = trainer.tracker + if phase not in self.phases or not created or tracker is None or tracker.backend != "wandb": + return + wandb = tracker.logger + future = self._executor.submit(upload_with_retries, wandb, wandb.run.id, phase, step, created, trained) + future.add_done_callback(_log_failure) + + def on_step_end(self, trainer: Any, callback_input: CallbackInput, control: TrainingControl) -> None: + ids = callback_input.trajectory_ids + trained = None if ids is None else {(str(t.instance_id), t.repetition_id) for t in ids} + self._submit(trainer, "train", callback_input.global_step, trained) + + def on_eval_end(self, trainer: Any, callback_input: CallbackInput, control: TrainingControl) -> None: + self._submit(trainer, "eval", callback_input.global_step, set()) + + def on_train_end(self, trainer: Any, callback_input: CallbackInput, control: TrainingControl) -> None: + # The tracker finishes the W&B run right after this event. + self._executor.shutdown(wait=True) + + +def _log_failure(future: Future) -> None: + if future.exception() is not None: + logger.opt(exception=future.exception()).error("uploading skycap records failed") diff --git a/examples/train_integrations/harbor_skycap/entrypoints/main_harbor_skycap.py b/examples/train_integrations/harbor_skycap/entrypoints/main_harbor_skycap.py index dc2d5774f0..28140b2d96 100644 --- a/examples/train_integrations/harbor_skycap/entrypoints/main_harbor_skycap.py +++ b/examples/train_integrations/harbor_skycap/entrypoints/main_harbor_skycap.py @@ -13,7 +13,7 @@ import os import sys from dataclasses import dataclass, field -from typing import Any, Optional +from typing import Any, List, Optional import ray import yaml @@ -28,10 +28,20 @@ HarborSkyRLConfig, _deep_merge, ) +from ..artifacts import RecordLog, SkycapUploads from ..harbor_generator import HarborSkycapGenerator from ..servers import SkycapServers, start_servers +@dataclass +class SkycapWandbConfig: + enabled: bool = True + """Upload each step's documents, without sidecars, as a version of ``skycap-records--``, + when ``trainer.logger`` is wandb. The trainer fetches them from the skycap servers.""" + phases: List[str] = field(default_factory=lambda: ["train"]) + """The training phases to upload: ``train``, ``eval``.""" + + @dataclass class SkycapConfig: num_servers: int = 1 @@ -43,6 +53,7 @@ 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.""" + wandb: SkycapWandbConfig = field(default_factory=SkycapWandbConfig) 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 @@ -59,6 +70,10 @@ class HarborSkycapConfig(HarborSkyRLConfig): skycap: SkycapConfig = field(default_factory=SkycapConfig) +def record_dir(cfg: Any) -> str: + return cfg.skycap.record_dir or os.path.join(cfg.trainer.export_path, "skycap") + + def start_skycap(cfg: Any, engine_url: str) -> SkycapServers: """skycap servers in token mode, in front of SkyRL's router.""" ie = cfg.generator.inference_engine @@ -86,24 +101,34 @@ def start_skycap(cfg: Any, engine_url: str) -> SkycapServers: num_servers=cfg.skycap.num_servers, num_cpus_per_server=cfg.skycap.num_cpus_per_server, placement_strategy=cfg.skycap.placement_strategy, - record_dir=cfg.skycap.record_dir or os.path.join(cfg.trainer.export_path, "skycap"), + record_dir=record_dir(cfg), ttl=cfg.skycap.ttl, ) class HarborSkycapExp(HarborExp): skycap: Optional[SkycapServers] = 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: + self.records = RecordLog() return HarborSkycapGenerator( generator_cfg=cfg.generator, harbor_cfg=cfg.harbor_trial_config, capture_urls=self.skycap.urls, inference_engine_client=inference_engine_client, + records=self.records, ) + def get_trainer(self, *args, **kwargs): + trainer = super().get_trainer(*args, **kwargs) + if self.records is not None: + trainer.add_callback(SkycapUploads(self.records, self.cfg.skycap.wandb.phases)) + return trainer + def run(self): try: super().run() diff --git a/examples/train_integrations/harbor_skycap/harbor_generator.py b/examples/train_integrations/harbor_skycap/harbor_generator.py index b1fce2eab4..f4b0dd69de 100644 --- a/examples/train_integrations/harbor_skycap/harbor_generator.py +++ b/examples/train_integrations/harbor_skycap/harbor_generator.py @@ -35,6 +35,7 @@ from skyrl.train.generators.utils import build_vllm_cache_salt from skyrl.train.utils.rate_limiter import create_rate_limiter +from .artifacts import Created, RecordLog from .compose import TrialOutcome, compose, split litellm.suppress_debug_info = True @@ -53,6 +54,7 @@ def __init__( harbor_cfg: Dict[str, Any], capture_urls: List[str], inference_engine_client: Any = None, + records: Optional[RecordLog] = None, ) -> None: """ Args: @@ -60,6 +62,7 @@ def __init__( harbor_cfg: Harbor's ``TrialConfig`` template. capture_urls: the skycap servers to spread trajectories over. inference_engine_client: read for its ``weight_version``, which keys the prefix-cache salt. + records: the log each created trajectory goes into, for the W&B uploads. """ if not getattr(generator_cfg, "step_wise_trajectories", False): raise ValueError( @@ -80,6 +83,7 @@ def __init__( self.generator_cfg = generator_cfg self.capture_urls = list(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 '/'") @@ -110,6 +114,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) @@ -125,7 +130,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False async with CapturePool(self.capture_urls) as pool: async def worker(index: int, prompt: ConversationType, trajectory_id: TrajectoryID) -> None: - outcomes[index] = await self._trial(pool, prompt, trajectory_id, cache_salt, step) + outcomes[index] = await self._trial(pool, prompt, trajectory_id, cache_salt, step, phase) progress.update(1) try: @@ -149,13 +154,14 @@ async def _trial( trajectory_id: TrajectoryID, cache_salt: Optional[str], step: Optional[int], + phase: str, ) -> TrialOutcome: """One rollout, retried on unknown errors. Never raises: one failure must not cancel the batch.""" started = time.monotonic() for attempt in range(MAX_NUM_RETRIES_PER_TRIAL): prefix = f"Trajectory {trajectory_id} attempt {attempt + 1}/{MAX_NUM_RETRIES_PER_TRIAL}" try: - outcome = await self._attempt(pool, prompt, trajectory_id, cache_salt, step, attempt) + outcome = await self._attempt(pool, prompt, trajectory_id, cache_salt, step, phase, attempt) except Exception as error: # noqa: BLE001 - retried, then masked logger.warning(f"{prefix} failed: {type(error).__name__}: {error}") continue @@ -172,6 +178,7 @@ async def _attempt( trajectory_id: TrajectoryID, cache_salt: Optional[str], step: Optional[int], + phase: str, attempt: int, ) -> TrialOutcome: """One attempt on its own trajectory, so a retry never shares a graph with the attempt it replaces.""" @@ -183,6 +190,11 @@ async def _attempt( "attempt": attempt, } async with pool.trajectory(meta) as trajectory: + if self.records is not None: + self.records.add( + phase, + Created(trajectory.server, trajectory.id, meta["instance_id"], meta["repetition_id"], attempt), + ) config = self._trial_config(prompt, trajectory.base_url, cache_salt) async with self._rate_limiter: results = await (await Trial.create(TrialConfig.model_validate(config))).run() diff --git a/examples/train_integrations/harbor_skycap/run_codecontests_modal.sh b/examples/train_integrations/harbor_skycap/run_codecontests_modal.sh index 9ba8b32a4b..5f659bf4b3 100755 --- a/examples/train_integrations/harbor_skycap/run_codecontests_modal.sh +++ b/examples/train_integrations/harbor_skycap/run_codecontests_modal.sh @@ -24,7 +24,8 @@ # LOGGER=console|wandb W&B (project harbor-compare, run named after the experiment) by # default when a key is found, else the console # -# Needs: the GPUs, and a file holding `MODAL_TOKEN_ID=... MODAL_TOKEN_SECRET=...`. For W&B, +# Needs: the GPUs, and MODAL_TOKEN_ID and MODAL_TOKEN_SECRET, a file holding them, or +# ~/.modal.toml (single node only). For W&B, # WANDB_API_KEY, or a file holding `WANDB_API_KEY=...` (WANDB_KEY_FILE). set -euo pipefail @@ -58,7 +59,7 @@ DATASET="${DATASET:-open-thoughts/CodeContests}" DATA_ROOT="/tmp/harbor/data" TASKS_DIR="$DATA_ROOT/$(basename "$DATASET")" -EXPERIMENT="harbor-$GENERATOR-$(python3 -c 'import uuid; print(uuid.uuid4().hex[:12])')" +EXPERIMENT="${EXPERIMENT:-harbor-$GENERATOR-$(date -u +%Y%m%d-%H%M%S)-$(python3 -c 'import uuid; print(uuid.uuid4().hex[:6])')}" RUN_DIR="/tmp/harbor/runs/$EXPERIMENT" SUBSET_DIR="$RUN_DIR/tasks" @@ -96,16 +97,20 @@ fi #----------------------- # Credentials #----------------------- -if [[ ! -f "$MODAL_KEY_FILE" ]]; then - echo "no Modal credentials at $MODAL_KEY_FILE (set MODAL_KEY_FILE)" >&2 +# The Modal SDK reads ~/.modal.toml itself, but only on this node. +if [[ -z "${MODAL_TOKEN_ID:-}" && -f "$MODAL_KEY_FILE" ]]; then + set -a + # shellcheck disable=SC1090 + source "$MODAL_KEY_FILE" + set +a +fi +if [[ -n "${MODAL_TOKEN_ID:-}" ]]; then + : "${MODAL_TOKEN_SECRET:?set MODAL_TOKEN_SECRET with MODAL_TOKEN_ID}" + export MODAL_TOKEN_ID MODAL_TOKEN_SECRET +elif [[ ! -f "$HOME/.modal.toml" ]]; then + echo "no Modal credentials in the environment, at $MODAL_KEY_FILE or in ~/.modal.toml" >&2 exit 1 fi -set -a -# shellcheck disable=SC1090 -source "$MODAL_KEY_FILE" -set +a -: "${MODAL_TOKEN_ID:?$MODAL_KEY_FILE must set MODAL_TOKEN_ID}" -: "${MODAL_TOKEN_SECRET:?$MODAL_KEY_FILE must set MODAL_TOKEN_SECRET}" # SkyRL reads the W&B key from the environment only. if [[ -z "${WANDB_API_KEY:-}" && -f "$WANDB_KEY_FILE" ]]; then diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 80a81efe8a..c18f20de60 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -444,7 +444,12 @@ async def train(self): # One profiler step per RL global step. self._profiler_step() - self._fire("on_step_end", batch=training_input, metrics=status) + self._fire( + "on_step_end", + batch=training_input, + metrics=status, + trajectory_ids=generator_output.get("trajectory_ids"), + ) step_started = False # Capture callback-driven triggers, then reset. diff --git a/skyrl/train/utils/callbacks.py b/skyrl/train/utils/callbacks.py index 731815ce9c..ed260c3f67 100644 --- a/skyrl/train/utils/callbacks.py +++ b/skyrl/train/utils/callbacks.py @@ -29,6 +29,7 @@ def on_event(self, trainer, callback_input, control): if TYPE_CHECKING: from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch + from skyrl.train.generators.base import TrajectoryID @dataclass @@ -49,6 +50,9 @@ class CallbackInput: # Step events batch: Optional["TrainingInputBatch"] = None + # on_step_end only - the trajectories behind ``batch``, one per row + trajectory_ids: Optional[List["TrajectoryID"]] = None + # Step end / eval end metrics: Optional[Dict[str, Any]] = None diff --git a/tests/integrations/harbor_skycap/test_harbor_skycap.py b/tests/integrations/harbor_skycap/test_harbor_skycap.py index 341f0bdc97..ca010d9b51 100644 --- a/tests/integrations/harbor_skycap/test_harbor_skycap.py +++ b/tests/integrations/harbor_skycap/test_harbor_skycap.py @@ -2,6 +2,7 @@ server in token mode, which calls a mock SkyRL router.""" import asyncio +import json from types import SimpleNamespace import numpy as np @@ -12,13 +13,19 @@ pytest.importorskip("harbor") import aiohttp # noqa: E402 +import orjson # noqa: E402 import pytest_asyncio # noqa: E402 +import wandb # noqa: E402 +import zstandard # noqa: E402 from aiohttp.test_utils import TestServer # noqa: E402 from examples.train_integrations.harbor.entrypoints.main_harbor import ( HARBOR_DEFAULT_CONFIG, # noqa: E402 ) -from examples.train_integrations.harbor_skycap import harbor_generator # noqa: E402 +from examples.train_integrations.harbor_skycap import ( # noqa: E402 + artifacts, + harbor_generator, +) from examples.train_integrations.harbor_skycap.compose import ( # noqa: E402 TrialOutcome, compose, @@ -33,6 +40,7 @@ pack_sample_support, # noqa: E402 ) from skyrl.train.generators.base import TrajectoryID # noqa: E402 +from skyrl.train.utils.callbacks import CallbackInput # noqa: E402 from skyrl.train.utils.trainer_utils import validate_generator_output # noqa: E402 from tests.integrations.harbor_skycap.fakes import ( # noqa: E402 EXPERTS_PER_TOKEN, @@ -123,9 +131,9 @@ def batch(*scripts: str, repetitions: int = 1) -> dict: return {"prompts": prompts, "trajectory_ids": ids, "batch_metadata": SimpleNamespace(global_step=3)} -def generator(skycap, **cfg) -> HarborSkycapGenerator: +def generator(skycap, records=None, **cfg) -> HarborSkycapGenerator: engine_client = SimpleNamespace(weight_version=7) - return HarborSkycapGenerator(generator_cfg(**cfg), harbor_cfg(), [skycap.url], engine_client) + return HarborSkycapGenerator(generator_cfg(**cfg), harbor_cfg(), [skycap.url], engine_client, records) @pytest.mark.asyncio @@ -215,6 +223,139 @@ async def test_a_trial_with_no_captured_tokens_is_retried_then_masked_not_reward assert len(trials.configs) == harbor_generator.MAX_NUM_RETRIES_PER_TRIAL + 1 +class FakeWandb: + """The ``wandb`` module as the tracker holds it: a run, ``Artifact`` and ``log_artifact``.""" + + Artifact = wandb.Artifact + + def __init__(self, failures=0) -> None: + self.run = SimpleNamespace(id="run/1") + self.logged = [] + self.failures = failures + + def log_artifact(self, artifact, aliases): + if self.failures: + self.failures -= 1 + raise ConnectionError("W&B unreachable") + self.logged.append((artifact, aliases)) + return SimpleNamespace(wait=lambda timeout: artifact) + + +def uploads(records, phases=("train",), backend="wandb", failures=0): + fake = FakeWandb(failures) + trainer = SimpleNamespace(tracker=SimpleNamespace(backend=backend, logger=fake)) + return artifacts.SkycapUploads(records, list(phases)), trainer, fake + + +def event(step=3, trajectory_ids=None) -> CallbackInput: + return CallbackInput(global_step=step, epoch=0, total_steps=9, steps_per_epoch=9, trajectory_ids=trajectory_ids) + + +def step_index(artifact) -> list: + return json.loads(open(artifact.manifest.entries["step.json"].local_path).read()) + + +@pytest.mark.asyncio +async def test_each_step_uploads_every_attempt_and_marks_the_trained_ones(skycap, trials) -> None: + records = artifacts.RecordLog() + out = await generator(skycap, records=records).generate(batch("linear", "crash"), disable_tqdm=True) + callback, trainer, fake = uploads(records) + callback.on_step_end(trainer, event(trajectory_ids=out["trajectory_ids"]), None) + callback.on_train_end(trainer, event(), None) + + ((artifact, aliases),) = fake.logged + assert artifact.name == "skycap-records-train-run-1" and artifact.type == "skycap-records" + assert aliases == ["step-3", "latest"] + ids = list(record.list_ids(skycap.server.record_dir)) + assert set(artifact.manifest.entries) == {f"{i}.json.zst" for i in ids} | {"step.json"} + rows = {(r["instance_id"], r["attempt"]): r for r in step_index(artifact)} + # The first crash attempt was retried, so only its second attempt is in the batch. + assert {key: r["trained"] for key, r in rows.items()} == { + ("linear", 0): True, + ("crash", 0): False, + ("crash", 1): True, + } + assert artifact.metadata == { + "global_step": 3, + "training_phase": "train", + "run_id": "run/1", + "contents": "documents", + "num_uploaded": 3, + "num_missing": 0, + "num_trained": 2, + } + # The upload is the document the server wrote, recompressed. + (linear,) = [r["id"] for r in step_index(artifact) if r["instance_id"] == "linear"] + uploaded = artifact.manifest.entries[f"{linear}.json.zst"].local_path + stored = record.read_document(skycap.server.record_dir, linear) + assert orjson.loads(zstandard.ZstdDecompressor().decompress(open(uploaded, "rb").read())) == stored + + +@pytest.mark.asyncio +async def test_a_trajectory_dropped_by_dynamic_sampling_is_uploaded_untrained(skycap, trials) -> None: + records = artifacts.RecordLog() + out = await generator(skycap, records=records).generate(batch("linear", "summarize"), disable_tqdm=True) + kept = [t for t in out["trajectory_ids"] if t.instance_id == "linear"] + callback, trainer, fake = uploads(records) + callback.on_step_end(trainer, event(trajectory_ids=kept), None) + callback.on_train_end(trainer, event(), None) + + ((artifact, _),) = fake.logged + assert {r["instance_id"]: r["trained"] for r in step_index(artifact)} == {"linear": True, "summarize": False} + + +@pytest.mark.asyncio +async def test_eval_uploads_to_its_own_artifact_only_when_asked(skycap, trials) -> None: + eval_batch = {**batch("linear"), "batch_metadata": SimpleNamespace(global_step=3, training_phase="eval")} + for phases, expected in ((["train"], []), (["train", "eval"], ["skycap-records-eval-run-1"])): + records = artifacts.RecordLog() + await generator(skycap, records=records).generate(eval_batch, disable_tqdm=True) + callback, trainer, fake = uploads(records, phases) + callback.on_eval_end(trainer, event(), None) + callback.on_train_end(trainer, event(), None) + assert [artifact.name for artifact, _ in fake.logged] == expected + assert records.take("eval") == [] + ((artifact, _),) = fake.logged + assert artifact.metadata["num_trained"] == 0 + assert not any(row["trained"] for row in step_index(artifact)) + + +@pytest.mark.asyncio +async def test_a_document_the_server_does_not_have_is_counted_missing(skycap, trials) -> None: + records = artifacts.RecordLog() + await generator(skycap, records=records).generate(batch("linear"), disable_tqdm=True) + records.add("train", artifacts.Created(skycap.url, "tr_nope", "gone", 0, 0)) + callback, trainer, fake = uploads(records) + callback.on_step_end(trainer, event(), None) + callback.on_train_end(trainer, event(), None) + + ((artifact, _),) = fake.logged + assert artifact.metadata["num_uploaded"] == 1 and artifact.metadata["num_missing"] == 1 + assert {r["instance_id"]: r["uploaded"] for r in step_index(artifact)} == {"linear": True, "gone": False} + + +@pytest.mark.asyncio +async def test_no_upload_without_wandb_logging(skycap, trials) -> None: + records = artifacts.RecordLog() + await generator(skycap, records=records).generate(batch("linear"), disable_tqdm=True) + callback, trainer, fake = uploads(records, backend="tensorboard") + callback.on_step_end(trainer, event(), None) + callback.on_train_end(trainer, event(), None) + assert fake.logged == [] + + +@pytest.mark.asyncio +async def test_a_failed_upload_is_retried_three_times(skycap, trials, monkeypatch) -> None: + monkeypatch.setattr(artifacts, "RETRY_DELAY", 0.0) + for failures, logged in ((2, 1), (3, 0)): + records = artifacts.RecordLog() + await generator(skycap, records=records).generate(batch("linear"), disable_tqdm=True) + callback, trainer, fake = uploads(records, failures=failures) + callback.on_step_end(trainer, event(), None) + callback.on_train_end(trainer, event(), None) + assert len(fake.logged) == logged and fake.failures == 0 + + def test_a_batch_with_nothing_to_train_still_carries_padded_support() -> None: outcome = TrialOutcome(trajectory_id=TrajectoryID("a", 0), stop_reason="error") diff --git a/tests/train/test_rl_callbacks.py b/tests/train/test_rl_callbacks.py index 3a077e381c..386b8bc8cb 100644 --- a/tests/train/test_rl_callbacks.py +++ b/tests/train/test_rl_callbacks.py @@ -22,6 +22,7 @@ import torch from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch +from skyrl.train.generators.base import TrajectoryID from skyrl.train.trainer import RayPPOTrainer from skyrl.train.utils.callbacks import ( CallbackInput, @@ -66,6 +67,7 @@ def _snap(self, name: str, ci: CallbackInput) -> None: "total_steps": ci.total_steps, "steps_per_epoch": ci.steps_per_epoch, "has_batch": ci.batch is not None, + "trajectory_ids": ci.trajectory_ids, "has_metrics": ci.metrics is not None, "metrics_keys": sorted((ci.metrics or {}).keys()), "has_logs": ci.logs is not None, @@ -221,7 +223,14 @@ def test_callbacks_fire_during_rl_training(monkeypatch): monkeypatch.setattr( trainer, "generate", - AsyncMock(return_value={"rollout_metrics": None, "response_ids": [[1]], "rewards": [0.0]}), + AsyncMock( + return_value={ + "rollout_metrics": None, + "response_ids": [[1]], + "rewards": [0.0], + "trajectory_ids": [TrajectoryID(instance_id="uid-0", repetition_id=0)], + } + ), ) monkeypatch.setattr(trainer, "eval", AsyncMock(return_value={"eval/score": 0.5})) @@ -299,6 +308,7 @@ def _record_load_checkpoints(): assert snap["has_batch"], "on_step_end should see the training batch" assert snap["has_metrics"], "on_step_end should see step metrics" assert "policy_loss" in snap["metrics_keys"], snap["metrics_keys"] + assert snap["trajectory_ids"] == [TrajectoryID(instance_id="uid-0", repetition_id=0)], snap # Both eval ends carry eval metrics for snap in snaps_by_event["on_eval_end"]: