diff --git a/docker/miles-ports/Dockerfile b/docker/miles-ports/Dockerfile index 295628294..940711fd7 100644 --- a/docker/miles-ports/Dockerfile +++ b/docker/miles-ports/Dockerfile @@ -19,7 +19,7 @@ FROM docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c ARG MILES_REPOSITORY=https://github.com/michaellchung/miles -ARG MILES_COMMIT=386ae3b07138ccd45ce4a07fc0383587f7968a3a +ARG MILES_COMMIT=44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9 ARG SGLANG_REPOSITORY=https://github.com/michaellchung/sglang ARG SGLANG_COMMIT=2fa880182eefdb31e64d8ea70317be27f9f32121 # setuptools-scm version of SGLANG_COMMIT (describe v0.5.20-69-g4e4148f1b); @@ -47,5 +47,17 @@ RUN cd /sgl-workspace/sglang \ && SETUPTOOLS_SCM_PRETEND_VERSION="$SGLANG_VERSION" \ SGLANG_BUILD_RUST_EXTS=none pip install -e "python[all]" --no-deps +# peft at the pinned MILES_PEFT_VERSION, no deps (as the build script's +# wheel unpack); island setup then reads it from the manifest, not by import. +ARG PEFT_VERSION=0.20.0 +RUN pip install --no-deps "peft==$PEFT_VERSION" + +# Lazy modelopt (rl-trainer-cold-start-cost L2.2), as the crane layer: the +# Megatron-Bridge package __init__ and the Megatron-LM patch. +COPY megatron_bridge_init.py /opt/sglang/lib/python3.12/site-packages/megatron/bridge/__init__.py +COPY megatron_lm_lazy_modelopt.patch /tmp/ +RUN cd /root/Megatron-LM && git apply /tmp/megatron_lm_lazy_modelopt.patch \ + && rm /tmp/megatron_lm_lazy_modelopt.patch + # Provenance (written by the build script: component SHAs + base digests). COPY --from=manifest image-manifest.json /opt/yeto/image-manifest.json diff --git a/docker/miles-ports/make_layer.py b/docker/miles-ports/make_layer.py index f2b428f43..31cdb726c 100755 --- a/docker/miles-ports/make_layer.py +++ b/docker/miles-ports/make_layer.py @@ -12,6 +12,15 @@ fork repository (opaque whiteout hides the base history), * rewrites SGLang's editable-install metadata (dist-info, .pth, finder, ``sglang/_version.py``) for the fork's setuptools-scm version, +* unpacks extra pure-Python wheels (``--wheel``, e.g. peft) into the base's + site-packages, as ``pip install --no-deps`` would, +* replaces the base's ``megatron/bridge/__init__.py`` (Megatron-Bridge + 8cd3466) with ``megatron_bridge_init.py``: the same file without its eager + ``import modelopt.torch`` (rl-trainer-cold-start-cost L2.2), +* rewrites the Megatron-LM files (/root/Megatron-LM, base commit a84b105) + touched by ``megatron_lm_lazy_modelopt.patch`` -- the patch applied to the + commit's files from ``--megatron-lm-src`` -- so a trainer process does not + import modelopt unless a model uses it (L2.2), * adds /opt/yeto/image-manifest.json with every component SHA. Approximate docker equivalent (unverified): docker/miles-ports/Dockerfile. @@ -30,6 +39,7 @@ import subprocess import tarfile import tempfile +import zipfile from pathlib import Path SITE = "opt/sglang/lib/python3.12/site-packages" @@ -38,6 +48,14 @@ # pyproject.toml, setup.cfg, requirements*, sgl-kernel/**, rust/**, csrc, # .cu/.cpp -- changes what a real install builds and needs a docker build. EXEMPT_PREFIXES = ("test/", "tests/", "docs/") +# Megatron-Bridge as the base installs it (pip, not editable) and the file the +# layer writes over its package __init__ (see megatron_bridge_init.py). +MEGATRON_BRIDGE_COMMIT = "8cd3466d14d2337c8492827b3712482c2b3e4866" +MEGATRON_BRIDGE_INIT = Path(__file__).with_name("megatron_bridge_init.py") +# Megatron-LM as the base checks it out (/root/Megatron-LM, on sys.path) and +# the patch the layer applies to it (see the comments it adds). +MEGATRON_LM_COMMIT = "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e" +MEGATRON_LM_PATCH = Path(__file__).with_name("megatron_lm_lazy_modelopt.patch") def git(repo: Path, *argv: str) -> str: @@ -93,6 +111,45 @@ def version_file(version: str) -> bytes: ).encode() +def wheel_files(wheel: Path) -> dict[str, bytes]: + """``{site-packages relative path: bytes}`` of a pure-Python wheel. + + Only ``*-none-any`` wheels without a ``.data`` directory are accepted: + those install by unpacking alone (no scripts, headers or compiled code). + An ``INSTALLER`` file is added, as pip writes one. + """ + + if not wheel.name.endswith("-none-any.whl"): + raise SystemExit(f"{wheel.name}: only pure-Python (*-none-any) wheels are supported") + files: dict[str, bytes] = {} + with zipfile.ZipFile(wheel) as archive: + for name in archive.namelist(): + if name.endswith("/"): + continue + if ".data/" in name.split(".dist-info/")[0]: + raise SystemExit(f"{wheel.name}: has a .data directory ({name}); not supported") + files[name] = archive.read(name) + dist_info = {n.split("/")[0] for n in files if n.split("/")[0].endswith(".dist-info")} + if len(dist_info) != 1: + raise SystemExit(f"{wheel.name}: expected one .dist-info directory, found {sorted(dist_info)}") + files[f"{dist_info.pop()}/INSTALLER"] = b"yeto-make-layer\n" + return files + + +def patched_files(src: Path, commit: str, patch: Path) -> dict[str, bytes]: + """``{repo path: bytes}`` of the files ``patch`` touches, as ``commit`` plus the patch.""" + paths = [line.split("\t")[-1] for line in subprocess.run( + ["git", "apply", "--numstat", str(patch)], check=True, capture_output=True, text=True + ).stdout.splitlines()] + with tempfile.TemporaryDirectory(prefix="img-patch-") as tmp: + archive = subprocess.run(["git", "-C", str(src), "archive", commit, *paths], + check=True, capture_output=True).stdout + with tarfile.open(fileobj=io.BytesIO(archive)) as tar: + tar.extractall(tmp, filter="data") + subprocess.run(["git", "apply", str(patch.resolve())], cwd=tmp, check=True, capture_output=True) + return {p: Path(tmp, p).read_bytes() for p in paths} + + def main() -> None: p = argparse.ArgumentParser() p.add_argument("--manifest-json", required=True, help="image-manifest.json content file") @@ -108,6 +165,10 @@ def main() -> None: p.add_argument("--sglang-version", required=True) p.add_argument("--sglang-base-metadata", type=Path, required=True, help="dir with the base image's sglang site-packages metadata") + p.add_argument("--wheel", type=Path, action="append", default=[], + help="pure-Python wheel unpacked into site-packages (repeatable)") + p.add_argument("--megatron-lm-src", type=Path, required=True, + help=f"clone containing Megatron-LM {MEGATRON_LM_COMMIT[:7]} (the base's checkout)") p.add_argument("--mtime", type=int, required=True) p.add_argument("--out", type=Path, required=True) a = p.parse_args() @@ -184,6 +245,14 @@ def add_tree(root: Path, dest: str) -> None: add_file(f"{SITE}/{name}", data) add_file("sgl-workspace/sglang/python/sglang/_version.py", version_file(new)) + for wheel in a.wheel: + for name, data in sorted(wheel_files(wheel).items()): + add_file(f"{SITE}/{name}", data) + + add_file(f"{SITE}/megatron/bridge/__init__.py", MEGATRON_BRIDGE_INIT.read_bytes()) + for name, data in patched_files(a.megatron_lm_src, MEGATRON_LM_COMMIT, MEGATRON_LM_PATCH).items(): + add_file(f"root/Megatron-LM/{name}", data) + add_dir("opt/yeto") add_file("opt/yeto/image-manifest.json", Path(a.manifest_json).read_bytes()) diff --git a/docker/miles-ports/megatron_bridge_init.py b/docker/miles-ports/megatron_bridge_init.py new file mode 100644 index 000000000..e5fb1f62c --- /dev/null +++ b/docker/miles-ports/megatron_bridge_init.py @@ -0,0 +1,46 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Megatron Bridge - A component of the Megatron ecosystem.""" + +# yeto miles-ports image: the eager ``import modelopt.torch`` that Megatron-Bridge +# 8cd3466 runs here is removed (rl-trainer-cold-start-cost L2.2). It cost +# 21-30 s per trainer process (coldstart-gpu-review §6.2). Every bridge module +# that uses modelopt (training.checkpointing, distillation, conversion.modelopt_utils) +# imports it itself, so its plugins still load when a modelopt feature is used. +# Written over the base image's file by docker/miles-ports/make_layer.py. + +import megatron.bridge.diffusion.models # noqa: F401 — registers diffusion bridges +import megatron.bridge.models # noqa: F401 — triggers all bridge and HF Auto class registrations +from megatron.bridge.models.conversion.auto_bridge import AutoBridge +from megatron.bridge.package_info import ( + __contact_emails__, + __contact_names__, + __download_url__, + __homepage__, + __package_name__, + __repository_url__, + __version__, +) + + +__all__ = [ + "__version__", + "__package_name__", + "__contact_names__", + "__contact_emails__", + "__homepage__", + "__repository_url__", + "__download_url__", + "AutoBridge", +] diff --git a/docker/miles-ports/megatron_lm_lazy_modelopt.patch b/docker/miles-ports/megatron_lm_lazy_modelopt.patch new file mode 100644 index 000000000..f7e306963 --- /dev/null +++ b/docker/miles-ports/megatron_lm_lazy_modelopt.patch @@ -0,0 +1,190 @@ +diff --git a/megatron/post_training/checkpointing.py b/megatron/post_training/checkpointing.py +index 19df61b..70106de 100644 +--- a/megatron/post_training/checkpointing.py ++++ b/megatron/post_training/checkpointing.py +@@ -5,15 +5,7 @@ import os + from pathlib import Path + from typing import Optional, Tuple, Union + +-import modelopt +-import modelopt.torch.opt as mto + import torch.nn as nn +-from modelopt.torch.opt.plugins import ( +- restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy, +-) +-from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( +- _load_extra_state_from_sharded_checkpoint, +-) + + from megatron.core import dist_checkpointing + from megatron.core.dist_checkpointing.serialization import _legacy_common_state_exists +@@ -22,7 +14,11 @@ from megatron.training import get_args + from megatron.training.checkpointing import _load_base_checkpoint, load_checkpoint + from megatron.training.utils import print_rank_0 + +-from .utils import print_distributed_quant_summary ++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported ++# modelopt (and .utils, which imports it) at module top. megatron.training's ++# get_model imports this module for has_modelopt_state, a file check that needs ++# no modelopt, so every trainer paid for modelopt. The functions that use ++# modelopt import it themselves. + + logger = logging.getLogger(__name__) + +@@ -109,6 +105,11 @@ def load_modelopt_state(model: nn.Module, load_dir: Optional[str] = None) -> Non + model: the model to load the modelopt_state into + load_dir: optionally provide a different loading path + """ ++ import modelopt.torch.opt as mto ++ from modelopt.torch.opt.plugins import ( ++ restore_sharded_modelopt_state as restore_sharded_modelopt_state_legacy, ++ ) ++ + args = get_args() + load_dir = load_dir or args.load + +@@ -162,6 +163,8 @@ def load_modelopt_checkpoint( + When loading an .nemo sharded checkpoint, this is usually `model.`. Otherwise, this is + typically an empty string. + """ ++ from .utils import print_distributed_quant_summary ++ + + def _remove_prefix_state_dict_pre_hook( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs +@@ -213,6 +216,12 @@ def restore_sharded_modelopt_state(model: list[nn.Module], checkpoint_name: str + """Temporary function. Copy of modelopt.torch.opt.plugins.restore_sharded_modelopt_state. + Will be removed once modelopt.torch.opt.plugins.restore_sharded_modelopt_state is up to date. + """ ++ import modelopt ++ import modelopt.torch.opt as mto ++ from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( ++ _load_extra_state_from_sharded_checkpoint, ++ ) ++ + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + +diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py +index 41865c5..8d4cc10 100644 +--- a/megatron/training/checkpointing.py ++++ b/megatron/training/checkpointing.py +@@ -3,6 +3,7 @@ + """Input/output checkpointing.""" + + import contextlib ++import importlib.util + import inspect + import multiprocessing + import os +@@ -67,14 +68,31 @@ except ImportError: + + + # [ModelOpt]: Import +-try: +- from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state +- +- from megatron.post_training.utils import print_distributed_quant_summary +- +- has_nvidia_modelopt = True +-except Exception: +- has_nvidia_modelopt = False ++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported ++# modelopt at module top here, which every trainer process paid for (21-30 s). ++# Every use below is a no-op for a model ModelOpt has not converted ++# (save_modelopt_state / save_sharded_modelopt_state return early on ++# ``not is_converted``; the quant summary has nothing to report), and only ++# modelopt's own code converts a model. So while modelopt is not imported in ++# this process, skipping them is exact; once it is, they run as before. ++has_nvidia_modelopt = importlib.util.find_spec("modelopt") is not None ++ ++ ++def _modelopt_api(): ++ """ModelOpt checkpoint helpers, or None when no model here can carry ModelOpt state.""" ++ if not has_nvidia_modelopt or "modelopt" not in sys.modules: ++ return None ++ try: ++ from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state ++ ++ from megatron.post_training.utils import print_distributed_quant_summary ++ except Exception: ++ return None ++ return types.SimpleNamespace( ++ save_modelopt_state=save_modelopt_state, ++ save_sharded_modelopt_state=save_sharded_modelopt_state, ++ print_distributed_quant_summary=print_distributed_quant_summary, ++ ) + + + _CHECKPOINT_VERSION = None +@@ -882,8 +900,8 @@ def save_checkpoint( + verify_integrity=args.verify_integrity, + ) + # [ModelOpt]: save sharded modelopt_state +- if has_nvidia_modelopt: +- save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1)) ++ if (modelopt_api := _modelopt_api()) is not None: ++ modelopt_api.save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1)) + elif ckpt_type == CheckpointType.GLOBAL and ckpt_format in ["torch_dcp", "fsdp_dtensor"]: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + # TODO Handle non-empty directories (e.g., after a crash during saving). +@@ -940,11 +958,11 @@ def save_checkpoint( + ) + else: + # [ModelOpt]: Inject modelopt_state into state_dict +- if has_nvidia_modelopt: ++ if (modelopt_api := _modelopt_api()) is not None: + if ckpt_type == CheckpointType.LOCAL: + print_rank_0('WARNING: Local checkpointing does not support nvidia_modelopt.') + else: +- save_modelopt_state(model, state_dict) ++ modelopt_api.save_modelopt_state(model, state_dict) + + end_ckpt = time() + logger.debug( +@@ -2570,7 +2588,8 @@ def load_checkpoint( + log_printed = True + + if has_nvidia_modelopt: +- print_distributed_quant_summary(model, msg="After loading checkpoint") ++ if (modelopt_api := _modelopt_api()) is not None: ++ modelopt_api.print_distributed_quant_summary(model, msg="After loading checkpoint") + + # Load teacher model in Distillation mode. + if getattr(args, "export_kd_teacher_load", None): +diff --git a/megatron/training/training.py b/megatron/training/training.py +index d64ae84..eb51d7d 100644 +--- a/megatron/training/training.py ++++ b/megatron/training/training.py +@@ -13,6 +13,7 @@ import argparse + import copy + import dataclasses + import functools ++import importlib.util + import gc + import inspect + import logging +@@ -204,12 +205,18 @@ try: + except ImportError: + has_rl_utils = False + +-try: +- from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation ++# yeto miles-ports image (rl-trainer-cold-start-cost L2.2): Megatron-LM a84b105 imported ++# modelopt at module top here, which every trainer process paid for (21-30 s). ++# The distillation helper is now imported where a modelopt-enabled run uses it. ++has_nvidia_modelopt = importlib.util.find_spec("modelopt") is not None + +- has_nvidia_modelopt = True +-except ImportError: +- has_nvidia_modelopt = False ++ ++def get_tensor_shapes_adjust_fn_for_distillation(*args, **kwargs): ++ from modelopt.torch.distill.plugins.megatron import ( ++ get_tensor_shapes_adjust_fn_for_distillation as adjust_fn, ++ ) ++ ++ return adjust_fn(*args, **kwargs) + + try: + from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP diff --git a/docs/MILES_RL.md b/docs/MILES_RL.md index c012328e1..d66178498 100644 --- a/docs/MILES_RL.md +++ b/docs/MILES_RL.md @@ -1021,6 +1021,8 @@ the run exit 3. A synced run still fetches its checkpoint first. | 4 | a fixed-roster RL island could not be recovered | | 5 | the Modal app was not confirmed stopped after teardown | | 6 | the run stalled: no island event for `--rl-stall-timeout` seconds (default 900, 0 disables) and not every island finalized | +| 7 | Modal moved an island to a different container mid-run | +| 8 | the trainer scheduler gave up: `--rl-trainer-max-failures` (default 3) trainer incarnations failed in a row | For exit 5, every row of `modal app list` with the run's app name must be `stopped` with 0 tasks; an earlier run's row with the same name counts too. diff --git a/examples/rl_algorithms/decoupled-ppo-c2-agentic.json b/examples/rl_algorithms/decoupled-ppo-c2-agentic.json new file mode 100644 index 000000000..ed0b982b3 --- /dev/null +++ b/examples/rl_algorithms/decoupled-ppo-c2-agentic.json @@ -0,0 +1,53 @@ +{ + "advantage": { + "estimator": "grpo", + "reward_binary": false, + "reward_postprocess": null, + "rewards_normalization": true, + "std_normalization": true, + "whiten": false + }, + "correction": { + "batch_normalize": false, + "function": null, + "level": "token", + "lower": 0.0, + "method": "tis", + "mismatch_metrics": false, + "opsm_delta": null, + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null + }, + "entropy_coef": 0.0, + "execution": { + "group_required": true, + "max_policy_staleness": 4, + "needs_critic": false, + "needs_rollout_logprobs": false + }, + "kl": { + "coef": null, + "estimator": null, + "placement": "none", + "unbiased": false + }, + "loss": { + "aggregation": "default", + "custom_loss": null, + "eps_clip": null, + "eps_clip_c": null, + "eps_clip_high": null, + "reducer": null, + "variant": "policy_loss" + }, + "plugins": [], + "sampling": { + "filter": null, + "max_replacements": null, + "over_sampling_batch_size": null, + "overlong_filter": false + }, + "schema": "yeto-rl-algorithm-spec-v2" +} \ No newline at end of file diff --git a/examples/rl_algorithms/decoupled-ppo-c2.json b/examples/rl_algorithms/decoupled-ppo-c2.json new file mode 100644 index 000000000..b874e8155 --- /dev/null +++ b/examples/rl_algorithms/decoupled-ppo-c2.json @@ -0,0 +1,53 @@ +{ + "advantage": { + "estimator": "grpo", + "reward_binary": false, + "reward_postprocess": null, + "rewards_normalization": true, + "std_normalization": true, + "whiten": false + }, + "correction": { + "batch_normalize": false, + "function": null, + "level": "token", + "lower": 0.0, + "method": "tis", + "mismatch_metrics": false, + "opsm_delta": null, + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null + }, + "entropy_coef": 0.0, + "execution": { + "group_required": true, + "max_policy_staleness": 2, + "needs_critic": false, + "needs_rollout_logprobs": false + }, + "kl": { + "coef": null, + "estimator": null, + "placement": "none", + "unbiased": false + }, + "loss": { + "aggregation": "default", + "custom_loss": null, + "eps_clip": null, + "eps_clip_c": null, + "eps_clip_high": null, + "reducer": null, + "variant": "policy_loss" + }, + "plugins": [], + "sampling": { + "filter": null, + "max_replacements": null, + "over_sampling_batch_size": null, + "overlong_filter": false + }, + "schema": "yeto-rl-algorithm-spec-v2" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/build-record.json new file mode 100644 index 000000000..94c2d3c92 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:17aeb5c-2fa8801", + "digest": "sha256:b304f85736eace7b4e9647b9af11def88e496f47c77d2c931798b06d467c9395", + "layer_tar_sha256": "928913a1327a17d7c1c37a06f7004b2a9fe49882712141cf7c0fe1974cfea445", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "17aeb5cb51635b9ba6ef4ca10db95762695c7e03", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/image-manifest.json new file mode 100644 index 000000000..4ebec48a5 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/image-manifest.json @@ -0,0 +1,42 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "miles": { + "commit": "17aeb5cb51635b9ba6ef4ca10db95762695c7e03", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "17aeb5c-2fa8801" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/build-record.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/build-record.json new file mode 100644 index 000000000..4196d8329 --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/build-record.json @@ -0,0 +1,12 @@ +{ + "tag": "ghcr.io/michaellchung/yeto-miles-ports:44d41ef-2fa8801", + "digest": "sha256:5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721", + "layer_tar_sha256": "f4c859de9b308834177afdb903a22c024cd2177a747e21feb478d52b49d7c990", + "base": "docker.io/radixark/miles@sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "base_index": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "miles_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "sglang_version": "0.5.21.dev70+g2fa8801", + "peft_version": "0.20.0", + "peft_wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" +} \ No newline at end of file diff --git a/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/image-manifest.json b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/image-manifest.json new file mode 100644 index 000000000..38f0adf5e --- /dev/null +++ b/openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/image-manifest.json @@ -0,0 +1,50 @@ +{ + "base": { + "index_digest": "sha256:90940828dcd4d54fd907ff668b43537cbd94778047580e4160d6560af548b74d", + "manifest_digest": "sha256:4d69750720eb1a4b99976a3fed45ed018e1fc4fff3d240667fa17d258dfc6e0c", + "miles_commit": "9e4260de047a704208535c0e90c531929879ab40", + "platform": "linux/amd64", + "repository": "docker.io/radixark/miles", + "sglang_commit": "e5bba1f9c433da5352e521a8bb826bab960e7836" + }, + "builder": "scripts/build_miles_ports_image.sh (crane append); docker/miles-ports/Dockerfile equivalent", + "image": "ghcr.io/michaellchung/yeto-miles-ports", + "megatron_bridge": { + "commit": "8cd3466d14d2337c8492827b3712482c2b3e4866", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_bridge_init.py)" + ] + }, + "megatron_lm": { + "commit": "a84b105473eca9c9c6e49dbf1c8aa460c5cf796e", + "patches": [ + "lazy-modelopt (docker/miles-ports/megatron_lm_lazy_modelopt.patch)" + ], + "path": "/root/Megatron-LM", + "repository": "https://github.com/radixark/Megatron-LM" + }, + "miles": { + "commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", + "install": "base's editable install (pip install -e /root/miles --no-deps) kept (same path); /root is on the base's sys.path, so without PYTHONPATH import miles is a namespace package over the /root/miles checkout -- yeto runs with PYTHONPATH=$HOME/miles (regular package)", + "path": "/root/miles", + "repository": "https://github.com/michaellchung/miles", + "upstream_commit": "9e4260de047a704208535c0e90c531929879ab40" + }, + "peft": { + "install": "wheel unpacked into the base's site-packages (pip install --no-deps equivalent)", + "version": "0.20.0", + "wheel": "peft-0.20.0-py3-none-any.whl", + "wheel_sha256": "0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233" + }, + "pure_python_overlay": true, + "schema": 1, + "sglang": { + "commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "install": "pip install -e /sgl-workspace/sglang/python[all] --no-deps (base's editable install, same path)", + "path": "/sgl-workspace/sglang", + "repository": "https://github.com/michaellchung/sglang", + "upstream_commit": "571212b636baca45e10fa3b4da11a289123f3235", + "version": "0.5.21.dev70+g2fa8801" + }, + "tag": "44d41ef-2fa8801" +} \ No newline at end of file diff --git a/openspec/changes/yeto-framework-decoupling/hash-migration.md b/openspec/changes/yeto-framework-decoupling/hash-migration.md index 2ccae04bf..ba687f4b4 100644 --- a/openspec/changes/yeto-framework-decoupling/hash-migration.md +++ b/openspec/changes/yeto-framework-decoupling/hash-migration.md @@ -272,3 +272,25 @@ design D6：去耦合后 `AlgorithmSpec.sha256()` 按中立名与新源码重新 - Only change: the source hash of `yeto/rl/adapters/miles/rollout_meta_hook.py` `1188451c…` → `a4fa982b…` (the three plugin entries in `plugins`, all 8 samples). Cause: fully_async rounds now report `execution_mode` on the carry-over event (ff6105fc). `algorithm_sha256`, contract hashes, `ports_runtime_fingerprint` and the Miles argv are unchanged. - Regenerated with `python tests/decoupling_golden.py --write`. + +## S24 L6.2 correction group out of the hash (rl-offpolicy-correction-passthrough D8, branch s24-l6-core, 2026-10-14) + +- Cause: in phase 1 the correction group is recorded in the run manifest and island tapes, not hashed. `AlgorithmSpec.sha256()` hashes a fixed default correction group (`identity_dict()`); `canonical_json()` still carries the full group, now with the neutral field names (`level`, `lower`, `upper`, `veto`, `batch_normalize`, `ratio_mask`; `tis_clip`/`tis_clip_low` are accepted input aliases). A custom correction function's source hash stays in the identity as a plugin reference. +- Affected golden samples (3 of 8): `grpo_tis.json` — `algorithm_sha256` `5a8a5a9f…` → `27df1133…` (the default-spec hash), with the derived launcher/learner/execution-profile fields; `drgrpo.json` and `seq_adv_maxrl.json` — `algorithm_spec` text only (new field names), hashes unchanged. The other 5 samples and `fake_engine_tapes.json` are unchanged. +- Hash boundary: runs with a TIS (or other non-default) correction group recorded before this change carry the old hash; runs after it carry the hash of the same spec without its correction group. Compare across the boundary by the recorded `algorithm_spec`, not by hash. Phase 2 (L6b.1) puts the group back into the hash at the same version boundary as the `rl-identity-tiering` semantic hash. +- Accepted by the main agent (SESSION24 GOAL §6). Regenerated with `python tests/decoupling_golden.py --write`; the diff touches only these 3 files. + +## S24 L5.6 island sync timing (branch s24-l5-measure, 2026-10-14) + +- Only change: `fake_engine_tapes.json` `strict_progress.tape` gains one `rl_sync_timing` event (rl-trainer-start-stop L5.6: `sync/submit_ts`, `sync/global_ready_ts`, `sync/recv_seconds`, `sync/sync_wait_seconds`; all four are scrubbed to ``). The 8 config samples, algorithm hashes and contract hashes are unchanged. +- Regenerated with `python tests/decoupling_golden.py --write`; the diff is 11 added lines. + +## S24 cold-start image pin (branch s24-image, 2026-10-14) + +- `MILES_NEXT_COMMIT` 386ae3b07 → 17aeb5cb5 (michaellchung/miles `s24-miles-import`: light Ray actors stay off `megatron_utils`, rl-trainer-cold-start-cost L2.1; master port held until use, S23 defect 9). No loss path change; appended to `MILES_FORK_LINEAGE`. `MILES_NEXT_IMAGE` `@sha256:bcd5ffd3…` → `@sha256:b304f857…` (tag 17aeb5c-2fa8801; the layer also carries peft 0.20.0 and the lazy-modelopt Megatron-Bridge `__init__.py`, L2.2–L2.4; build record `openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-17aeb5c-2fa8801/`). +- Golden: in all 8 samples only `backend_identity` (engine_commit, sha256) and `ports_runtime_fingerprint` change with the Miles commit. Algorithm hashes, contract hashes, plugins, Miles argv and `fake_engine_tapes.json` are unchanged. Regenerated with `python tests/decoupling_golden.py --write`; the diff is 24 lines. + +## S24 lazy-modelopt image pin (branch s24-image, 2026-10-15) + +- `MILES_NEXT_COMMIT` 17aeb5cb5 → 44d41efd0 (michaellchung/miles `s24-miles-import2`: the `megatron_utils` package `__init__` no longer imports deep_ep and the Megatron-Bridge plugins; they moved to `megatron_utils.runtime_patches`, imported by the modules that need them. RayWorkerManager unpickles placement-group and worker-spec objects whose modules import `megatron_config`, so it loaded `megatron.bridge`. rl-trainer-cold-start-cost L2.2). No loss path change; appended to `MILES_FORK_LINEAGE`. `MILES_NEXT_IMAGE` `@sha256:b304f857…` → `@sha256:5b9dc2a9…` (tag 44d41ef-2fa8801; the layer also patches Megatron-LM a84b105 so `megatron.training` imports modelopt only where a model uses it, `docker/miles-ports/megatron_lm_lazy_modelopt.patch`; build record `openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801/`). +- Golden: in all 8 samples only `backend_identity` (engine_commit, sha256) and `ports_runtime_fingerprint` change with the Miles commit. Algorithm hashes, contract hashes, plugins, Miles argv and `fake_engine_tapes.json` are unchanged. Regenerated with `python tests/decoupling_golden.py --write`. diff --git a/pyproject.toml b/pyproject.toml index d8ee97b16..1ca14343d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,3 +71,4 @@ include = ["yeto*"] [tool.setuptools.package-data] "yeto.dashboard" = ["*.json", "static/*"] "yeto.cloud" = ["*.json"] +"yeto.rl.adapters.miles" = ["*.json"] diff --git a/scripts/build_miles_ports_image.sh b/scripts/build_miles_ports_image.sh index 447c70381..82beb498b 100755 --- a/scripts/build_miles_ports_image.sh +++ b/scripts/build_miles_ports_image.sh @@ -8,6 +8,8 @@ # # Inputs (env, defaults = this checkout's pins): # MILES_SRC / SGLANG_SRC local clones containing the commits +# MEGATRON_LM_SRC clone containing the base's Megatron-LM commit +# (default: shallow fetch into $OUT/megatron-lm) # MILES_COMMIT / SGLANG_COMMIT override the pinned fork commits # CRANE crane binary (go-containerregistry >= 0.20) # DEST_REPO ghcr.io/michaellchung/yeto-miles-ports (private) @@ -36,7 +38,8 @@ pins=$(python3 - <<'PY' import re src = open("yeto/rl/adapters/miles/pins.py").read() for k in ("MILES_NEXT_REPOSITORY", "MILES_NEXT_COMMIT", "MILES_NEXT_UPSTREAM_COMMIT", - "SGLANG_NEXT_REPOSITORY", "SGLANG_NEXT_COMMIT", "SGLANG_NEXT_UPSTREAM_COMMIT"): + "SGLANG_NEXT_REPOSITORY", "SGLANG_NEXT_COMMIT", "SGLANG_NEXT_UPSTREAM_COMMIT", + "MILES_PEFT_VERSION"): print(f"{k}={re.search(rf'^{k} = \"([^\"]+)\"', src, re.M).group(1)}") PY ) @@ -73,6 +76,29 @@ import json, sys; print(json.load(sys.stdin)["layers"][104]["digest"])') mv "$OUT/meta-x/opt/sglang/lib/python3.12/site-packages" "$META" fi +# Megatron-LM files patched in the layer (docker/miles-ports/ +# megatron_lm_lazy_modelopt.patch) start from the base's checkout commit. +MEGATRON_LM_REPOSITORY=https://github.com/radixark/Megatron-LM +MEGATRON_LM_COMMIT=a84b105473eca9c9c6e49dbf1c8aa460c5cf796e # /root/Megatron-LM +if [ -z "${MEGATRON_LM_SRC:-}" ]; then + MEGATRON_LM_SRC="$OUT/megatron-lm" + if [ ! -d "$MEGATRON_LM_SRC" ]; then + git init -q "$MEGATRON_LM_SRC" + git -C "$MEGATRON_LM_SRC" fetch -q --depth 1 "$MEGATRON_LM_REPOSITORY" "$MEGATRON_LM_COMMIT" + fi +fi +git -C "$MEGATRON_LM_SRC" cat-file -e "$MEGATRON_LM_COMMIT^{commit}" + +# peft (pinned MILES_PEFT_VERSION) goes into the image, so island setup +# neither installs nor imports it (rl-trainer-cold-start-cost L2.3/L2.4): +# the pure-Python wheel is unpacked into the base's site-packages, no deps +# (the island setup used to pip-install it into the same environment). +PIP=${PIP:-python3 -m pip} +mkdir -p "$OUT/wheels" +$PIP download -q --no-deps --only-binary=:all: --dest "$OUT/wheels" "peft==$MILES_PEFT_VERSION" +PEFT_WHEEL=$(ls "$OUT/wheels"/peft-"$MILES_PEFT_VERSION"-*.whl) +PEFT_WHEEL_SHA=$(sha256sum "$PEFT_WHEEL" | cut -d' ' -f1) + python3 - "$OUT/image-manifest.json" <, absent: Vec, + /// Deltas rejected since the previous outer step: (island, reason). + rejected: Vec<(u32, &'static str)>, }, } @@ -782,20 +784,22 @@ impl TapeEvent { f64_json(*cap_arrived), f64_json(*cap_total) ), - Self::OuterStep { base_version, arrived, cap_arrived, cap_total, timed_out, raw_weights, weights, absent } => { + Self::OuterStep { base_version, arrived, cap_arrived, cap_total, timed_out, raw_weights, weights, absent, rejected } => { let fmt = |v: &Vec<(u32, u64, f64)>| -> String { v.iter().map(|(i, b, w)| format!("[{i},{b},{}]", f64_json(*w))).collect::>().join(",") }; let w = fmt(raw_weights); let nw = fmt(weights); let a: Vec = absent.iter().map(u32::to_string).collect(); + let rj: Vec = rejected.iter().map(|(i, r)| format!("[{i},\"{r}\"]")).collect(); format!( - "\"base_version\":{base_version},\"arrived\":{arrived},\"cap_arrived\":{},\"cap_total\":{},\"timed_out\":{timed_out},\"raw_weights\":[{}],\"weights\":[{}],\"absent\":[{}]", + "\"base_version\":{base_version},\"arrived\":{arrived},\"cap_arrived\":{},\"cap_total\":{},\"timed_out\":{timed_out},\"raw_weights\":[{}],\"weights\":[{}],\"absent\":[{}],\"rejected\":[{}]", f64_json(*cap_arrived), f64_json(*cap_total), w, nw, - a.join(",") + a.join(","), + rj.join(",") ) } }; @@ -816,6 +820,55 @@ impl TapeEvent { } } +/// Wall-clock fields added to each tape line when it is written (L5.7). +/// +/// The coordinator stays clock-free (its tape is compared with the Python +/// ledger); the server stamps each event as it flushes it, under the same +/// lock and right after the event was created. The round starts at the +/// previous `outer_step` line (or when the clock was made). +#[derive(Debug, Clone, PartialEq)] +pub struct TapeClock { + round_start_unix: f64, + last_arrival: Option, +} + +impl TapeClock { + pub fn new(now_unix: f64) -> Self { + Self { round_start_unix: now_unix, last_arrival: None } + } + + /// Extra JSON fields for `ev` written at `now_unix` (leading comma + /// included): `time_unix` on every event; `arrival_offset_s` on + /// accepted and carried deltas; `step_seconds` and + /// `last_arrival_island` on `outer_step`, which also starts a new round. + pub fn stamp(&mut self, ev: &TapeEvent, now_unix: f64) -> String { + let mut out = format!(",\"time_unix\":{}", f64_json(now_unix)); + let offset = now_unix - self.round_start_unix; + match ev { + TapeEvent::DeltaAccepted { island_id } | TapeEvent::DeltaCarriedOver { island_id, .. } => { + self.last_arrival = Some(*island_id); + out += &format!(",\"arrival_offset_s\":{}", f64_json(offset)); + } + TapeEvent::OuterStep { .. } => { + let last = self.last_arrival.take().map(|i| i.to_string()).unwrap_or_else(|| "null".into()); + out += &format!(",\"step_seconds\":{},\"last_arrival_island\":{last}", f64_json(offset)); + self.round_start_unix = now_unix; + } + _ => {} + } + out + } +} + +impl TapeEvent { + /// `to_json` plus the clock fields from `TapeClock::stamp`. + pub fn to_json_stamped(&self, syncer_epoch: u64, clock: &mut TapeClock, now_unix: f64) -> String { + let mut line = self.to_json(syncer_epoch); + line.pop(); // closing brace + line + &clock.stamp(self, now_unix) + "}" + } +} + #[derive(Debug, Clone, PartialEq)] pub struct StepPlan { pub base_version: u64, @@ -832,6 +885,9 @@ pub struct ElasticCoordinator { members: BTreeMap, pending: BTreeMap, carried: BTreeMap, + /// Rejections since the last outer step (L5.5 merge metrics; not + /// checkpointed, a restart starts the count again). + rejected: Vec<(u32, &'static str)>, pub tape: Vec, /// Ledger of uncommitted deltas dropped by LEAVE / lease expiry /// (persisted by the D-S6 checkpoint). @@ -1019,6 +1075,7 @@ impl ElasticCoordinator { members: BTreeMap::new(), pending: BTreeMap::new(), carried: BTreeMap::new(), + rejected: Vec::new(), tape: Vec::new(), dropped: Vec::new(), published: BTreeMap::new(), @@ -1191,6 +1248,9 @@ impl ElasticCoordinator { } } }; + if let TapeEvent::DeltaRejected { island_id, reason } = ev { + self.rejected.push((island_id, reason)); + } self.tape.push(ev); self.tape.last().unwrap() } @@ -1252,6 +1312,7 @@ impl ElasticCoordinator { raw_weights: raw, weights: weights.clone(), absent, + rejected: std::mem::take(&mut self.rejected), }); Some(StepPlan { base_version, weights }) } @@ -1751,4 +1812,67 @@ mod tests { let plan = c.try_advance(false).unwrap(); assert_eq!(plan.weights, vec![(1, 2, 0.5), (2, 2, 0.5)]); } + + /// L5.5: weights, carried discount and rejected reasons per outer step, + /// against hand-computed values (w = tokens^2 / steps). + #[test] + fn outer_step_records_merge_shares_carry_and_rejections() { + let mut c = ElasticCoordinator::new(params(), 30.0, 0).unwrap(); + for i in 1..=3 { + c.join(i, 1.0, 0.0).unwrap(); + } + c.submit(Delta { island_id: 1, base_version: 0, c_tokens: 10, c_steps: 1 }); + c.submit(Delta { island_id: 2, base_version: 0, c_tokens: 20, c_steps: 2 }); + let plan = c.try_advance(true).unwrap(); // v0 -> v1, island 3 absent + assert_eq!(plan.weights, vec![(1, 0, 100.0 / 300.0), (2, 0, 200.0 / 300.0)]); + match c.tape.last().unwrap() { + TapeEvent::OuterStep { rejected, .. } => assert!(rejected.is_empty()), + e => panic!("{e:?}"), + } + + c.submit(Delta { island_id: 3, base_version: 0, c_tokens: 30, c_steps: 1 }); // lag 1 + c.submit(Delta { island_id: 1, base_version: 1, c_tokens: 10, c_steps: 1 }); + c.submit(Delta { island_id: 9, base_version: 1, c_tokens: 10, c_steps: 1 }); + c.submit(Delta { island_id: 2, base_version: 5, c_tokens: 10, c_steps: 1 }); + let plan = c.try_advance(true).unwrap(); + // raw: island 1 = 100; island 3 carried = 900 * 0.5 = 450. + let mut got = plan.weights.clone(); + got.sort_by_key(|x| x.0); + assert_eq!(got, vec![(1, 1, 100.0 / 550.0), (3, 0, 450.0 / 550.0)]); + let ev = c.tape.last().unwrap(); + let TapeEvent::OuterStep { raw_weights, rejected, .. } = ev else { panic!("{ev:?}") }; + assert!(raw_weights.contains(&(3, 0, 450.0))); + assert_eq!(rejected, &vec![(9, "not_member"), (2, "unknown_version")]); + assert!(ev.to_json(0).ends_with(",\"rejected\":[[9,\"not_member\"],[2,\"unknown_version\"]]}")); + } + + /// L5.7: time_unix on every line; arrival offset on deltas; step + /// duration and last arrival on outer_step, which restarts the round. + #[test] + fn tape_clock_stamps_arrival_offsets_and_step_duration() { + let mut clock = TapeClock::new(100.0); + let join = TapeEvent::PoolJoin { island_id: 1, catch_up: false, base_version: 0, membership_epoch: 1 }; + assert_eq!(clock.stamp(&join, 100.5), ",\"time_unix\":100.5"); + let acc = TapeEvent::DeltaAccepted { island_id: 1 }; + assert_eq!(clock.stamp(&acc, 101.5), ",\"time_unix\":101.5,\"arrival_offset_s\":1.5"); + let carry = TapeEvent::DeltaCarriedOver { island_id: 3, base: 0, lag: 1, discount: 0.5 }; + assert_eq!(clock.stamp(&carry, 103.0), ",\"time_unix\":103.0,\"arrival_offset_s\":3.0"); + let step = TapeEvent::OuterStep { + base_version: 1, + arrived: 1, + cap_arrived: 1.0, + cap_total: 2.0, + timed_out: false, + raw_weights: vec![], + weights: vec![], + absent: vec![], + rejected: vec![], + }; + let line = step.to_json_stamped(7, &mut clock, 104.0); + assert!(line.starts_with("{\"kind\":\"outer_step\",\"syncer_epoch\":7,")); + assert!(line.ends_with(",\"time_unix\":104.0,\"step_seconds\":4.0,\"last_arrival_island\":3}")); + assert_eq!(clock.stamp(&acc, 104.25), ",\"time_unix\":104.25,\"arrival_offset_s\":0.25"); + clock.stamp(&step, 105.0); + assert!(clock.stamp(&step, 106.0).ends_with("\"step_seconds\":1.0,\"last_arrival_island\":null")); + } } diff --git a/syncer/src/elastic_server.rs b/syncer/src/elastic_server.rs index 64232fccd..30fc6aadd 100644 --- a/syncer/src/elastic_server.rs +++ b/syncer/src/elastic_server.rs @@ -62,6 +62,11 @@ fn group_label(g: &str) -> &str { if g.is_empty() { "未声明" } else { g } } +/// Seconds since the Unix epoch (tape `time_unix`). +fn unix_now() -> f64 { + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(f64::NAN, |d| d.as_secs_f64()) +} + fn hex32(h: &[u8; 32]) -> String { if h.iter().all(|b| *b == 0) { return "".into(); @@ -121,6 +126,7 @@ struct Shared { flushed: usize, round_started: Instant, tape: Option, + tape_clock: crate::elastic::TapeClock, key: Vec, outer_lr: f32, outer_momentum: f32, @@ -269,7 +275,7 @@ impl Shared { let new = &self.coord.tape[self.flushed..]; if let Some(f) = self.tape.as_mut() { for ev in new { - writeln!(f, "{}", ev.to_json(self.coord.syncer_epoch))?; + writeln!(f, "{}", ev.to_json_stamped(self.coord.syncer_epoch, &mut self.tape_clock, unix_now()))?; } f.flush()?; } @@ -306,8 +312,8 @@ impl Shared { let island_json = island.map(|i| i.to_string()).unwrap_or_else(|| "null".into()); let secs = if seconds.is_finite() { format!("{seconds:?}") } else { "null".into() }; let line = format!( - "{{\"kind\":\"transfer\",\"syncer_epoch\":{},\"direction\":\"{direction}\",\"msg\":\"{name}\",\"island_id\":{island_json},\"bytes\":{bytes},\"seconds\":{secs},\"outer_version\":{}}}", - self.coord.syncer_epoch, self.coord.outer_version + "{{\"kind\":\"transfer\",\"syncer_epoch\":{},\"direction\":\"{direction}\",\"msg\":\"{name}\",\"island_id\":{island_json},\"bytes\":{bytes},\"seconds\":{secs},\"outer_version\":{},\"time_unix\":{:?}}}", + self.coord.syncer_epoch, self.coord.outer_version, unix_now() ); if let Some(f) = self.tape.as_mut() { writeln!(f, "{line}")?; @@ -600,6 +606,7 @@ pub async fn run(listener: TcpListener, cfg: ElasticServerConfig) -> Result<(u64 compat_group: None, flushed: 0, round_started: Instant::now(), + tape_clock: crate::elastic::TapeClock::new(unix_now()), tape, key: cfg.key.clone(), outer_lr: cfg.outer_lr, diff --git a/tests/buffer_rollout_procs.py b/tests/buffer_rollout_procs.py new file mode 100644 index 000000000..391672118 --- /dev/null +++ b/tests/buffer_rollout_procs.py @@ -0,0 +1,59 @@ +"""Fake sampler / trainer processes for tests/test_rl_buffer_rollout.py (no Ray, no Miles). + +``python -m tests.buffer_rollout_procs sampler STORE ID GROUPS VERSION`` + writes GROUPS groups (2 samples each) generated at engine VERSION; +``python -m tests.buffer_rollout_procs trainer STORE CUT LOG ROUNDS GROUPS [CRASH_ROUND]`` + resumes at the cut's watermark, trains ROUNDS rounds of GROUPS groups (the trained + sample ids go to LOG), cuts after each round; exits hard after taking CRASH_ROUND. +""" + +from __future__ import annotations + +import json +import os +import sys + +from yeto.rl.adapters.miles.buffer_rollout import BUFFER_ROOT_ENV, BatchReader, GroupWriter, open_buffer + + +def sample(index: int, group: int, version: int) -> dict: + return {"index": index, "group_index": group, "response_length": 3, "tokens": [1, 2, 3, 4, 5], + "rollout_log_probs": [-0.1, -0.2, -0.3], "status": "completed", + "weight_versions": [[{"version": str(version), "abs_start": 2, "abs_end": 5}]]} + + +def sampler(store: str, sampler_id: str, groups: int, version: int) -> None: + buffer = open_buffer({BUFFER_ROOT_ENV: store}) + writer = GroupWriter(buffer, sampler_id) + start = len(buffer) + for g in range(start, start + groups): + writer.write([sample(2 * g, g, version), sample(2 * g + 1, g, version)], group_key=f"g{g}") + + +def trainer(store: str, cut_path: str, log_path: str, rounds: int, groups: int, crash: int | None) -> None: + cut = json.loads(open(cut_path).read()) if os.path.exists(cut_path) else {"next": 0, "watermark": 0} + reader = BatchReader(open_buffer({BUFFER_ROOT_ENV: store}, batch_size=groups), max_policy_age=2, + wait_s=0.0, poll_s=0.0) + reader.reset(cut["watermark"]) + rid = cut["next"] + for rid in range(cut["next"], cut["next"] + rounds): + taken, report = reader.take(rid, groups, current_version=rid + 1) + if rid == crash: + os._exit(3) # crash after consuming, before the train step and the cut + with open(log_path, "a") as fh: # the train step + fh.write(json.dumps({"rollout_id": rid, "groups": [g[0]["group_index"] for _, g in taken], + "stale": report.stale_groups}) + "\n") + cut = {"next": rid + 1, "watermark": reader.watermark_before(rid + 1)} + tmp = cut_path + ".tmp" + with open(tmp, "w") as fh: + json.dump(cut, fh) + os.replace(tmp, cut_path) + + +if __name__ == "__main__": + mode, *rest = sys.argv[1:] + if mode == "sampler": + sampler(rest[0], rest[1], int(rest[2]), int(rest[3])) + else: + trainer(rest[0], rest[1], rest[2], int(rest[3]), int(rest[4]), + int(rest[5]) if len(rest) > 5 else None) diff --git a/tests/golden/decoupling/codex_harness.json b/tests/golden/decoupling/codex_harness.json index 110192f55..1da3b18e8 100644 --- a/tests/golden/decoupling/codex_harness.json +++ b/tests/golden/decoupling/codex_harness.json @@ -14,10 +14,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:67aee42a1a530e9b32695039266e1cef5393816354f3dd9363a6467a31823d43", @@ -269,7 +269,7 @@ "source_sha256": "6055fcc8d5922d5c8c4caac07db8baef4a473188adb571974835b491168ea1d0" } ], - "ports_runtime_fingerprint": "sha256:87d0d31cc84996342a8b73b3c596e987d0bc30f4e841a36f8985ee67a10dec42", + "ports_runtime_fingerprint": "sha256:2909542e2c68e09214ab6e13f916b8740dd55dcfea14c34d41681aa609b3101b", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/decoupled.json b/tests/golden/decoupling/decoupled.json index a8e202ff2..5df13d16c 100644 --- a/tests/golden/decoupling/decoupled.json +++ b/tests/golden/decoupling/decoupled.json @@ -14,10 +14,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:05aafa6daecde7385ca10a91faed2b9a0e6e66e1d9d29ba30347ae9feb85e001", @@ -230,7 +230,7 @@ "source_sha256": "a4fa982bd6d6dfebbe69657b40e260091e738babc3f711ee12688f49c20731b9" } ], - "ports_runtime_fingerprint": "sha256:a264d5a6993a1f15754fc174e4a9f21d55c9b133db556787a8939f012c2d0102", + "ports_runtime_fingerprint": "sha256:10e3b45be7113b25e5bb4c9ed4b73de4f1d2afb735ceb9c7f9dae7f00d8cfefc", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/drgrpo.json b/tests/golden/decoupling/drgrpo.json index cd44dbd79..e41802e86 100644 --- a/tests/golden/decoupling/drgrpo.json +++ b/tests/golden/decoupling/drgrpo.json @@ -12,13 +12,17 @@ "whiten": false }, "correction": { + "batch_normalize": false, "function": null, + "level": "token", + "lower": null, "method": "none", "mismatch_metrics": false, "opsm_delta": null, - "tis_clip": null, - "tis_clip_low": null, - "use_rollout_logprobs": false + "ratio_mask": "none", + "upper": null, + "use_rollout_logprobs": false, + "veto": null }, "entropy_coef": 0.0, "execution": { @@ -59,10 +63,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:bf82d9e79e11047714673c97a0c60a9c80ac6673763a0052ed90d86a3cd2cf93", @@ -282,7 +286,7 @@ "spec_sha256": "253856acaefbea8de936b03ea89bf5c50039719732f21709aa66358f9f78e4ed" } ], - "ports_runtime_fingerprint": "sha256:f09780299090f49f7a1beaed1bc05936ec2a220d83ba7f82a5f636085086d90e", + "ports_runtime_fingerprint": "sha256:962e9c27f65c4b97d6466363157dc304aaaf602b4079c76f5a907b335e5d175c", "runtime_attrs": { "yeto_algo_plugins": { "config": { diff --git a/tests/golden/decoupling/elastic.json b/tests/golden/decoupling/elastic.json index f4e56fd92..91b319876 100644 --- a/tests/golden/decoupling/elastic.json +++ b/tests/golden/decoupling/elastic.json @@ -14,10 +14,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:67aee42a1a530e9b32695039266e1cef5393816354f3dd9363a6467a31823d43", @@ -236,7 +236,7 @@ "source_sha256": "a4fa982bd6d6dfebbe69657b40e260091e738babc3f711ee12688f49c20731b9" } ], - "ports_runtime_fingerprint": "sha256:f1f204812fd1a14d134a0bb53bf3c28370220fa86bb0d0e5ab3c68d9e05ae4f0", + "ports_runtime_fingerprint": "sha256:db29c7712ed2bf1a578c4c1ffd5308e85991b89a6976c6588b9c79d642f8836f", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/fake_engine_tapes.json b/tests/golden/decoupling/fake_engine_tapes.json index 88a6a51b8..4c5791928 100644 --- a/tests/golden/decoupling/fake_engine_tapes.json +++ b/tests/golden/decoupling/fake_engine_tapes.json @@ -2160,6 +2160,17 @@ "policy_version": 1, "time_unix": "" }, + { + "event": "rl_sync_timing", + "island_id": 0, + "outer_sync": "strict", + "rollout_id": 0, + "sync/global_ready_ts": "", + "sync/recv_seconds": "", + "sync/submit_ts": "", + "sync/sync_wait_seconds": "", + "time_unix": "" + }, { "event": "rl_driver_phase", "island_id": 0, diff --git a/tests/golden/decoupling/fn_2x8.json b/tests/golden/decoupling/fn_2x8.json index 206328029..70d90bc30 100644 --- a/tests/golden/decoupling/fn_2x8.json +++ b/tests/golden/decoupling/fn_2x8.json @@ -14,10 +14,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:f88f978bfa3e826b91a31484bc1f8ad5c3f82ff375a77e5a6890e7a4a1114ac1", @@ -392,7 +392,7 @@ "source_sha256": "e06179ddc416d31e644823b0ca713f8cb408ac294382b34b9b9addb3fe983051" } ], - "ports_runtime_fingerprint": "sha256:f788bec4c5ce6dea6425bb0a298a3d0b688033d59a6c0c32ac91be6d034b4f13", + "ports_runtime_fingerprint": "sha256:57f0d3444b149bdd4cc68297cd82906e5e263b1b4f3ed354e1db2e0668149324", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/grpo_default.json b/tests/golden/decoupling/grpo_default.json index 0341b5323..948774bfb 100644 --- a/tests/golden/decoupling/grpo_default.json +++ b/tests/golden/decoupling/grpo_default.json @@ -14,10 +14,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:67aee42a1a530e9b32695039266e1cef5393816354f3dd9363a6467a31823d43", @@ -226,7 +226,7 @@ "source_sha256": "a4fa982bd6d6dfebbe69657b40e260091e738babc3f711ee12688f49c20731b9" } ], - "ports_runtime_fingerprint": "sha256:1ccec75e5149d087dddec88d0130529ec919944b0c3121b5049e062f3fa41e59", + "ports_runtime_fingerprint": "sha256:d0fbe08d5197c39dd2faf2852dd863cf427697a3cb4362df7e1dfab5751cceb4", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/grpo_tis.json b/tests/golden/decoupling/grpo_tis.json index 81c20fa33..9df475cc4 100644 --- a/tests/golden/decoupling/grpo_tis.json +++ b/tests/golden/decoupling/grpo_tis.json @@ -1,7 +1,7 @@ { - "algorithm_sha256": "5a8a5a9ff5abb317fc8159ab56c07293dc3bc99e084542880b6d32cc5a02b441", - "algorithm_sha256_launcher_expected": "5a8a5a9ff5abb317fc8159ab56c07293dc3bc99e084542880b6d32cc5a02b441", - "algorithm_sha256_learner_ports_spec": "5a8a5a9ff5abb317fc8159ab56c07293dc3bc99e084542880b6d32cc5a02b441", + "algorithm_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", + "algorithm_sha256_launcher_expected": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", + "algorithm_sha256_learner_ports_spec": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", "algorithm_spec": { "advantage": { "estimator": "grpo", @@ -12,13 +12,17 @@ "whiten": false }, "correction": { + "batch_normalize": false, "function": null, + "level": "token", + "lower": 0.0, "method": "tis", "mismatch_metrics": false, "opsm_delta": null, - "tis_clip": 2.0, - "tis_clip_low": 0.0, - "use_rollout_logprobs": false + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null }, "entropy_coef": 0.0, "execution": { @@ -55,16 +59,16 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { - "contract_hash": "sha256:ff3574071f62f2b71fb1eb71f5e32b71e272d1345b5cb280100429ab8a8c10f8", + "contract_hash": "sha256:67aee42a1a530e9b32695039266e1cef5393816354f3dd9363a6467a31823d43", "profile": { "algorithm_contract": "on-policy", - "algorithm_spec_sha256": "5a8a5a9ff5abb317fc8159ab56c07293dc3bc99e084542880b6d32cc5a02b441", + "algorithm_spec_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", "allowed_overlap": [], "execution_mode": "partitioned-serial", "groups_per_batch": 4, @@ -80,7 +84,7 @@ "schema": "yeto-rl-execution-profile-v1" } }, - "launch_algorithm_sha256": "5a8a5a9ff5abb317fc8159ab56c07293dc3bc99e084542880b6d32cc5a02b441", + "launch_algorithm_sha256": "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea", "launch_cli": [ "--rl-placement", "fixed-partition", @@ -274,7 +278,7 @@ "source_sha256": "a4fa982bd6d6dfebbe69657b40e260091e738babc3f711ee12688f49c20731b9" } ], - "ports_runtime_fingerprint": "sha256:a3b38f3bc0f2059195089d52dbd12184a923486eb0562243700c7fff4d927332", + "ports_runtime_fingerprint": "sha256:6562e7cb113358b8b9edf2caf454df065939c365c0617c3354d44b1f0fecfcb1", "runtime_attrs": { "yeto_rl_dynamic_sampling_max_replacements": null }, diff --git a/tests/golden/decoupling/seq_adv_maxrl.json b/tests/golden/decoupling/seq_adv_maxrl.json index 3e6435658..af43aa566 100644 --- a/tests/golden/decoupling/seq_adv_maxrl.json +++ b/tests/golden/decoupling/seq_adv_maxrl.json @@ -16,13 +16,17 @@ "whiten": false }, "correction": { + "batch_normalize": false, "function": null, + "level": "token", + "lower": null, "method": "none", "mismatch_metrics": false, "opsm_delta": null, - "tis_clip": null, - "tis_clip_low": null, - "use_rollout_logprobs": false + "ratio_mask": "none", + "upper": null, + "use_rollout_logprobs": false, + "veto": null }, "entropy_coef": 0.0, "execution": { @@ -68,10 +72,10 @@ "compat_group": "nvidia-h100", "device_family": "nvidia", "engine": "miles", - "engine_commit": "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + "engine_commit": "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", "param_map_sha256": "36c37d691b6372ee86485cf0f7d4f8d5e185a6346b875aac33380518ee8f259a", "schema": "yeto-backend-identity-v2", - "sha256": "3ee705c99863a03c2b9ceed4495efa807b74a7c29d26cbfc36973ff029c29d93" + "sha256": "aeee22abefe0a1e756fe8e23e5f2af3931a3eee946b355bc0044cace2a365ea6" }, "execution_profile": { "contract_hash": "sha256:66569e2a362185594a35dc65fd87d5e4c3a52c31b9b8505bb72cf2bf503d44c4", @@ -302,7 +306,7 @@ "spec_sha256": "c292fe3356eb3ca14fa93960c2da93303edf8d2755e62fef6b479bd833a885b8" } ], - "ports_runtime_fingerprint": "sha256:bce178d60a6fb40347df33f25e65ade9d01ba26d1223e667ff7f50f21999f140", + "ports_runtime_fingerprint": "sha256:12b422b970551519613a6459388fccb5325b8cc389d51d0e599889e655db9694", "runtime_attrs": { "yeto_algo_plugins": { "config": { diff --git a/tests/rl_cut_fakes.py b/tests/rl_cut_fakes.py index c7d953a05..8e9d097c3 100644 --- a/tests/rl_cut_fakes.py +++ b/tests/rl_cut_fakes.py @@ -113,10 +113,24 @@ def forward(self, x): return self.base(x) + self.drop(x) @ self.lora_A.T @ self.lora_B.T -def make_rank(seed=0, *, args=None, coord=None): +class FullModule(torch.nn.Module): + """A full-parameter trainer: every parameter trains, none is a LoRA adapter.""" + + def __init__(self, seed): + super().__init__() + g = torch.Generator().manual_seed(seed) + self.first = torch.nn.Parameter(torch.randn(4, 6, generator=g) * 0.1) + self.second = torch.nn.Parameter(torch.randn(6, 4, generator=g) * 0.1) + self.drop = torch.nn.Dropout(0.3) + + def forward(self, x): + return self.drop(x) @ self.first.T @ self.second.T + + +def make_rank(seed=0, *, args=None, coord=None, full=False): torch.manual_seed(1234) # frozen base identical in every build - module = LoraModule(seed) - opt = torch.optim.Adam([module.lora_A, module.lora_B], lr=0.1) + module = FullModule(seed) if full else LoraModule(seed) + opt = torch.optim.Adam([p for p in module.parameters() if p.requires_grad], lr=0.1) return SimpleNamespace( args=args or SimpleNamespace(fp16=False, bf16=False, global_batch_size=GBS), model=[module], @@ -138,7 +152,8 @@ def train_step(rank, batch): def params(rank): - return {n: p.detach().clone() for n, p in rank.model[0].named_parameters() if "lora" in n} + """The trainable parameters (LoRA adapters, or every parameter of a full-parameter rank).""" + return {n: p.detach().clone() for n, p in rank.model[0].named_parameters() if p.requires_grad} class LazyStateDistOptBackend(TorchCutBackend): diff --git a/tests/test_codepack.py b/tests/test_codepack.py new file mode 100644 index 000000000..c86a55d58 --- /dev/null +++ b/tests/test_codepack.py @@ -0,0 +1,174 @@ +"""rl-identity-tiering L9.4-L9.6: run code packs (no cloud; directory store).""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from yeto import codepack +from yeto.codepack import CodePackError, DirPackStore + + +def _git(repo: Path, *argv: str) -> str: + return subprocess.run(["git", *argv], cwd=repo, check=True, capture_output=True, + text=True).stdout + + +@pytest.fixture +def repo(tmp_path: Path) -> Path: + root = tmp_path / "repo" + (root / "yeto").mkdir(parents=True) + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@example.com") + _git(root, "config", "user.name", "t") + (root / ".gitignore").write_text("ignored.txt\n") + (root / "yeto" / "a.py").write_text("A = 1\n") + (root / "run.sh").write_text("#!/bin/sh\n") + (root / "run.sh").chmod(0o755) + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "init") + return root + + +def test_clean_pack_is_deterministic(repo: Path) -> None: + first, data1 = codepack.build_pack(repo) + (repo / "yeto" / "a.py").touch() # mtime only + second, data2 = codepack.build_pack(repo) + assert data1 == data2 and first.digest == second.digest + assert not first.dirty and first.commit == _git(repo, "rev-parse", "HEAD").strip() + + +def test_dirty_patch_includes_untracked_and_skips_ignored(repo: Path) -> None: + clean, _ = codepack.build_pack(repo) + (repo / "ignored.txt").write_text("x") + assert codepack.build_pack(repo)[0].digest == clean.digest + (repo / "yeto" / "new.py").write_text("NEW = 2\n") + index_before = (repo / ".git" / "index").read_bytes() + dirty, data = codepack.build_pack(repo) + assert dirty.dirty and dirty.digest != clean.digest + assert b"yeto/new.py" in codepack.dirty_patch(repo) + assert (repo / ".git" / "index").read_bytes() == index_before # real index untouched + tree = codepack.unpack(data, dirty.digest, repo.parent / "cache") + assert (tree / "yeto" / "new.py").read_text() == "NEW = 2\n" + assert not (tree / "ignored.txt").exists() + assert (tree / "run.sh").stat().st_mode & 0o111 + + +def test_deleted_tracked_file_is_not_packed(repo: Path) -> None: + (repo / "run.sh").unlink() + pack, data = codepack.build_pack(repo) + tree = codepack.unpack(data, pack.digest, repo.parent / "cache") + assert pack.dirty and not (tree / "run.sh").exists() + + +def test_verify_tree_rejects_changed_tree(repo: Path) -> None: + pack, data = codepack.build_pack(repo) + tree = codepack.unpack(data, pack.digest, repo.parent / "cache") + assert codepack.verify_tree(tree, pack.digest) == pack.digest + assert codepack.main(["verify", "--dir", str(tree), "--digest", pack.digest]) == 0 + (tree / "yeto" / "a.py").write_text("A = 2\n") + with pytest.raises(CodePackError, match="digest mismatch"): + codepack.verify_tree(tree, pack.digest) + assert codepack.main(["verify", "--dir", str(tree), "--digest", pack.digest]) == 1 + (tree / "yeto" / "a.py").unlink() + with pytest.raises(CodePackError, match="missing"): + codepack.verify_tree(tree, pack.digest) + + +def test_unpack_rejects_wrong_bytes(repo: Path) -> None: + pack, data = codepack.build_pack(repo) + with pytest.raises(CodePackError, match="does not have digest"): + codepack.unpack(data + b"\0", pack.digest, repo.parent / "cache") + + +def test_resume_takes_the_manifest_pack_after_local_change(repo: Path, tmp_path: Path) -> None: + store = DirPackStore(tmp_path / "store", uri_base="file-store://run") + first = codepack.resolve_run_pack(previous=None, repo=repo, store=store, + cache_root=tmp_path / "c1", semantic_sha256="s", + log=lambda _m: None) + assert not first.reused + assert first.pack.pack_uri == f"file-store://run/code_packs/{first.pack.digest}.tar" + (repo / "yeto" / "a.py").write_text("A = 99\n") # local tree drifts after start + manifest = {"code_pack": first.pack.record(), "semantic_sha256": "s"} + again = codepack.resolve_run_pack(previous=manifest, repo=repo, store=store, + cache_root=tmp_path / "c2", semantic_sha256="s", + log=lambda _m: None) + assert again.reused and again.pack == first.pack + assert (again.tree / "yeto" / "a.py").read_text() == "A = 1\n" + + +def test_upgrade_needs_equal_semantic_hash(repo: Path, tmp_path: Path) -> None: + store = DirPackStore(tmp_path / "store") + first = codepack.resolve_run_pack(previous=None, repo=repo, store=store, + cache_root=tmp_path / "c", semantic_sha256="s1", + log=lambda _m: None) + (repo / "yeto" / "a.py").write_text("A = 3\n") + manifest = {"code_pack": first.pack.record(), "semantic_sha256": "s1"} + with pytest.raises(CodePackError, match="semantic hash changes"): + codepack.resolve_run_pack(previous=manifest, repo=repo, store=store, + cache_root=tmp_path / "c", semantic_sha256="s2", + upgrade_reason="fix", log=lambda _m: None) + with pytest.raises(CodePackError, match="unknown"): + codepack.resolve_run_pack(previous={"code_pack": first.pack.record()}, repo=repo, + store=store, cache_root=tmp_path / "c", + semantic_sha256="s1", upgrade_reason="fix", + log=lambda _m: None) + up = codepack.resolve_run_pack(previous=manifest, repo=repo, store=store, + cache_root=tmp_path / "c", semantic_sha256="s1", + upgrade_reason="logging fix", log=lambda _m: None) + assert not up.reused and up.pack.digest != first.pack.digest + (step,) = up.pack.upgrades + assert (step["from"], step["to"], step["reason"]) == (first.pack.digest, up.pack.digest, + "logging fix") + assert codepack.CodePack.from_record(up.pack.record()) == up.pack + with pytest.raises(CodePackError, match="reason"): + codepack.upgrade_record(first.pack, up.pack, old_semantic="s", new_semantic="s", + reason=" ") + + +def test_store_path_must_be_visible(tmp_path: Path) -> None: + with pytest.raises(CodePackError, match="not visible"): + codepack.pack_store_for(str(tmp_path / "absent")) + with pytest.raises(CodePackError, match="no code pack store"): + codepack.pack_store_for("gs://bucket/x") + assert isinstance(codepack.pack_store_for(str(tmp_path)), DirPackStore) + + +def test_modal_volume_store_creates_the_volume_on_put(monkeypatch: pytest.MonkeyPatch) -> None: + # the first start stores the pack before the island creates the checkpoint volume + import io + import sys + import types + + calls: list[tuple[str, bool]] = [] + stored: dict[str, bytes] = {} + + class _Batch: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def put_file(self, f: io.BytesIO, path: str) -> None: + stored[path] = f.read() + + class _Volume: + def batch_upload(self, force: bool = False) -> _Batch: + return _Batch() + + def read_file(self, path: str): + yield stored[path] + + def from_name(name: str, create_if_missing: bool = False) -> _Volume: + calls.append((name, create_if_missing)) + return _Volume() + + monkeypatch.setitem(sys.modules, "modal", types.SimpleNamespace(Volume=types.SimpleNamespace(from_name=from_name))) + store = codepack._ModalVolumePackStore("modal-volume://vol/run") + uri = store.put("abc", b"pack") + assert uri == "modal-volume://vol/run/code_packs/abc.tar" + assert store.get(uri) == b"pack" + assert calls == [("vol", True), ("vol", False)] diff --git a/tests/test_codex_bridge_compaction.py b/tests/test_codex_bridge_compaction.py index 970a812b0..504c61fa6 100644 --- a/tests/test_codex_bridge_compaction.py +++ b/tests/test_codex_bridge_compaction.py @@ -299,6 +299,10 @@ async def _drive(cls, monkeypatch, **kw): def test_off_stock_bridge_unchanged_and_compactionrl_without_trigger_is_byte_identical(monkeypatch): + # created_at is wall-clock seconds; a run that crosses a second boundary + # would differ, so pin the clock this module reads (monotonic stays real). + monkeypatch.setattr(harness, "time", SimpleNamespace(time=lambda: 1_700_000_000.0, + monotonic=harness.time.monotonic)) _env(monkeypatch, on=False) stock_requests, stock_raw, stock_metrics = _run(_drive(harness._ResponsesBridge, monkeypatch)) _env(monkeypatch, on=True, t_comp=10) # never triggers below C-10 diff --git a/tests/test_coldstart_nebius.py b/tests/test_coldstart_nebius.py index acc145111..c28516398 100644 --- a/tests/test_coldstart_nebius.py +++ b/tests/test_coldstart_nebius.py @@ -122,7 +122,8 @@ def test_island_task_model_store(monkeypatch): task = _island_task(args, monkeypatch) env = model_store_env("org/model", "a" * 40) assert env in task.setup - assert '[ -n "$YETO_MODEL_STORE_HIT" ] || (nohup huggingface-cli download' in task.setup + assert '[ -n "$YETO_MODEL_STORE_HIT" ] || (nohup sh -c ' in task.setup + assert 'snapshot_download' in task.setup and 'huggingface-cli' not in task.setup # both ranks export it before the learner / worker snapshot fetch assert task.run.index(env) < task.run.index("set -e") < task.run.index("snapshot_download") plain = _cli(["--gpu", "nebius:2x1xl40s@eu-north1"]) diff --git a/tests/test_rl_algorithm_flags.py b/tests/test_rl_algorithm_flags.py index 7537beab9..917864ae2 100644 --- a/tests/test_rl_algorithm_flags.py +++ b/tests/test_rl_algorithm_flags.py @@ -218,7 +218,6 @@ def test_default_and_v1_argv_unchanged(): tis_clip_low=0.5, mismatch_metrics=True)), ["--use-tis", "--tis-clip", "5.0", "--tis-clip-low", "0.5", "--custom-tis-function-path", REF.path, "--get-mismatch-metrics"]), - (dict(correction=CorrectionSpec(use_rollout_logprobs=True)), ["--use-rollout-logprobs"]), (dict(correction=CorrectionSpec(method="opsm", opsm_delta=1e-4)), ["--use-opsm", "--opsm-delta", "0.0001"]), ] diff --git a/tests/test_rl_algorithm_provenance.py b/tests/test_rl_algorithm_provenance.py index c636bf82f..85f61bbbb 100644 --- a/tests/test_rl_algorithm_provenance.py +++ b/tests/test_rl_algorithm_provenance.py @@ -287,11 +287,11 @@ def test_export_cli_reads_spec_file(tmp_path, monkeypatch): def test_selection_event_records_absorbed_flags_and_allowance(): - launch = _launch(AlgorithmSpec(), ("--eps-clip-high", "0.28", "--use-rollout-logprobs")) + launch = _launch(AlgorithmSpec(), ("--eps-clip-high", "0.28", "--normalize-advantages")) event = selection_event(launch=launch, algorithm=launch.algorithm, miles_commit="x", unverified_mechanisms=("features:clip_higher",)) assert event["rl/algorithm_absorbed_flags"] == {"--eps-clip-high": "0.28", - "--use-rollout-logprobs": True} + "--normalize-advantages": True} assert event["rl/unverified_mechanisms"] == ["features:clip_higher"] assert event["rl/algorithm_spec_sha256"] == launch.algorithm.sha256() plain = selection_event(launch=_launch(AlgorithmSpec()), algorithm=AlgorithmSpec(), @@ -889,6 +889,9 @@ def __init__(self, **kw): def run(self): return {name: "JobStatus.SUCCEEDED" for name in self.learners} + def abort(self, error): # the trainer scheduler's give-up hook + raise AssertionError(f"unexpected abort: {error!r}") + monkeypatch.setattr(launcher, "FleetController", Controller) args = _launcher_args("ports", ("--controller", "local", "--rl-image", "docker:ghcr.io/x/y@sha256:" + "a" * 64), diff --git a/tests/test_rl_algorithm_spec_v2.py b/tests/test_rl_algorithm_spec_v2.py index e4d89957f..1593bc183 100644 --- a/tests/test_rl_algorithm_spec_v2.py +++ b/tests/test_rl_algorithm_spec_v2.py @@ -69,10 +69,10 @@ (lambda: KlSpec(unbiased="yes"), "kl.unbiased"), (lambda: AdvantageSpec(estimator="dapo"), "advantage.estimator"), (lambda: AdvantageSpec(whiten=1), "advantage.whiten"), - (lambda: CorrectionSpec(method="tis"), "correction.tis_clip"), + (lambda: CorrectionSpec(method="tis"), "correction.upper"), (lambda: CorrectionSpec(method="opsm"), "correction.opsm_delta"), (lambda: CorrectionSpec(method="bogus"), "correction.method"), - (lambda: CorrectionSpec(method="tis", tis_clip=-1, tis_clip_low=0), "correction.tis_clip"), + (lambda: CorrectionSpec(method="tis", tis_clip=-1, tis_clip_low=0), "correction.upper"), (lambda: SamplingSpec(max_replacements=-1, filter=BOUNDED_NONZERO_STD_FILTER), "sampling.max_replacements"), (lambda: SamplingSpec(over_sampling_batch_size=0), "sampling.over_sampling_batch_size"), @@ -155,7 +155,6 @@ def test_r0_golden_hashes_unchanged(): dict(advantage=AdvantageSpec(std_normalization=False)), dict(advantage=AdvantageSpec(estimator="gspo")), dict(kl=KlSpec(placement="loss", coef=0.01, estimator="k1")), - dict(correction=CorrectionSpec(method="opsm", opsm_delta=1e-4)), dict(sampling=SamplingSpec(over_sampling_batch_size=8)), dict(execution=ExecutionSpec(needs_rollout_logprobs=True)), dict(entropy_coef=0.01), @@ -169,6 +168,14 @@ def test_new_field_switches_to_v2_and_changes_hash(change): assert spec.sha256() != BOUNDED.sha256() +def test_correction_group_switches_to_v2_but_keeps_hash(): + # rl-offpolicy-correction-passthrough D8 (phase 1): the group is recorded, + # not hashed. + spec = AlgorithmSpec(correction=CorrectionSpec(method="opsm", opsm_delta=1e-4)) + assert json.loads(spec.canonical_json())["schema"] == ALGORITHM_SPEC_SCHEMA_V2 + assert spec.sha256() == AlgorithmSpec().sha256() + + def test_int_float_and_order_equivalence(): a = AlgorithmSpec.from_dict({"schema": ALGORITHM_SPEC_SCHEMA_V2, "loss": {"eps_clip": 0.2, "eps_clip_high": 0.28}, @@ -219,12 +226,12 @@ def test_plugin_hash_matches(): assert FILE_REF.sha256 == hashlib.sha256(handle.read()).hexdigest() -def test_plugin_hash_mismatch_reports_both(): +def test_plugin_hash_mismatch_warns_with_both(): + # rl-identity-tiering L9.2: a stale hash is a warning, not a reject. bad = PluginRef(FILE_REF.path, "0" * 64) - with pytest.raises(PluginIdentityError, match=f"spec {'0' * 64}, runtime {FILE_REF.sha256}"): - bad.verify() - with pytest.raises(PluginIdentityError): - AlgorithmSpec(plugins=(bad,)).verify_plugins() + assert f"spec {'0' * 64}, runtime {FILE_REF.sha256}" in bad.verify() + assert AlgorithmSpec(plugins=(bad,)).verify_plugins() == (bad.verify(),) + assert FILE_REF.verify() is None def test_plugin_import_failure(): diff --git a/tests/test_rl_buffer_rollout.py b/tests/test_rl_buffer_rollout.py new file mode 100644 index 000000000..690becc35 --- /dev/null +++ b/tests/test_rl_buffer_rollout.py @@ -0,0 +1,344 @@ +"""rl-trainer-start-stop L4.5 (CPU, no Ray, no Miles): the Miles rollout through +the durable sample buffer -- payload codec, sampler writes, trainer selection with +the version gate, the watermark in the round cut, crash and restart in processes.""" + +from __future__ import annotations + +import ast +import json +import subprocess +import sys +from pathlib import Path + +import numpy +import pytest +import torch +from test_rl_multinode_round_cut import NEXT, _checkpoint, _Driver, _Ledger, _trainer +from test_rl_reconfig_recovery import _ctl + +from tests.buffer_rollout_procs import sample +from tests.rl_cut_fakes import make_rank, params, train_step +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles.entry import _wire_sample_buffer +from yeto.rl.engine.recovery import CUT_KEY +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer + +REPO = Path(__file__).resolve().parents[1] + + +def _buffer(root, batch_size=2): + return FileSampleBuffer(CheckpointStore(root), batch_size=batch_size) + + +def _group(g, version): + return [sample(2 * g, g, version), sample(2 * g + 1, g, version)] + + +# ---------------------------------------------------------------- argv mapping +def test_argv_selects_the_buffer_rollout_functions(): + assert br.rollout_argv(None, max_policy_age=0) == () + assert br.rollout_argv("sampler", max_policy_age=0) == ( + "--rollout-function-path", "yeto.rl.adapters.miles.buffer_rollout_miles.BufferWriterRolloutFn") + assert br.rollout_argv("trainer", max_policy_age=2) == ( + "--rollout-function-path", "yeto.rl.adapters.miles.buffer_rollout_miles.BufferReaderRolloutFn") + + +@pytest.mark.parametrize("kwargs, message", [ + ({"role": "trainer", "max_policy_age": 0}, "max_policy_age >= 1"), + ({"role": "sampler", "max_policy_age": 1, "fully_async": True}, "--fully-async"), + ({"role": "trainer", "max_policy_age": 1, "rollout_function_path": "x.Y"}, "replaces"), + ({"role": "both", "max_policy_age": 1}, "expected one of"), +]) +def test_argv_refuses_unsupported_combinations(kwargs, message): + with pytest.raises(ValueError, match=message): + br.rollout_argv(**kwargs) + + +def test_rollout_function_paths_name_classes_of_the_plugin_module(): + """The Miles classes exist under the dotted paths Miles imports (the module needs + Miles, so its source is parsed instead of imported).""" + tree = ast.parse((REPO / "yeto/rl/adapters/miles/buffer_rollout_miles.py").read_text()) + classes = {node.name: [b.id for b in node.bases] for node in tree.body if isinstance(node, ast.ClassDef)} + for path in br.ROLLOUT_FUNCTION_PATHS.values(): + module, name = path.rsplit(".", 1) + assert module == "yeto.rl.adapters.miles.buffer_rollout_miles" and name in classes + assert classes["BufferWriterRolloutFn"] == ["InferenceRolloutFn"] + assert classes["BufferReaderRolloutFn"] == ["BaseRolloutFn"] + + +def _import_plugin_with_stub_miles(monkeypatch, registry): + """Import buffer_rollout_miles against minimal stand-ins for the Miles modules.""" + import importlib + import types + + class _Base: + def __init__(self, input): + self.constructor_input = input + + stubs = { + "miles.rollout.base_types": dict(BaseRolloutFn=_Base, RolloutFnConstructorInput=object, + RolloutFnInput=object, RolloutFnTrainOutput=object), + "miles.rollout.inference_rollout.inference_rollout_common": dict(InferenceRolloutFn=_Base), + "miles.utils.function_registry": dict(load_function=lambda path: registry.get(path) if path else None), + "miles.utils.types": dict(Sample=object), + } + for name, attrs in stubs.items(): + monkeypatch.setitem(sys.modules, name, types.SimpleNamespace(**attrs)) + plugin = "yeto.rl.adapters.miles.buffer_rollout_miles" + monkeypatch.setitem(sys.modules, plugin, None) # teardown restores the state before the test + del sys.modules[plugin] + return importlib.import_module(plugin) + + +def test_reader_round_calls_the_metadata_hooks(monkeypatch): + """L5.4 j: a trainer that reads the buffer must still run Miles' end-of-round hooks, + or the driver finds no rollout metadata for the round.""" + calls = [] + registry = {"filter": lambda args, data: calls.append(("filter", data)), + "all": lambda args, samples, source: calls.append(("all", samples, source))} + plugin = _import_plugin_with_stub_miles(monkeypatch, registry) + reader = plugin.BufferReaderRolloutFn.__new__(plugin.BufferReaderRolloutFn) + reader.args = args = type("A", (), {"rollout_sample_filter_path": "filter", + "rollout_all_samples_process_path": "all"})() + groups = [["s0", "s1"], ["s2", "s3"]] + reader._run_round_hooks(groups) + assert calls == [("filter", groups), ("all", groups, None)] + args.rollout_sample_filter_path = args.rollout_all_samples_process_path = None + reader._run_round_hooks(groups) + assert len(calls) == 2 + + +def test_island_env(): + assert br.island_env("sampler", store_root="/s", sampler_id="s0") == { + br.BUFFER_ROOT_ENV: "/s", br.SAMPLER_ID_ENV: "s0"} + assert br.island_env("trainer", store_root="/s", max_policy_age=2, wait_s=30) == { + br.BUFFER_ROOT_ENV: "/s", br.MAX_POLICY_AGE_ENV: "2", br.WAIT_ENV: "30"} + with pytest.raises(ValueError): + br.island_env("sampler", store_root="/s") + + +# ---------------------------------------------------------------- payload and entry +def test_payload_keeps_per_token_logprobs_versions_and_arrays(): + s = sample(0, 0, 4) + s["rollout_routed_experts"] = numpy.arange(6, dtype=numpy.int32).reshape(2, 3) + s["weight_versions"] = [[{"version": "3", "abs_start": 2, "abs_end": 3}, + {"version": "4", "abs_start": 3, "abs_end": 5}]] + back = br.decode_group(br.encode_group([s]))[0] + assert back["rollout_log_probs"] == s["rollout_log_probs"] + assert back["weight_versions"] == s["weight_versions"] + assert back["rollout_routed_experts"].dtype == numpy.int32 + assert numpy.array_equal(back["rollout_routed_experts"], s["rollout_routed_experts"]) + assert br.group_entry_fields([s, sample(1, 0, 5)]) == (3, 6) # oldest token version + + +def test_entry_needs_logprobs_and_a_numeric_version(): + bad = sample(0, 0, 1) + bad["rollout_log_probs"] = None + with pytest.raises(br.SampleBufferRolloutError, match="rollout_log_probs missing"): + br.group_entry_fields([bad]) + bad = sample(0, 0, 1) + bad["weight_versions"] = [] + with pytest.raises(br.SampleBufferRolloutError, match="numeric engine weight version"): + br.group_entry_fields([bad]) + + +def test_writer_writes_payload_then_entry(tmp_path): + buffer = _buffer(tmp_path) + entry = br.GroupWriter(buffer, "s0").write(_group(0, 7), group_key="r0-g0") + assert (entry.sample_id, entry.behavior_version, entry.tokens) == ("s0-r0-g0", 7, 6) + reopened = _buffer(tmp_path) + assert len(reopened) == 1 + assert br.decode_group(reopened.read_payload(entry.payload_ref))[1]["index"] == 1 + + +# ---------------------------------------------------------------- reader +def test_reader_passes_over_groups_older_than_the_limit(tmp_path): + buffer = _buffer(tmp_path) + writer = br.GroupWriter(buffer, "s0") + for g, version in enumerate((1, 3, 2, 4)): + writer.write(_group(g, version), group_key=f"g{g}") + reader = br.BatchReader(_buffer(tmp_path), max_policy_age=2, wait_s=0.0, poll_s=0.0) + groups, report = reader.take(0, 2, current_version=5) # floor = 3 + assert [(pos, g[0]["group_index"]) for pos, g in groups] == [(1, 1), (3, 3)] + assert (report.stale_groups, report.start, report.end, report.behavior_versions) == (2, 0, 4, {3: 1, 4: 1}) + + +def test_reader_waits_for_samples_then_times_out(tmp_path): + buffer = _buffer(tmp_path) + writer = br.GroupWriter(buffer, "s0") + now = [0.0] + + def sleep(dt): # a sampler appends while the reader waits + now[0] += 1.0 + if now[0] == 2.0: + writer.write(_group(0, 1), group_key="g0") + + reader = br.BatchReader(_buffer(tmp_path), max_policy_age=1, wait_s=5.0, poll_s=1.0, + clock=lambda: now[0], sleep=sleep) + groups, report = reader.take(0, 1, current_version=1) + assert len(groups) == 1 and report.wait_s == 2.0 + with pytest.raises(br.SampleBufferRolloutError, match="0/1 groups after 5 s"): + reader.take(1, 1, current_version=1) + + +def test_reader_returns_the_rest_after_the_data_end(tmp_path): + buffer = _buffer(tmp_path) + br.GroupWriter(buffer, "s0").write(_group(0, 1), group_key="g0") + buffer.mark_data_end() + reader = br.BatchReader(_buffer(tmp_path), max_policy_age=1, wait_s=0.0, poll_s=0.0) + groups, _ = reader.take(0, 4, current_version=1) + assert len(groups) == 1 + with pytest.raises(br.SampleBufferRolloutError, match="ended with no trainable group"): + reader.take(1, 4, current_version=1) + + +def test_watermark_covers_only_rounds_below_the_cut(tmp_path): + """A round selected ahead of the cut (prefetch) is not in the watermark.""" + buffer = _buffer(tmp_path) + writer = br.GroupWriter(buffer, "s0") + for g in range(6): + writer.write(_group(g, 1), group_key=f"g{g}") + reader = br.BatchReader(_buffer(tmp_path), max_policy_age=1, wait_s=0.0, poll_s=0.0) + for rid in range(3): + reader.take(rid, 2, current_version=1) + assert reader.watermark_before(0) == 0 + assert reader.watermark_before(2) == 4 # rounds 0 and 1 trained, round 2 selected + assert reader.buffer.watermark == 4 + assert reader.reset(2) == 2 + groups, _ = reader.take(1, 2, current_version=1) + assert [pos for pos, _ in groups] == [2, 3] + + +# ---------------------------------------------------------------- round cut +def _reader_link(tmp_path): + for g in range(6): + br.GroupWriter(_buffer(tmp_path / "buf"), "s0").write(_group(g, 1), group_key=f"g{g}") + reader = br.BatchReader(_buffer(tmp_path / "buf"), max_policy_age=1, wait_s=0.0, poll_s=0.0) + br.register_reader(reader) + pool = type("Pool", (), {"_in_executor": staticmethod(lambda fn, *a: fn(None, *a))})() + return reader, pool + + +def test_round_cut_records_the_watermark_and_resume_continues_there(tmp_path): + reader, pool = _reader_link(tmp_path) + for rid in range(NEXT + 1): # rounds 0..NEXT-1 trained, round NEXT selected ahead + reader.take(rid, 2, current_version=1) + store = tmp_path / "store" + rank = make_rank(0) + g = torch.Generator().manual_seed(3) + for _ in range(NEXT): + train_step(rank, torch.randn(3, 6, generator=g)) + ctl = _ctl(tmp_path / "a/state", {"t": 1000.0}, checkpoint_store=store) + driver = _Driver(rank, published=NEXT, local_step=NEXT) + cp = _checkpoint(ctl, driver, _trainer(rank), _Ledger(NEXT)) + args = type("Args", (), {"rollout_function_path": br.ROLLOUT_FUNCTION_PATHS["trainer"]})() + _wire_sample_buffer(cp, driver=type("D", (), {"rollout": pool})(), miles_args=args) + info = cp.save(driver, rollout_id=NEXT) + assert info[CUT_KEY]["buffer_watermark"] == 2 * NEXT + ctl.close() + + rank_b = make_rank(9) + ctl_b = _ctl(tmp_path / "b/state", {"t": 2000.0}, checkpoint_store=store) + reader_b, pool_b = _reader_link(tmp_path / "b") + driver_b = _Driver(rank_b) + cp_b = _checkpoint(ctl_b, driver_b, _trainer(rank_b), _Ledger(NEXT)) + _wire_sample_buffer(cp_b, driver=type("D", (), {"rollout": pool_b})(), miles_args=args) + cp_b.resume(driver_b) + assert reader_b.buffer.watermark == reader_b.buffer.read_cursor == 2 * NEXT + resumed = [e for e, f in driver_b.events if e == "rl_resume"] + assert resumed and dict(driver_b.events)["rl_resume"]["checks"]["buffer_watermark"] == 2 * NEXT + ctl_b.close() + + +def test_wiring_needs_round_cuts_and_skips_other_rollouts(): + args = type("Args", (), {"rollout_function_path": br.ROLLOUT_FUNCTION_PATHS["trainer"]})() + with pytest.raises(RuntimeError, match="needs round cuts"): + _wire_sample_buffer(None, driver=None, miles_args=args) + _wire_sample_buffer(None, driver=None, miles_args=type("Args", (), {"rollout_function_path": None})()) + + +# ---------------------------------------------------------------- processes +def _run(*args, check=True): + return subprocess.run([sys.executable, "-m", "tests.buffer_rollout_procs", *map(str, args)], cwd=REPO, check=check, + capture_output=True, text=True, timeout=120) + + +def test_sampler_and_trainer_processes_crash_and_continue_at_the_watermark(tmp_path): + store, cut, log = tmp_path / "store", tmp_path / "cut.json", tmp_path / "trained.jsonl" + _run("sampler", store, "s0", 6, 9) + crashed = _run("trainer", store, cut, log, 3, 2, 1, check=False) + assert crashed.returncode == 3, crashed.stderr + assert json.loads(cut.read_text()) == {"next": 1, "watermark": 2} # round 1 consumed, not cut + _run("trainer", store, cut, log, 2, 2) + _run("sampler", store, "s1", 2, 9) # a restarted sampler appends after the existing index + _run("trainer", store, cut, log, 1, 2) + trained = [json.loads(line) for line in log.read_text().splitlines()] + assert [(t["rollout_id"], t["groups"]) for t in trained] == [ + (0, [0, 1]), (1, [2, 3]), (2, [4, 5]), (3, [6, 7])] # round 1 trained once, after the restart + assert json.loads(cut.read_text()) == {"next": 4, "watermark": 8} + index = [json.loads(line) for line in (store / "sample-buffer/index.jsonl").read_text().splitlines()] + assert [e["sample_id"].split("-")[0] for e in index] == ["s0"] * 6 + ["s1"] * 2 + + +# ---------------------------------------------------------------- full-parameter stop and restart +def _round_batch(rollout_id): + return torch.randn(3, 6, generator=torch.Generator().manual_seed(100 + rollout_id)) + + +def test_full_parameter_trainer_stops_at_a_scheduled_cut_and_continues_without_retraining(tmp_path): + """L5.4 k blocker (T-FCUT): a full-parameter trainer stopped by the scheduler cuts + (trigger SCHEDULED); the next incarnation restores it, continues the policy version + and the buffer position, trains no group twice and ends bitwise equal to a run + that never stopped.""" + from yeto.rl.engine.recovery import SCHEDULED + + rounds, stop_at = 4, NEXT + torch.manual_seed(99) + reference = make_rank(0, full=True) + for rid in range(rounds): + train_step(reference, _round_batch(rid)) + + reader, pool = _reader_link(tmp_path) # 6 groups of version 1, 2 per round + writer = br.GroupWriter(_buffer(tmp_path / "buf"), "s1") + for g in range(6, 2 * rounds): + writer.write(_group(g, 1), group_key=f"g{g}") + args = type("Args", (), {"rollout_function_path": br.ROLLOUT_FUNCTION_PATHS["trainer"]})() + store, trained = tmp_path / "store", [] + torch.manual_seed(99) + rank_a = make_rank(0, full=True) + for rid in range(stop_at): + groups, _ = reader.take(rid, 2, current_version=1) + trained += [pos for pos, _ in groups] + train_step(rank_a, _round_batch(rid)) + reader.take(stop_at, 2, current_version=1) # selected ahead, never trained by A + ctl_a = _ctl(tmp_path / "a/state", {"t": 1000.0}, checkpoint_store=store) + driver_a = _Driver(rank_a, published=stop_at, local_step=stop_at) + cp_a = _checkpoint(ctl_a, driver_a, _trainer(rank_a), _Ledger(stop_at)) + _wire_sample_buffer(cp_a, driver=type("D", (), {"rollout": pool})(), miles_args=args) + info = cp_a.save(driver_a, rollout_id=stop_at, final=True, trigger=SCHEDULED) + assert info[CUT_KEY]["trigger"] == SCHEDULED and info["final"] + assert info[CUT_KEY]["buffer_watermark"] == 2 * stop_at + ctl_a.close() + + torch.manual_seed(12345) # the new incarnation's process starts with other RNG state + rank_b = make_rank(9, full=True) # and other initial weights + ctl_b = _ctl(tmp_path / "b/state", {"t": 2000.0}, checkpoint_store=store) + reader_b = br.BatchReader(_buffer(tmp_path / "buf"), max_policy_age=1, wait_s=0.0, poll_s=0.0) + br.register_reader(reader_b) + driver_b = _Driver(rank_b) + cp_b = _checkpoint(ctl_b, driver_b, _trainer(rank_b), _Ledger(stop_at)) + _wire_sample_buffer(cp_b, driver=type("D", (), {"rollout": pool})(), miles_args=args) + resumed = cp_b.resume(driver_b) + assert resumed["next_rollout_id"] == stop_at and driver_b.local_step == stop_at + event = dict(driver_b.events)["rl_resume"] + assert event["policy_version"] == stop_at and event["checks"]["policy_hash"] == "same" + for rid in range(stop_at, rounds): + groups, _ = reader_b.take(rid, 2, current_version=1) + trained += [pos for pos, _ in groups] + train_step(rank_b, _round_batch(rid)) + ctl_b.close() + + assert trained == list(range(2 * rounds)) # every group once, in order: none trained twice + for name, value in params(reference).items(): + assert torch.equal(value, params(rank_b)[name]), name diff --git a/tests/test_rl_code_pack_launch.py b/tests/test_rl_code_pack_launch.py new file mode 100644 index 000000000..e25ebbffc --- /dev/null +++ b/tests/test_rl_code_pack_launch.py @@ -0,0 +1,139 @@ +"""rl-identity-tiering L9.4-L9.6: the launcher side of the run code pack (fake sky/Modal, +directory checkpoint store; no cloud).""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import types +from pathlib import Path + +import pytest + +import yeto.launcher as launcher +from yeto import codepack +from yeto.codepack import CodePackError +from yeto.launch import code_pack +from test_rl_algorithm_provenance import _fake_sky, _launcher_args # noqa: F401 (fixture) +from test_rl_launcher_island_failure import Clock, _setup + +_REAL_RUN = subprocess.run + + +def _git(repo: Path, *argv: str) -> None: + _REAL_RUN(["git", *argv], cwd=repo, check=True, capture_output=True) + + +def _repo(root: Path) -> Path: + (root / "yeto").mkdir(parents=True) + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@example.com") + _git(root, "config", "user.name", "t") + (root / "yeto" / "a.py").write_text("A = 1\n") + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "init") + return root + + +@pytest.fixture +def launch(monkeypatch, tmp_path): + """``launch(*extra)`` runs one fake no-sync Modal launch with a code pack and + returns ``(args, manifest)``.""" + + repo = _repo(tmp_path / "repo") + store = tmp_path / "store" + store.mkdir() + record = [] + _setup(monkeypatch, tmp_path, failing=set(), succeeding={0}, clock=Clock(), record=record) + # _setup stubs subprocess.run for the launcher; the pack needs real git. + monkeypatch.setattr(codepack, "subprocess", types.SimpleNamespace( + run=_REAL_RUN, CalledProcessError=subprocess.CalledProcessError)) + # a fake modal-volume:// pack store backed by a local directory + monkeypatch.setitem(codepack._PACK_STORES, "modal-volume", + lambda uri: codepack.DirPackStore(store, uri_base=uri.rstrip("/"))) + monkeypatch.setattr(launcher, "_verify_modal_app_stopped", lambda *a, **k: True) + monkeypatch.setattr(code_pack, "REPO_ROOT", repo) + monkeypatch.setattr(code_pack, "CODE_PACK_CACHE", tmp_path / "cache") + + def run(*extra): + from yeto import runs + + args = _launcher_args("ports", ("--controller", "local", "--rl-single-island-no-sync", + "--rl-checkpoint-store", "modal-volume://vol/run", + "--rl-code-pack", "git-patch", *extra), + gpu="modal:1xa100") + args.keep, args.recover_timeout = False, 0.0 + # a no-sync run refuses old tapes; a later start here is about code only + shutil.rmtree(runs.run_dir(args.cluster_prefix) / "events", ignore_errors=True) + launcher.run(args) + path = runs.run_dir(args.cluster_prefix) / "run_manifest.json" + return args, json.loads(path.read_text()) + + run.repo, run.store = repo, store + return run + + +def test_first_start_writes_pack_and_resume_reuses_it(launch): + args, manifest = launch() + pack = manifest["code_pack"] + assert manifest["code_pack_mode"] == "git-patch" == pack["mode"] + assert manifest["code_pack_digest"] == pack["digest"] == args.code_pack.digest + assert pack["pack_uri"] == f"modal-volume://vol/run/code_packs/{pack['digest']}.tar" + assert (launch.store / "code_packs" / f"{pack['digest']}.tar").is_file() + assert pack["upgrades"] == [] + assert code_pack.code_workdir(args) == args.code_pack_tree + assert code_pack.code_pack_island_flag(args) == f" --code-pack-digest {pack['digest']}" + # the local tree changes; the next start of the same run keeps the run pack + (launch.repo / "yeto" / "a.py").write_text("A = 2\n") + args2, manifest2 = launch() + assert manifest2["code_pack"] == pack + assert (Path(args2.code_pack_tree) / "yeto" / "a.py").read_text() == "A = 1\n" + + +def test_upgrade_records_old_to_new(launch): + _, first = launch() + (launch.repo / "yeto" / "a.py").write_text("A = 3\n") + _, manifest = launch("--rl-code-pack-upgrade", "hotfix") + (step,) = manifest["code_pack"]["upgrades"] + assert step["from"] == first["code_pack_digest"] != step["to"] + assert step["to"] == manifest["code_pack_digest"] and step["reason"] == "hotfix" + + +def test_upgrade_refused_when_semantic_hash_changes(launch): + from yeto import runs + + args, manifest = launch() + manifest["semantic_sha256"] = "0" * 64 # the run started with another algorithm + path = runs.run_dir(args.cluster_prefix) / "run_manifest.json" + path.write_text(json.dumps(manifest)) + (launch.repo / "yeto" / "a.py").write_text("A = 4\n") + with pytest.raises(CodePackError, match="semantic hash changes"): + launch("--rl-code-pack-upgrade", "hotfix") + + +def test_off_by_default_and_upgrade_needs_pack(): + args = types.SimpleNamespace(training_mode="rl", rl_code_pack="off", + rl_code_pack_upgrade=None) + code_pack.prepare_run_code_pack(args) + assert not hasattr(args, "code_pack") and code_pack.code_pack_island_flag(args) == "" + args.rl_code_pack_upgrade = "x" + with pytest.raises(ValueError, match="requires --rl-code-pack"): + code_pack.prepare_run_code_pack(args) + args = types.SimpleNamespace(training_mode="rl", rl_code_pack="git-patch", + rl_checkpoint_store=None) + with pytest.raises(ValueError, match="requires --rl-checkpoint-store"): + code_pack.prepare_run_code_pack(args) + + +def test_island_refuses_a_tree_that_is_not_the_pack(tmp_path): + from yeto.rl.adapters.miles.island_entry import build_parser, verify_island_code_pack + + repo = _repo(tmp_path / "repo") + pack, data = codepack.build_pack(repo) + tree = codepack.unpack(data, pack.digest, tmp_path / "cache") + verify_island_code_pack(None) # no pack: the transition source-tree check only + verify_island_code_pack(pack.digest, tree) + with pytest.raises(CodePackError, match="digest mismatch"): + verify_island_code_pack("f" * 64, tree) + assert any("--code-pack-digest" in a.option_strings for a in build_parser()._actions) diff --git a/tests/test_rl_convertible.py b/tests/test_rl_convertible.py new file mode 100644 index 000000000..61de9c188 --- /dev/null +++ b/tests/test_rl_convertible.py @@ -0,0 +1,499 @@ +"""Convertible nodes (rl-trainer-start-stop L7.1–L7.3) with fake engines.""" + +from __future__ import annotations + +import argparse +import types + +import pytest + +from yeto.launch import trainer_loop as launch_loop +from yeto.rl.adapters.miles import LoopRunner +from yeto.rl.adapters.miles.convertible import ( + MilesConversionPort, + engine_requests_inflight, + miles_conversion_port, +) +from yeto.rl.adapters.miles.publish import MilesPublisher, PublicationError, engine_checksum_digest +from yeto.rl.engine.convertible import ( + CONVERSION_DECISION_KIND, + ConversionError, + ConversionScheduler, + ConvertibleNode, + InflightTrajectory, + LastRolloutNodeError, + NodeClass, + NodeClassError, + NodePool, + NodeSpec, + Role, + apply_cut_rule, + parse_node_spec, + validate_nodes, +) +from yeto.rl.engine.trainer_lifecycle import BacklogStats, Hold, SampleIndexEntry +from yeto.rl.engine.trainer_scheduler import S1Config, S1Scheduler + + +class Clock: + def __init__(self) -> None: + self.t = 0.0 + + def __call__(self) -> float: + return self.t + + +class FakeGpu: + """One GPU that holds either the rollout engine or the trainer, never both.""" + + def __init__(self, clock: Clock, inflight=(), checksums=None) -> None: + self.clock = clock + self.holder = "rollout" + self.accepting = True + self.calls: list[str] = [] + self.inflight = list(inflight) + self.checksums = checksums or {} + + def _step(self, name: str, seconds: float = 1.0) -> None: + self.calls.append(name) + self.clock.t += seconds + + def stop_requests(self): + self._step("stop_requests") + self.accepting = False + out, self.inflight = self.inflight, [] + return out + + def release_rollout(self): + assert self.holder == "rollout" and not self.accepting + self._step("release_rollout") + self.holder = None + + def onload_trainer(self): + assert self.holder is None, "trainer onload before rollout released its memory" + self._step("onload_trainer", 2.0) + self.holder = "trainer" + + def offload_trainer(self): + assert self.holder == "trainer" + self._step("offload_trainer") + self.holder = None + + def resume_rollout(self, version): + assert self.holder is None, "rollout resume before trainer offload" + self._step("resume_rollout", 3.0) + self.holder = "rollout" + return self.checksums.get(version, f"sha-{version}") + + def accept_requests(self): + self._step("accept_requests") + self.accepting = True + + +class Buffer: + def __init__(self) -> None: + self.entries: list[SampleIndexEntry] = [] + + def append(self, entry): + self.entries.append(entry) + + +def _entry(sid, version=3): + return SampleIndexEntry(sample_id=sid, behavior_version=version, payload_ref=f"p/{sid}") + + +def _node(name="c0", clock=None, **gpu): + clock = clock or Clock() + return ConvertibleNode(NodeSpec(name, NodeClass.CONVERTIBLE), FakeGpu(clock, **gpu), clock=clock) + + +# ---------------------------------------------------------------- L7.1 classes +def test_convertible_spot_or_low_precision_is_refused_by_name(): + with pytest.raises(NodeClassError, match="node c1: a spot node cannot be convertible"): + validate_nodes([NodeSpec("c1", NodeClass.CONVERTIBLE, spot=True)], trainer_precision="bf16") + with pytest.raises(NodeClassError, match="node c2: rollout precision fp8"): + validate_nodes([NodeSpec("c2", NodeClass.CONVERTIBLE, precision="fp8")], trainer_precision="bf16") + validate_nodes([NodeSpec("r", NodeClass.ROLLOUT, spot=True, precision="fp8"), + NodeSpec("c", NodeClass.CONVERTIBLE, precision="bf16")], trainer_precision="bf16") + + +def test_parse_node_spec(): + assert parse_node_spec("island2=rollout:spot:fp8:capacity=2") == NodeSpec( + "island2", NodeClass.ROLLOUT, spot=True, precision="fp8", capacity=2.0) + with pytest.raises(NodeClassError, match="class 'gpu'"): + parse_node_spec("x=gpu") + + +def _launch_args(argv): + parser = argparse.ArgumentParser() + launch_loop.add_arguments(parser) + parser.add_argument("--on-demand", dest="spot", action="store_false", default=True) + return parser.parse_args(argv) + + +def test_launcher_refuses_convertible_on_a_spot_run(): + args = _launch_args(["--rl-node-class", "island1=convertible"]) # --spot is the default + with pytest.raises(SystemExit, match="node island1: a spot node cannot be convertible"): + launch_loop.check(args) + ok = _launch_args(["--on-demand", "--rl-node-class", "island1=convertible", + "--rl-node-class", "island2=rollout:spot:fp8"]) + launch_loop.check(ok) + assert [s.spot for s in launch_loop.node_specs(ok)] == [False, True] + + +# ---------------------------------------------------------------- L7.2 cut rule +def test_cut_rule_returns_finished_segments_and_drops_unfinished(): + buf = Buffer() + out = apply_cut_rule([ + InflightTrajectory("t0", finished=(_entry("t0.s0"), _entry("t0.s1")), unfinished=True), + InflightTrajectory("t1", finished=(_entry("t1.s0"),), unfinished=False), + InflightTrajectory("t2"), + ], buf) + assert out.returned == ("t0.s0", "t0.s1", "t1.s0") + assert out.dropped == ("t0", "t2") + assert [e.sample_id for e in buf.entries] == list(out.returned) + assert {e.behavior_version for e in buf.entries} == {3} # behavior version kept + + +# ---------------------------------------------------------------- L7.2 both directions +def test_both_directions_in_order_with_timing_and_mid_conversion_trajectories(): + inflight = [InflightTrajectory("t0", finished=(_entry("t0.s0"),))] + node = _node(inflight=inflight) + gpu, buf = node.port, Buffer() + + r1 = node.to_trainer(buf) + assert gpu.calls == ["stop_requests", "release_rollout", "onload_trainer"] + assert node.role is Role.TRAINER and gpu.holder == "trainer" + assert r1.seconds == 4.0 and r1.phases == {"stop_requests": 1.0, "release_rollout": 1.0, + "onload_trainer": 2.0} + assert r1.cut.returned == ("t0.s0",) and r1.cut.dropped == ("t0",) + assert [e.sample_id for e in buf.entries] == ["t0.s0"] + + r2 = node.to_rollout(7, expected_checksum="sha-7") + assert gpu.calls[3:] == ["offload_trainer", "resume_rollout", "accept_requests"] + assert node.role is Role.ROLLOUT and gpu.holder == "rollout" and gpu.accepting + assert r2.seconds == 5.0 and r2.weight_checksum == "sha-7" + ev = r2.event() + assert ev["event"] == "node_conversion" and ev["direction"] == "to_rollout" and ev["version"] == 7 + + +def test_checksum_mismatch_fails_and_engine_stays_closed(): + node = _node(checksums={7: "sha-other"}) + node.to_trainer(Buffer()) + with pytest.raises(ConversionError, match="weight checksum sha-other != published sha-7"): + node.to_rollout(7, expected_checksum="sha-7") + assert node.role is Role.FAILED and not node.port.accepting + + +def test_wrong_role_and_port_failure(): + node = _node() + with pytest.raises(ConversionError, match="needs role trainer"): + node.to_rollout(1) + + def boom(): + raise RuntimeError("cuda oom") + node.port.onload_trainer = boom + with pytest.raises(ConversionError, match="rollout -> trainer failed: cuda oom"): + node.to_trainer(Buffer()) + assert node.role is Role.FAILED + + +def test_only_convertible_class_gets_a_node(): + with pytest.raises(NodeClassError, match="not convertible"): + ConvertibleNode(NodeSpec("r", NodeClass.ROLLOUT), FakeGpu(Clock())) + + +# ---------------------------------------------------------------- L7.3 capacity and guard +def _pool(clock, *, rollout=1, convertible=2, capacity=1.0): + specs = [NodeSpec(f"r{i}", NodeClass.ROLLOUT, spot=True) for i in range(rollout)] + specs.append(NodeSpec("t0", NodeClass.TRAINER)) + nodes = [ConvertibleNode(NodeSpec(f"c{i}", NodeClass.CONVERTIBLE, capacity=capacity), + FakeGpu(clock), clock=clock) for i in range(convertible)] + return NodePool(specs, nodes) + + +def test_last_rollout_node_is_never_converted(): + clock = Clock() + pool = _pool(clock, rollout=0, convertible=2) + assert pool.check_conversion(["c0"]) == 1.0 + with pytest.raises(LastRolloutNodeError): + pool.check_conversion(["c0", "c1"]) + pool.convertible["c0"].to_trainer(Buffer()) + with pytest.raises(LastRolloutNodeError): + pool.check_conversion(["c1"]) + + +def test_capacity_counts_conversion_and_trainer_class_does_no_rollout(): + clock = Clock() + pool = _pool(clock, rollout=2, convertible=1, capacity=0.5) + assert pool.total_capacity() == 2.5 # t0 (trainer class) is not counted + pool.convertible["c0"].to_trainer(Buffer()) + assert pool.capacity() == 2.0 and pool.doing_rollout() == ["r0", "r1"] + + +def test_idle_counts_only_time_with_neither_rollout_nor_training(): + clock = Clock() + node = _node(clock=clock) + clock.t += 100 # rollout: not idle + node.to_trainer(Buffer()) + node.set_training(True) + clock.t += 50 # training: not idle + node.set_training(False) + clock.t += 7 # neither + assert node.idle_seconds() == 7 + node.to_rollout(1) # the conversion time is not idle either + clock.t += 30 + assert node.idle_seconds() == 7 + + +def _sched(clock, pool, records=None, group_size=1): + s1 = S1Scheduler(S1Config(a_hard=1, start_seconds=10.0), clock=clock) + rec = (lambda kind, tx_id=None, **f: records.append({"kind": kind, **f})) if records is not None else None + return ConversionScheduler(s1, pool, group_size=group_size, record=rec) + + +def test_scheduler_converts_on_backlog_and_back_by_ski_rental(): + clock, records, buf = Clock(), [], Buffer() + pool = _pool(clock, rollout=1, convertible=2) + sched = _sched(clock, pool, records) + + assert sched.decide(BacklogStats(batches=1, samples=4)).to_trainer == () + d = sched.decide(BacklogStats(batches=2, samples=8)) # Q_on = a_hard + 1 = 2 + assert d.to_trainer == ("c0",) and d.lost_capacity == 1.0 and d.capacity_after == 2.0 + sched.apply(d, buf) + assert pool.capacity() == 2.0 + + clock.t += 1 + assert sched.decide(BacklogStats(batches=0, samples=0)).to_rollout == () # idle starts + clock.t += 10 # T_idle = T_s with an empty backlog + back = sched.decide(BacklogStats(batches=0, samples=0)) + assert back.to_rollout == ("c0",) and back.lost_capacity == -1.0 and back.note == "stop:STOPPED" + receipts = sched.apply(back, buf, version=4, expected_checksum="sha-4") + assert receipts[0].weight_checksum == "sha-4" and pool.capacity() == 3.0 + assert [e["direction"] for e in sched.events] == ["to_trainer", "to_rollout"] + assert [r["kind"] for r in records] == [CONVERSION_DECISION_KIND] * 2 + + +def test_scheduler_holds_instead_of_converting_the_last_rollout_node(): + clock, records = Clock(), [] + pool = _pool(clock, rollout=0, convertible=2) + sched = _sched(clock, pool, records, group_size=2) + d = sched.decide(BacklogStats(batches=5, samples=20)) + assert d.to_trainer == () and d.note == "hold:last_rollout_node" + assert records[-1]["decision"]["note"] == "hold:last_rollout_node" + sched.group_size = 1 + assert sched.decide(BacklogStats(batches=5, samples=20)).to_trainer == ("c0",) + + +def test_scheduler_finishes_at_data_end(): + clock = Clock() + pool = _pool(clock) + sched = _sched(clock, pool) + sched.apply(sched.decide(BacklogStats(batches=3, samples=12)), Buffer()) + d = sched.decide(BacklogStats(batches=0, samples=0, data_ended=True)) + assert d.to_rollout == ("c0",) and d.note == "stop:FINISHED" + + +def test_inner_decision_types_are_the_lifecycle_ones(): + # A transition state holds: the inner scheduler is asked with a non-startable state. + clock = Clock() + pool = _pool(clock) + pool.convertible["c0"].role = Role.TO_TRAINER + sched = _sched(clock, pool) + assert sched.decide(BacklogStats(batches=9, samples=36)).to_trainer == () + assert isinstance(sched.inner.decide(BacklogStats(batches=0, samples=0), + sched.lifecycle_state()), Hold) + + +# ---------------------------------------------------------------- Miles port +class FakeController: + def __init__(self, log): + self.log = log + + async def abort_all(self): + self.log.append("abort_all") + + async def offload(self, tags=None): + self.log.append(("offload", tuple(tags))) + + async def onload(self, tags=None): + self.log.append(("onload", tuple(tags))) + + async def onload_kv(self): + self.log.append("onload_kv") + + +class FakeTrainer: + def __init__(self, log): + self.log = log + + def onload(self): + self.log.append("trainer.onload") + + def offload(self): + self.log.append("trainer.offload") + + +def _run(coro): + try: + coro.send(None) + except StopIteration as stop: + return stop.value + raise AssertionError("fake coroutine did not finish") + + +def test_miles_port_call_order_both_directions(): + log: list = [] + port = MilesConversionPort( + controller=FakeController(log), trainer=FakeTrainer(log), run=_run, + load_weights=lambda v: log.append(("load_weights", v)), + read_checksum=lambda: (log.append("read_checksum"), "sha-5")[1], + collect_inflight=lambda: (log.append("collect"), [InflightTrajectory("t9")])[1], + set_admission=lambda on: log.append(("admission", on)), + memory_tags=lambda: ["weights", "kv_cache", "cuda_graph"], weight_tags=lambda: ["weights"]) + clock = Clock() + node = ConvertibleNode(NodeSpec("c0", NodeClass.CONVERTIBLE), port, clock=clock) + r = node.to_trainer(Buffer()) + assert r.cut.dropped == ("t9",) + node.to_rollout(5, expected_checksum="sha-5") + assert log == [("admission", False), "collect", "abort_all", + ("offload", ("weights", "kv_cache", "cuda_graph")), "trainer.onload", + "trainer.offload", ("onload", ("weights",)), ("load_weights", 5), "onload_kv", + "read_checksum", ("admission", True)] + + +class FakeEngine: + def __init__(self, url): + self.server_url = url + self.version = None + + async def update_weight_version(self, stamp): + self.version = stamp + + async def get_weight_version(self): + return self.version + + +class ChecksumController(FakeController): + """Engines whose ``check_weights(checksum)`` body is the last loaded version.""" + + def __init__(self, log, engines=1): + super().__init__(log) + self.loaded = None + self.engines = [FakeEngine(f"http://e{i}") for i in range(engines)] + self.skew = False + + async def start_update_weights(self): + self.log.append("start_update_weights") + return types.SimpleNamespace(rollout_engines=self.engines, + snapshot_cell_id_to_hashes={f"c{i}": "h" for i in range(len(self.engines))}) + + async def end_update_weights(self, snapshot_cell_id_to_hashes): + self.log.append("end_update_weights") + + async def abort_update_weights(self): + self.log.append("abort_update_weights") + + async def get_cell_statuses(self): + return {f"c{i}": types.SimpleNamespace(phase="Running") for i in range(len(self.engines))} + + async def check_weights(self, action): + assert action == "checksum" + return [{"w": f"v{self.loaded}", "e": i if self.skew else 0} for i in range(len(self.engines))] + + +def _publisher(log, controller, *, lora_rank=0): + async def update_weights(args, actor, executor, ctrl): + log.append("update_weights") + ctrl.loaded = 7 + + publisher = MilesPublisher(args=argparse.Namespace(offload_rollout=True, lora_rank=lora_rank), + actor_model=None, rollout_executor=None, + inference_controller=controller, update_weights=update_weights, + flatten_checksums=list, runner=LoopRunner()) + + async def warmup(engine, adapter, timeout_s): + log.append(("lora_warmup", engine.server_url)) + return {} + + publisher.lora_warmup = warmup + return publisher + + +def test_publisher_load_version_returns_the_engine_checksum_digest(): + log: list = [] + controller = ChecksumController(log, engines=2) + publisher = _publisher(log, controller) + assert publisher.published_checksum is None + with pytest.raises(PublicationError, match="needs a load_version"): + publisher.served_checksum() + publisher.load_version(7) + digest = publisher.served_checksum() + assert digest == engine_checksum_digest({"e": 0, "w": "v7"}) + # no onload, no KV: the port does those; no LoRA warm-up for a full-parameter run + assert log == ["update_weights", "start_update_weights", "end_update_weights"] + assert [e.version for e in controller.engines] == ["7", "7"] + publisher._reference = ("token", {"w": "v7", "e": 0}) + assert publisher.published_checksum == digest + + +def test_publisher_load_version_warms_the_lora_adapter_before_the_read_back(): + log: list = [] + controller = ChecksumController(log, engines=2) + publisher = _publisher(log, controller, lora_rank=16) + publisher.load_version(7) + publisher.served_checksum() + assert log == ["update_weights", "start_update_weights", "end_update_weights", + ("lora_warmup", "http://e0"), ("lora_warmup", "http://e1")] + + +def test_publisher_load_version_fails_when_engines_differ(): + log: list = [] + controller = ChecksumController(log, engines=2) + controller.skew = True + with pytest.raises(PublicationError, match="2 different weight checksums"): + publisher = _publisher(log, controller) + publisher.load_version(7) + publisher.served_checksum() + + +class FakePool: + def __init__(self, log, sample): + self.log = log + self.sample = sample + + def load_sample(self): + return self.sample + + def set_admission(self, open_): + self.log.append(("admission", open_)) + + +def test_engine_requests_inflight_counts_running_and_queued_or_nothing_when_unknown(): + assert [t.trajectory_id for t in engine_requests_inflight( + FakePool([], {"running_requests": 2, "queued_requests": 1}))] == [ + "engine-request-0", "engine-request-1", "engine-request-2"] + assert engine_requests_inflight(FakePool([], None)) == [] + assert engine_requests_inflight(FakePool([], {"running_requests": None})) == [] + + +def test_miles_conversion_port_wires_publisher_pool_and_trainer(): + log: list = [] + controller = ChecksumController(log) + publisher = _publisher(log, controller, lora_rank=16) + port = miles_conversion_port(controller=controller, + rollout=FakePool(log, {"running_requests": 1, "queued_requests": 0}), + trainer=FakeTrainer(log), publisher=publisher, run=_run) + port._memory_tags = lambda: ["weights", "kv_cache", "cuda_graph"] + port._weight_tags = lambda: ["weights"] + node = ConvertibleNode(NodeSpec("c0", NodeClass.CONVERTIBLE), port, clock=Clock()) + assert node.to_trainer(Buffer()).cut.dropped == ("engine-request-0",) + expected = engine_checksum_digest({"e": 0, "w": "v7"}) + receipt = node.to_rollout(7, expected_checksum=expected) + assert receipt.weight_checksum == expected + assert log == [("admission", False), "abort_all", + ("offload", ("weights", "kv_cache", "cuda_graph")), "trainer.onload", + "trainer.offload", ("onload", ("weights",)), "update_weights", + "start_update_weights", "end_update_weights", "onload_kv", + # the LoRA warm-up generates: only with the KV cache back (G-L7.4) + ("lora_warmup", "http://e0"), ("admission", True)] diff --git a/tests/test_rl_elastic_debug_pause.py b/tests/test_rl_elastic_debug_pause.py index e5832573d..76df5267b 100644 --- a/tests/test_rl_elastic_debug_pause.py +++ b/tests/test_rl_elastic_debug_pause.py @@ -53,6 +53,7 @@ def test_default_off_island_task_envs_unchanged(monkeypatch, tmp_path): from yeto.gpu_spec import parse_gpu_spec _fake_sky(monkeypatch) + monkeypatch.setattr("yeto.launch.rl_island_task.submitted_unix", lambda: 1.0) monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "k") base = ("--gpu", "modal:1xh100,modal:1xh100", "--rl-engine", "ports", "--rl-island-scheduling", "elastic") diff --git a/tests/test_rl_engine_driver.py b/tests/test_rl_engine_driver.py index 40340514c..cf48504d1 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -947,3 +947,20 @@ def test_miles_trainer_group_publish_offloaded_follows_offload_train(): assert prop.fget(SimpleNamespace(_args=SimpleNamespace(colocate=True, offload_train=True))) assert not prop.fget(SimpleNamespace(_args=SimpleNamespace(colocate=True, offload_train=False))) assert not prop.fget(SimpleNamespace(_args=SimpleNamespace(colocate=False, offload_train=True))) + + +def test_first_optimizer_step_records_the_last_start_phase_once(tmp_path): + from yeto.rl.engine.start_phases import START_PHASES, StartPhases + + emitted = [] + clock = iter(range(100, 200)) + phases = StartPhases(lambda event, **f: emitted.append((event, f)), clock=lambda: float(next(clock))) + for name in START_PHASES[:-1]: + phases.mark(name) + driver = _driver(_engine(), LocalOnlySync(3), tmp_path) + driver.start_phases = phases + driver.run() + names = [f["phase"] for e, f in emitted if e == "trainer_start_phase"] + assert names == list(START_PHASES) + [summary] = [f for e, f in emitted if e == "trainer_start_summary"] + assert summary["total_s"] == 6.0 and summary["missing"] == [] diff --git a/tests/test_rl_full_parameter_local_state.py b/tests/test_rl_full_parameter_local_state.py new file mode 100644 index 000000000..6084acf71 --- /dev/null +++ b/tests/test_rl_full_parameter_local_state.py @@ -0,0 +1,155 @@ +"""CPU tests for the no-sync full-parameter policy state (T-FPUB, L5.4 run i). + +GPU run s24-l54-pair-20261014i: a full-parameter trainer island without an outer +sync published through the LoRA digest export and failed with "Megatron produced +an empty or duplicate LoRA mapping". A full-parameter island now gets +``MilesFullParameterLocalPolicyState``: a trainer-side content digest of every +model parameter, no LoRA export. +""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from test_rl_miles_adapter_state import FakeTrainGroup # noqa: E402 +from yeto.rl.adapters.miles import state_plugin as sp # noqa: E402 +from yeto.rl.adapters.miles.entry import _policy_state_for # noqa: E402 +from yeto.rl.adapters.miles.publish import MilesPublisher, PublicationError, _payload_of # noqa: E402 +from yeto.rl.adapters.miles.state import ( # noqa: E402 + MilesFullParameterLocalPolicyState, + MilesPolicyState, + PolicyStateError, +) +from yeto.rl.engine import bridges # noqa: E402 +from yeto.rl.engine.driver import IslandDriver # noqa: E402 +from yeto.rl.engine.policy_digest import UnsupportedLayoutError, is_resident # noqa: E402 +from yeto.rl.engine.trainable_state import LAYOUT_FRAGMENTS # noqa: E402 + +REV = "r" * 40 +CFG = "c" * 64 +LAYOUT = "d" * 64 + + +class FullRank: + """A full-parameter Megatron rank: plain bf16 weights, no LoRA adapters.""" + + def __init__(self, main: bool, seed: int = 0) -> None: + torch.manual_seed(seed) + self.model = [torch.nn.Linear(4, 3).to(torch.bfloat16)] + self.args = SimpleNamespace(tensor_model_parallel_size=1, pipeline_model_parallel_size=1, + expert_model_parallel_size=1, hf_checkpoint="/model") + self._is_first_replica_megatron_main_rank = main + + def train_step(self) -> None: + with torch.no_grad(): + self.model[0].weight.add_(1.0) + sp.bump_weights_version() + + +def _ranks(n=2): + return [FullRank(main=(i == 0), seed=i) for i in range(n)] + + +def _state(ranks): + group = FakeTrainGroup(ranks) + return group, MilesFullParameterLocalPolicyState( + actor_model=group, base_model_revision=REV, config_hash=CFG, expected_layout_hash=LAYOUT) + + +class _Driver: + def __init__(self, ps): + self.policy_state = ps + self.events = [] + + def export_local_resident(self, *, policy_version): + return IslandDriver.export_local_resident(self, policy_version=policy_version) + + def export_local(self): + return self.policy_state.export() + + def emit(self, *a, **k): + self.events.append((a, k)) + + +def test_lora_state_on_a_full_parameter_trainer_reproduces_the_l54_failure(monkeypatch): + """Before the fix the no-sync island used the LoRA state: the bridge finds no + adapter on a full-parameter model and the digest export fails.""" + + hf_export = ModuleType("miles.backends.megatron_utils.hf_export") + no_adapters = SimpleNamespace(build_adapter_conversion_tasks=lambda model: {}) + hf_export._get_hf_bridge = lambda path: SimpleNamespace(_model_bridge=no_adapters) + monkeypatch.setitem(sys.modules, "miles.backends.megatron_utils.hf_export", hf_export) + group = FakeTrainGroup(_ranks()) + lora = MilesPolicyState(actor_model=group, base_model_revision=REV, config_hash=CFG) + with pytest.raises(sp.StatePluginError, match="empty or duplicate LoRA mapping"): + IslandDriver.export_local_resident(_Driver(lora), policy_version=0) + + +def test_policy_state_selection_follows_the_parameter_mode(): + group = FakeTrainGroup(_ranks()) + kwargs = dict(actor_model=group, base_model_revision=REV, config_hash=CFG, + expected_layout_hash=LAYOUT, runner=None) + assert isinstance(_policy_state_for(None, full_parameter=True, **kwargs), + MilesFullParameterLocalPolicyState) + lora = _policy_state_for(None, **kwargs) + assert type(lora) is MilesPolicyState + + +def test_full_parameter_digest_is_resident_and_content_addressed(): + ranks = _ranks() + group, ps = _state(ranks) + first = ps.export_digest(policy_version=2) + assert group.calls == [sp.EXPORT_FULL_DIGEST] + assert is_resident(first) and first.layout == LAYOUT_FRAGMENTS + assert first.policy_version == 2 and first.layout_hash == LAYOUT == ps.layout_hash + assert first.weights_mark == sp.current_weights_version() + # payload bytes: every rank's bf16 weight (3x4) and bias (3) + assert _payload_of(first)[1] == 2 * (12 + 3) * 2 + # same weights -> same hash, whatever the version; a rebuilt state agrees + assert ps.export(policy_version=5).policy_tensor_hash() == first.policy_tensor_hash() + assert _state(_ranks())[1].export_digest(policy_version=0).policy_tensor_hash() == \ + first.policy_tensor_hash() + # rank order is part of the identity + assert _state(list(reversed(_ranks())))[1].export_digest(policy_version=0).policy_tensor_hash() != \ + first.policy_tensor_hash() + ranks[1].train_step() + assert ps.export_digest(policy_version=2).policy_tensor_hash() != first.policy_tensor_hash() + + +def test_publisher_holds_check_and_refusals(): + ranks = _ranks() + group, ps = _state(ranks) + publisher = MilesPublisher(args=SimpleNamespace(), actor_model=None, rollout_executor=None, + inference_controller=None, export_trainer_state=ps.export) + resident = ps.export_digest(policy_version=1) + group.calls.clear() + assert publisher._check_trainer_holds(resident) == resident.policy_tensor_hash() + assert group.calls == [sp.WEIGHTS_VERSION] + ranks[0].train_step() + with pytest.raises(PublicationError, match="trainer weights differ"): + publisher._check_trainer_holds(resident) + with pytest.raises(UnsupportedLayoutError): + resident.materialize() + with pytest.raises(PolicyStateError, match="never applies"): + ps.apply(resident, optimizer="preserve", local_step=0) + with pytest.raises(PolicyStateError, match="bound layout hash"): + MilesFullParameterLocalPolicyState(actor_model=group, base_model_revision=REV, + config_hash=CFG, expected_layout_hash=None) + + +@pytest.mark.parametrize("fastpath", ["1", "0"]) +def test_local_only_sync_publishes_a_full_parameter_state(monkeypatch, fastpath): + monkeypatch.setenv("YETO_RL_PUBLISH_FASTPATH", fastpath) + _group, ps = _state(_ranks()) + sync = bridges.LocalOnlySync(num_rollout=3) + start = sync.start(_Driver(ps)) + assert is_resident(start.state) and start.state.policy_version == 0 + monkeypatch.setattr(bridges, "asdict", lambda s: {}) + boundary = sync.boundary(_Driver(ps), rollout_id=0, stats=SimpleNamespace()) + assert is_resident(boundary.state) and boundary.state.policy_version == 1 + assert boundary.state.policy_tensor_hash() == start.state.policy_tensor_hash() diff --git a/tests/test_rl_grpo_knobs.py b/tests/test_rl_grpo_knobs.py index 9045b7371..402c1c642 100644 --- a/tests/test_rl_grpo_knobs.py +++ b/tests/test_rl_grpo_knobs.py @@ -523,7 +523,7 @@ def test_examples_build_and_translate(name): spec = AlgorithmSpec.from_json_file(str(EXAMPLES / f"{name}.json")) # plugin hashes are the current sources (a dispatcher/reducer edit must # regenerate the examples: the hash is the plugin identity) - spec.verify_plugins() + assert spec.verify_plugins() == () assert spec.rejections() == [] cfg = make_config() if spec.sampling.over_sampling_batch_size is not None: @@ -576,10 +576,12 @@ def test_post_process_rechecks_pipeline_identity(): rp.post_process(args, []) -def test_stale_plugin_hash_rejected(): +def test_stale_plugin_hash_warns_not_rejected(): spec = pipeline_spec(reward_shapers=[OVERLONG]) stale = spec.replace(plugins=[{"path": p.path, "sha256": "1" * 64} for p in spec.plugins]) - assert any("source hash differs" in p for p in stale.rejections()) + assert stale.rejections() == [] + assert all("source hash mismatch" in w for w in stale.verify_plugins()) + assert len(stale.verify_plugins()) == len(spec.plugins) def test_unknown_transform_refused_at_launch(): diff --git a/tests/test_rl_identity_tiering.py b/tests/test_rl_identity_tiering.py new file mode 100644 index 000000000..93ecfca7c --- /dev/null +++ b/tests/test_rl_identity_tiering.py @@ -0,0 +1,80 @@ +"""rl-identity-tiering L9.1/L9.2: semantic hash, dual write, tolerant PluginRef.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from yeto.provenance import CODE_PACK_MODE_SOURCE_TREE +from yeto.rl.algos import mismatch_correction as mc +from yeto.rl.engine.algorithm import ( + AlgorithmSpec, + AlgorithmSpecError, + PluginRef, + plugin_source_sha256, +) +from yeto.rl.engine.driver import IslandDriver + +OBSERVE = mc.OBSERVE_PATH +STALE = "1" * 64 + + +def _custom(function: dict, **correction) -> AlgorithmSpec: + return AlgorithmSpec(correction={"method": "custom", "mismatch_metrics": True, + "function": function, **correction}) + + +def test_plugin_source_change_keeps_semantic_hash(): + current = _custom({"path": OBSERVE, "sha256": plugin_source_sha256(OBSERVE)}) + changed = _custom({"path": OBSERVE, "sha256": STALE}) + unpinned = _custom({"path": OBSERVE}) + assert current.sha256() != changed.sha256() # the full spec hash still sees the source + assert current.semantic_sha256() == changed.semantic_sha256() == unpinned.semantic_sha256() + # top-level plugins list: the path stays, the hash goes + listed = AlgorithmSpec(plugins=[{"path": OBSERVE, "sha256": STALE}]) + assert listed.semantic_dict()["spec"]["plugins"] == [{"path": OBSERVE}] + + +def test_configuration_change_changes_semantic_hash(): + base = _custom({"path": OBSERVE}) + assert base.semantic_sha256() != AlgorithmSpec().semantic_sha256() + assert AlgorithmSpec(kl_coef=0.01).semantic_sha256() != AlgorithmSpec().semantic_sha256() + other_path = _custom({"path": mc.ICEPOP_PATH}) + assert other_path.semantic_sha256() != base.semantic_sha256() + + +def test_plugin_ref_without_sha256_keeps_old_canonical_forms(): + ref = PluginRef.from_dict({"path": OBSERVE}) + assert ref.sha256 is None and ref.to_dict() == {"path": OBSERVE} + assert ref.verify(import_callable=True) is None + pinned = PluginRef.from_path(OBSERVE) + assert pinned.to_dict() == {"path": OBSERVE, "sha256": plugin_source_sha256(OBSERVE)} + with pytest.raises(AlgorithmSpecError, match="optional sha256"): + PluginRef.from_dict({"sha256": STALE}) + with pytest.raises(AlgorithmSpecError, match="optional sha256"): + PluginRef.from_dict({"path": OBSERVE, "sha256": STALE, "extra": 1}) + with pytest.raises(AlgorithmSpecError, match="64 hex"): + PluginRef.from_dict({"path": OBSERVE, "sha256": "abc"}) + + +def test_old_spec_file_with_stale_plugin_hash_loads_and_warns(tmp_path): + path = tmp_path / "old.json" + path.write_text(json.dumps(_custom({"path": OBSERVE, "sha256": STALE}).to_dict())) + spec = AlgorithmSpec.from_json_file(str(path)) + warnings = spec.verify_plugins() + assert len(warnings) == 1 + assert f"spec {STALE}, runtime {plugin_source_sha256(OBSERVE)}" in warnings[0] + + +def test_driver_start_code_identity(): + spec = _custom({"path": OBSERVE}) + without = SimpleNamespace(algorithm=spec, code_pack_digest=None) + assert IslandDriver._code_identity(without) == {} # rl_driver_start unchanged + with_digest = SimpleNamespace(algorithm=spec, code_pack_digest="a" * 64) + assert IslandDriver._code_identity(with_digest) == { + "semantic_sha256": spec.semantic_sha256(), + "code_pack_digest": "a" * 64, + "code_pack_mode": CODE_PACK_MODE_SOURCE_TREE, + } diff --git a/tests/test_rl_launcher_island_failure.py b/tests/test_rl_launcher_island_failure.py index 17544f3b1..eaa3845f8 100644 --- a/tests/test_rl_launcher_island_failure.py +++ b/tests/test_rl_launcher_island_failure.py @@ -276,6 +276,13 @@ def test_run_manifest_records_image_and_pins(monkeypatch, tmp_path, capsys): assert manifest["rl_image"] == rl.MILES_NEXT_IMAGE # the ports default, no override assert manifest["miles_commit"] == rl.MILES_NEXT_COMMIT assert manifest["sglang_commit"] == rl.SGLANG_NEXT_COMMIT + # rl-identity-tiering L9.1 dual write (transition: digest = source tree hash) + from yeto.rl.engine.algorithm import AlgorithmSpec + + assert manifest["code_pack_digest"] == manifest["source_sha256"] == args.source_sha256 + assert manifest["code_pack_mode"] == "source-tree" + assert manifest["semantic_sha256"] == AlgorithmSpec.from_dict( + _json.loads(args.rl_algorithm_spec_json)).semantic_sha256() out = capsys.readouterr().out assert f"image {rl.MILES_NEXT_IMAGE}" in out and rl.MILES_NEXT_COMMIT in out diff --git a/tests/test_rl_launcher_verda.py b/tests/test_rl_launcher_verda.py index 5d5a9c45f..ddfa9769c 100644 --- a/tests/test_rl_launcher_verda.py +++ b/tests/test_rl_launcher_verda.py @@ -86,6 +86,7 @@ def test_verda_island_without_registry_login_pulls_anonymously(monkeypatch): def test_other_clouds_are_untouched_by_the_in_vm_path(monkeypatch, gpu): """The feature is keyed on IN_VM_DOCKER_CLOUDS only: with the set emptied the Nebius/AWS task is identical, and the normal task carries no in-VM markers.""" + monkeypatch.setattr("yeto.launch.rl_island_task.submitted_unix", lambda: 1.0) _, task = _task_for(gpu, monkeypatch, env=LOGIN) monkeypatch.setattr("yeto.launch.rl_island_task.IN_VM_DOCKER_CLOUDS", frozenset()) _, plain = _task_for(gpu, monkeypatch, env=LOGIN) diff --git a/tests/test_rl_merge_metrics.py b/tests/test_rl_merge_metrics.py new file mode 100644 index 000000000..489a7d6d6 --- /dev/null +++ b/tests/test_rl_merge_metrics.py @@ -0,0 +1,165 @@ +"""rl-trainer-start-stop L5.5 / L5.7: merge metrics folded from the elastic +syncer tape, and the tape clock fields. No Ray, no GPU.""" +from __future__ import annotations + +import json +import os +import socket +import subprocess +import time + +import pytest + +from yeto.rl.engine.merge_metrics import MergeMetricsFolder +from yeto.wandb_tape import record_metrics, tape_metrics + + +def _fold(records): + folder = MergeMetricsFolder() + return [m for m in (folder.feed(r) for r in records) if m is not None] + + +def test_folder_hand_computed_shares_carry_and_rejections(): + # Round v1: island 1 on time (w = 10^2/1 = 100), island 3 late by one + # (w = 30^2/1 * 0.5 = 450), island 9 not a member, island 2 unknown version. + records = [ + {"kind": "delta_carried_over", "island_id": 3, "base": 0, "lag": 1, "discount": 0.5}, + {"kind": "delta_accepted", "island_id": 1}, + {"kind": "delta_rejected", "island_id": 9, "reason": "not_member"}, + {"kind": "delta_rejected", "island_id": 2, "reason": "unknown_version"}, + {"kind": "outer_step", "base_version": 1, "arrived": 1, "timed_out": True, + "raw_weights": [[1, 1, 100.0], [3, 0, 450.0]], + "weights": [[1, 1, 100 / 550], [3, 0, 450 / 550]], "absent": [2, 3], + "rejected": [[9, "not_member"], [2, "unknown_version"]], + "time_unix": 1000.0, "step_seconds": 4.0, "last_arrival_island": 1}, + {"kind": "delta_accepted", "island_id": 1}, + {"kind": "outer_step", "base_version": 2, "arrived": 1, "timed_out": False, + "raw_weights": [[1, 2, 100.0]], "weights": [[1, 2, 1.0]], "absent": [], "rejected": []}, + ] + first, second = _fold(records) + assert first["global_step"] == 2 and first["merge/base_version"] == 1 + assert first["merge/island/1/weight_share"] == pytest.approx(100 / 550) + assert first["merge/island/3/weight_share"] == pytest.approx(450 / 550) + assert first["merge/carried_weight_share"] == pytest.approx(450 / 550) + assert first["merge/submitted"] == 4 + assert first["merge/carried_fraction"] == 0.25 + assert first["merge/carried_discount_mean"] == 0.5 + assert first["merge/rejected_fraction"] == 0.5 + assert first["merge/rejected_fraction/not_member"] == 0.25 + assert first["merge/rejected_fraction/unknown_version"] == 0.25 + assert first["merge/step_seconds"] == 4.0 and first["merge/last_arrival_island"] == 1 + # The folder starts again after each step; absent values stay absent. + assert second["merge/submitted"] == 1 and second["merge/rejected_fraction"] == 0.0 + assert second["merge/carried_weight_share"] == 0.0 + assert "merge/carried_discount_mean" not in second and "merge/step_seconds" not in second + + +def test_old_tape_without_rejected_list_counts_rejection_events(): + records = [ + {"kind": "delta_accepted", "island_id": 1}, + {"kind": "delta_rejected", "island_id": 2, "reason": "carry_lag_exceeded"}, + {"kind": "outer_step", "base_version": 0, "arrived": 1, "timed_out": False, + "raw_weights": [[1, 0, 1.0]], "weights": [[1, 0, 1.0]], "absent": [2]}, + ] + (m,) = _fold(records) + assert m["merge/rejected_fraction/carry_lag_exceeded"] == 0.5 + + +def test_wandb_routes_elastic_records_through_the_folder_only(): + folder = MergeMetricsFolder() + assert record_metrics({"kind": "pool_join", "island_id": 1}, folder) is None + assert record_metrics({"kind": "delta_accepted", "island_id": 1}, folder) is None + step = {"kind": "outer_step", "base_version": 0, "weights": [[1, 0, 1.0]]} + assert record_metrics(step, folder)["merge/island/1/weight_share"] == 1.0 + strict = {"step": 3, "responders": []} + assert record_metrics(strict, folder) == tape_metrics(strict) + + +def _wait_tape(tape, predicate, timeout=10.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if tape.exists(): + rows = [json.loads(line) for line in tape.read_text().splitlines() if line.strip()] + if predicate(rows): + return rows + time.sleep(0.05) + raise AssertionError(f"tape condition not met: {tape.read_text() if tape.exists() else 'no tape'}") + + +@pytest.mark.skipif(not os.environ.get("YETO_TEST_ELASTIC_SYNCER"), + reason="needs YETO_TEST_ELASTIC_SYNCER (syncer binary with elastic mode)") +def test_real_syncer_tape_merge_shares_and_clock(tmp_path): + """Real elastic syncer, two client islands (theta 0.5: one arrival steps). + + v0: island 1 alone -> share 1.0. Island 2 then sends on base 0 (lag 1, + discount 0.5, w = 20^2/2 * 0.5 = 100); island 1 on base 1 (w = 100) -> + v1 shares 0.5 / 0.5. Island 2 sends on base 7 (unknown_version), then on + base 2 -> v2 rejected fraction 1/2.""" + from yeto.rl.elastic_client import ElasticClientConfig, ElasticIslandClient + + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + key = "merge-metrics" + tape = tmp_path / "syncer-tape.jsonl" + syncer = subprocess.Popen( + [os.environ["YETO_TEST_ELASTIC_SYNCER"], "--port", str(port), "--learners", "2", + "--total-steps", "5", "--event-tape", str(tape), + "--island-scheduling-mode", "elastic", "--quorum-theta", "0.5", "--carry-gamma", "0.5", + "--soft-deadline-s", "60", "--q-min", "1", "--max-carry-lag", "2", + "--island-lease-s", "6", "--syncer-epoch", "1"], + env=dict(os.environ, YETO_ISLAND_HMAC_KEY=key), + stdout=open(tmp_path / "syncer.log", "w"), stderr=subprocess.STDOUT) + islands = {i: ElasticIslandClient(ElasticClientConfig(("127.0.0.1", port), i, lease_s=6.0, + syncer_epoch=1), key.encode()) + for i in (1, 2)} + started = time.time() + try: + for c in islands.values(): + for _ in range(50): + try: + c.join() + break + except OSError: + time.sleep(0.1) + islands[1].elastic_init([0.0] * 4) + assert islands[1].wait_base(newer_than=None, timeout_s=10.0) is not None + update = [0.1] * 4 + islands[1].delta_tensor(base_version=0, c_tokens=10, c_steps=1, update=update) + assert islands[2].wait_base(newer_than=0, timeout_s=10.0) is not None + islands[2].delta_tensor(base_version=0, c_tokens=20, c_steps=2, update=update) + _wait_tape(tape, lambda rows: any(r["kind"] == "delta_carried_over" for r in rows)) + islands[1].delta_tensor(base_version=1, c_tokens=10, c_steps=1, update=update) + assert islands[2].wait_base(newer_than=1, timeout_s=10.0) is not None + islands[2].delta_tensor(base_version=7, c_tokens=20, c_steps=2, update=update) + _wait_tape(tape, lambda rows: any(r["kind"] == "delta_rejected" for r in rows)) + islands[2].delta_tensor(base_version=2, c_tokens=20, c_steps=2, update=update) + rows = _wait_tape(tape, lambda rows: sum(r["kind"] == "outer_step" for r in rows) >= 3) + finally: + for c in islands.values(): + c.close() + syncer.kill() + syncer.wait(5) + + steps = _fold(rows) + assert [m["merge/base_version"] for m in steps[:3]] == [0, 1, 2] + v0, v1, v2 = steps[:3] + assert v0["merge/island/1/weight_share"] == 1.0 and v0["merge/submitted"] == 1 + assert v1["merge/island/1/weight_share"] == pytest.approx(0.5) + assert v1["merge/island/2/weight_share"] == pytest.approx(0.5) + assert v1["merge/carried_weight_share"] == pytest.approx(0.5) + assert v1["merge/carried_fraction"] == 0.5 and v1["merge/carried_discount_mean"] == 0.5 + assert v2["merge/rejected_fraction/unknown_version"] == 0.5 + assert v2["merge/island/2/weight_share"] == 1.0 + + # L5.7 clock: every line has time_unix within the test window; deltas + # carry their arrival offset; outer steps their duration and last island. + ended = time.time() + assert all(started - 1 <= r["time_unix"] <= ended + 1 for r in rows) + for r in rows: + if r["kind"] in ("delta_accepted", "delta_carried_over"): + assert r["arrival_offset_s"] >= 0.0 + outer = [r for r in rows if r["kind"] == "outer_step"][:3] + assert [r["last_arrival_island"] for r in outer] == [1, 1, 2] + assert all(r["step_seconds"] >= 0.0 for r in outer) + assert v1["merge/step_seconds"] == outer[1]["step_seconds"] diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 78b5e2a5a..9121feea0 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -323,6 +323,8 @@ def test_leaf_policy_matches_run_config_fields(): leaves |= {"eval", "geometry.moe", "algorithm.lr_schedule"} # the None variants # rl-algo-critic-family 2.4: critic run config (None here) and its leaves leaves |= {"algorithm.critic", "algorithm.critic.critic_load", "algorithm.critic.init_sha256"} + # L4.6: the sample buffer config (None here) and its leaves + leaves |= {f"sample_buffer.{f.name}" for f in dataclasses.fields(rc.SampleBufferConfig)} assert leaves == set(mc.LEAF_POLICY) diff --git a/tests/test_rl_miles_carry_over.py b/tests/test_rl_miles_carry_over.py index 1e11921f0..e94f3d267 100644 --- a/tests/test_rl_miles_carry_over.py +++ b/tests/test_rl_miles_carry_over.py @@ -239,11 +239,12 @@ def test_governor_fallback_lowers_the_rollout_limit_through_the_sink(sink, tmp_p def test_miles_stage_two_support_and_single_turn_check(): from yeto.rl.adapters.miles.policy_age import SUPPORT, check_task, policy_age_argv - # agentic-rollout-utilization 5 (stage 3) raised the declared stage; limit still <= 1 - assert (SUPPORT.stage, SUPPORT.max_policy_age) == (3, 1) + # agentic-rollout-utilization 5 (stage 3); L6.9 raised the cap to 4 (A_hard >= 2) + assert (SUPPORT.stage, SUPPORT.max_policy_age) == (3, 4) SUPPORT.check(1) - with pytest.raises(PolicyAgeError, match="at most 1"): - SUPPORT.check(2) + SUPPORT.check(2) + with pytest.raises(PolicyAgeError, match="at most 4"): + SUPPORT.check(5) assert policy_age_argv(0) == () and policy_age_argv(1) == ("--partial-rollout",) check_task(1) check_task(0, custom_agent="x.y", custom_generate="x.z") diff --git a/tests/test_rl_miles_cut_plugin.py b/tests/test_rl_miles_cut_plugin.py index 285d2b88c..706983f3f 100644 --- a/tests/test_rl_miles_cut_plugin.py +++ b/tests/test_rl_miles_cut_plugin.py @@ -24,21 +24,24 @@ def _batches(n): return [torch.randn(3, 6, generator=g) for _ in range(n)] -def test_same_shape_restore_matches_continuous_run_bitwise(tmp_path): +@pytest.mark.parametrize("full", [False, True], ids=["lora", "full-parameter"]) +def test_same_shape_restore_matches_continuous_run_bitwise(tmp_path, full): batches = _batches(3) # continuous run: 2 steps, then the frozen third batch - cont = make_rank(seed=0) + cont = make_rank(seed=0, full=full) torch.manual_seed(99) for b in batches: train_step(cont, b) # the RNG is global: finish the continuous run first # cut run: same 2 steps, save, "rebuild" (fresh rank, different init), restore, third batch - run = make_rank(seed=0) + run = make_rank(seed=0, full=full) torch.manual_seed(99) for b in batches[:2]: train_step(run, b) summary = save_cut_shard(run, directory=str(tmp_path), cut_id="c1") assert summary["has_optimizer_state"] and summary["has_rng"] and summary["scheduler_samples"] == 8 - fresh = make_rank(seed=5) + assert summary["tuning"] == ("full" if full else "lora") + assert summary["adapter_names"] == sorted(params(run)) + fresh = make_rank(seed=5, full=full) torch.manual_seed(12345) # the rebuilt process has other RNG state restored = restore_cut_shard(fresh, directory=str(tmp_path), files=[summary], cut_id="c1") assert restored["state_digest"] == summary["state_digest"] @@ -117,6 +120,20 @@ def test_non_lora_trainable_parameters_are_refused(tmp_path): assert re.search("other than LoRA", save_cut_shard(rank, directory=str(tmp_path), cut_id="c1")["refused"]) +def test_a_trainer_without_trainable_parameters_is_refused(tmp_path): + rank = make_rank(0, full=True) + rank.model[0].requires_grad_(False) + assert re.search("no trainable parameters", save_cut_shard(rank, directory=str(tmp_path), cut_id="c1")["refused"]) + + +def test_a_lora_cut_does_not_restore_into_a_full_parameter_trainer(tmp_path): + rank = make_rank(0) + train_step(rank, _batches(1)[0]) + s = save_cut_shard(rank, directory=str(tmp_path), cut_id="c1") + refused = restore_cut_shard(make_rank(0, full=True), directory=str(tmp_path), files=[s], cut_id="c1") + assert re.search("adapter names differ", refused["refused"]) + + @pytest.mark.parametrize( "args, match", [ diff --git a/tests/test_rl_mismatch_correction.py b/tests/test_rl_mismatch_correction.py index cce300a96..c303184a4 100644 --- a/tests/test_rl_mismatch_correction.py +++ b/tests/test_rl_mismatch_correction.py @@ -122,7 +122,7 @@ def test_observe_translation_and_identity(): other = AlgorithmSpec(correction={"method": "custom", "mismatch_metrics": True, "function": {"path": mc.OBSERVE_PATH, "sha256": "1" * 64}}) assert other.sha256() != s.sha256() != AlgorithmSpec().sha256() - s.verify_plugins(import_callable=True) + assert s.verify_plugins(import_callable=True) == () @pytest.mark.parametrize( @@ -130,7 +130,7 @@ def test_observe_translation_and_identity(): [ (lambda: spec(method="custom", function=mc.observe_ref().to_dict()), "mismatch_metrics"), (lambda: spec(method="custom", function=mc.observe_ref().to_dict(), - mismatch_metrics=True, tis_clip=2.0), "tis_clip"), + mismatch_metrics=True, tis_clip=2.0), "correction.upper"), (lambda: spec(method="custom", function=mc.observe_ref().to_dict(), mismatch_metrics=True, use_rollout_logprobs=True), "tis_with_rollout_logprobs"), @@ -193,15 +193,15 @@ def test_icepop_translation_and_hash(): "--custom-tis-function-path", mc.ICEPOP_PATH, "--get-mismatch-metrics", ] assert s.correction.function.sha256 == mc.ICEPOP_SOURCE_SHA256 - assert icepop(low=0.4).sha256() != s.sha256() + assert icepop(low=0.4).sha256() == s.sha256() # phase 1 (D8): bounds are not hashed @pytest.mark.parametrize( "build, needle", [ (lambda: spec(method="custom", function=mc.icepop_ref().to_dict(), tis_clip=5.0), - "tis_clip_low"), - (lambda: spec(method="custom", function=mc.icepop_ref().to_dict()), "tis_clip"), + "correction.lower"), + (lambda: spec(method="custom", function=mc.icepop_ref().to_dict()), "correction.upper"), (lambda: icepop(low=5.0, high=0.5), "5.0"), (lambda: icepop(low=5.0, high=5.0), "smaller"), ], @@ -310,7 +310,7 @@ def test_opsm_source_identity(): trainer, rollout = opsm(), opsm(source="rollout") assert trainer.correction.opsm_old_logprob_source == "trainer" assert opsm(source="trainer").sha256() == trainer.sha256() - assert trainer.sha256() != rollout.sha256() + assert trainer.sha256() == rollout.sha256() # phase 1 (D8): recorded, not hashed assert json.loads(rollout.canonical_json())["correction"]["opsm_old_logprob_source"] == "rollout" # Not selecting OPSM: the field (at its default) does not affect the hash. assert spec(opsm_old_logprob_source="trainer").sha256() == AlgorithmSpec().sha256() @@ -358,7 +358,8 @@ def test_mis_translation(): "tis_lower_bound": 0.999, "tis_mode": "mask", "tis_upper_bound": 1.001, "use_rs": False, } assert s.correction.function.sha256 == PluginRef.from_path(mc.MIS_PATH).sha256 - assert mis(mode="clip", low=0.5).sha256() != mis(mode="mask", low=0.5).sha256() + # phase 1 (D8): MIS fields are recorded, not hashed + assert mis(mode="clip", low=0.5).sha256() == mis(mode="mask", low=0.5).sha256() def test_runtime_attrs_unchanged_without_mis(): @@ -376,7 +377,7 @@ def test_runtime_attrs_unchanged_without_mis(): (lambda: mis(mode="truncate", low=0.5), "unused"), (lambda: mis(mode="clip", low=3.0, high=2.0), "smaller"), (lambda: mis(level="geometric", mis_batch_normalize=True), "token|sequence"), - (lambda: mis(tis_clip=2.0), "tis_clip"), + (lambda: mis(tis_clip=2.0), "correction.lower/upper"), (lambda: spec(mis_mode="mask"), "require the MIS function"), (lambda: tis(mis_upper_bound=2.0), "require the MIS function"), ], @@ -543,6 +544,6 @@ def test_doc_examples_pin_current_plugin_sources(tmp_path): spec = AlgorithmSpec.from_json_file(str(path)) for ref in spec.referenced_plugins(): if ref.path.startswith("yeto."): - ref.verify(import_callable=False) + assert ref.verify(import_callable=False) is None else: assert ref.sha256 == mc.ICEPOP_SOURCE_SHA256 diff --git a/tests/test_rl_mismatch_observe.py b/tests/test_rl_mismatch_observe.py index 72d181115..0f3387842 100644 --- a/tests/test_rl_mismatch_observe.py +++ b/tests/test_rl_mismatch_observe.py @@ -188,7 +188,7 @@ def test_icepop_weights(): def test_icepop_source_hash_pinned(): ref = PluginRef.from_path(mc.ICEPOP_PATH) assert ref.sha256 == mc.ICEPOP_SOURCE_SHA256 - ref.verify(import_callable=True) + assert ref.verify(import_callable=True) is None def test_icepop_all_out_of_interval_zero_gradient(local_log_probs): diff --git a/tests/test_rl_offpolicy_correction.py b/tests/test_rl_offpolicy_correction.py new file mode 100644 index 000000000..fd9460e59 --- /dev/null +++ b/tests/test_rl_offpolicy_correction.py @@ -0,0 +1,448 @@ +"""rl-offpolicy-correction-passthrough phase 1 (CPU): neutral correction fields, +phase-1 hash, Miles mapping, decoupled PPO refusals, cross-island warning, +offpolicy_gate, A_hard >= 2, start-version gap and observe-only bounds.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from yeto.rl.adapters.miles import carry_over +from yeto.rl.adapters.miles import offpolicy_correction as mo +from yeto.rl.adapters.miles.algorithm_flags import algorithm_argv +from yeto.rl.adapters.miles.config import MilesConfigError, check_extra_argv +from yeto.rl.algos import mismatch_correction as mc +from yeto.rl.engine import offpolicy_correction as oc +from yeto.rl.engine.algorithm import AlgorithmSpec, AlgorithmSpecError, CorrectionSpec +from yeto.rl.engine.island_ledger import StalenessPolicy +from yeto.rl.engine.policy_age import PolicyAgeError +from yeto.rl.engine.version_segments import TokenProvenance + + +def decoupled(upper: float = 2.0, **kw) -> AlgorithmSpec: + """Decoupled PPO with C = upper (the example-config value is 2.0).""" + return AlgorithmSpec(correction=CorrectionSpec(method="tis", lower=0.0, upper=upper, **kw)) + + +# -- L6.1 neutral fields ------------------------------------------------------ + + +def test_fields_written_in_full_with_defaults(): + group = CorrectionSpec(method="tis", lower=0.0, upper=2.0).to_dict() + for name in ("method", "level", "lower", "upper", "veto", "batch_normalize", "ratio_mask"): + assert name in group + assert (group["level"], group["ratio_mask"], group["veto"], group["batch_normalize"]) == ( + "token", "none", None, False) + + +def test_old_bound_names_are_aliases(): + old = CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.1) + assert old == CorrectionSpec(method="tis", upper=2.0, lower=0.1) + assert (old.tis_clip, old.tis_clip_low) == (2.0, 0.1) + assert CorrectionSpec.from_dict({"method": "tis", "tis_clip": 2, "tis_clip_low": 0}) == \ + CorrectionSpec(method="tis", upper=2.0, lower=0.0) + with pytest.raises(AlgorithmSpecError, match="old name of correction.upper"): + CorrectionSpec(method="tis", tis_clip=2.0, upper=3.0, lower=0.0) + + +@pytest.mark.parametrize("kw, needle", [ + (dict(method="tis", upper=2.0), "correction.upper and correction.lower"), + (dict(method="tis", lower=0.0), "correction.upper and correction.lower"), + (dict(ratio_mask="icepop"), "requires correction.method='tis'"), + (dict(method="tis", lower=0.0, upper=2.0, ratio_mask="bogus"), "correction.ratio_mask"), + (dict(method="tis", lower=0.0, upper=2.0, level="turn"), "correction.level"), + (dict(method="tis", lower=0.0, upper=2.0, level="geometric"), "ratio_mask='mis'"), + (dict(veto=1e-4), "correction.veto"), + (dict(batch_normalize=True), "correction.batch_normalize"), + (dict(level="sequence"), "correction.level"), + (dict(method="tis", lower=0.0, upper=2.0, veto=0.0), "correction.veto"), + (dict(method="tis", lower=0.0, upper=2.0, batch_normalize=1), "correction.batch_normalize"), +]) +def test_invalid_specs(kw, needle): + with pytest.raises(AlgorithmSpecError, match=needle.replace(".", r"\.").replace("(", r"\(")): + CorrectionSpec(**kw) + + +def test_valid_specs(): + CorrectionSpec(method="tis", lower=0.5, upper=5.0, ratio_mask="icepop") + CorrectionSpec(method="tis", lower=0.5, upper=2.0, ratio_mask="mis", level="geometric", + veto=1e-4, batch_normalize=True) + assert not decoupled().rejections() + assert any("smaller" in p for p in + AlgorithmSpec(correction=CorrectionSpec(method="tis", lower=3.0, upper=2.0)) + .rejections()) + + +# -- L6.2 phase-1 hash and records ------------------------------------------- + + +@pytest.mark.parametrize("correction", [ + CorrectionSpec(method="tis", lower=0.0, upper=2.0), + CorrectionSpec(method="tis", lower=0.0, upper=5.0), + CorrectionSpec(method="tis", lower=0.5, upper=5.0, ratio_mask="icepop"), + CorrectionSpec(method="tis", lower=0.5, upper=2.0, ratio_mask="mis", veto=1e-4), + CorrectionSpec(method="opsm", opsm_delta=1e-4), +]) +def test_correction_group_does_not_change_the_hash(correction): + base = AlgorithmSpec(entropy_coef=0.01) + spec = base.replace(correction=correction) + assert spec.sha256() == base.sha256() + # the full group still travels to the islands + assert AlgorithmSpec.from_dict(json.loads(spec.canonical_json())) == spec + + +def test_default_hashes_are_unchanged(): + # values from before the change (base c4ae38a6) + assert AlgorithmSpec().sha256() == \ + "27df1133c924e7a337e90246c7e1e30699017f461e9f35b1dd165a0263e65bea" + assert AlgorithmSpec(loss={"variant": "policy_loss", "eps_clip": 0.3}).sha256() == \ + "64f4876cb7a99d2650731844dfc7da29dbbe4b638dc3837f5916cf4e98c4f2f4" + assert AlgorithmSpec(entropy_coef=0.1).sha256() == \ + "e5a86b389f80e46d9e36bf5942687fc9d9a14106aac34159c680952724251e82" + + +def test_custom_function_source_hash_stays_in_the_identity(): + ref = mc.icepop_ref() + spec = AlgorithmSpec(correction={"method": "custom", "function": ref.to_dict(), + "tis_clip": 5.0, "tis_clip_low": 0.5}) + other = AlgorithmSpec(correction={"method": "custom", "tis_clip": 5.0, "tis_clip_low": 0.5, + "function": {"path": ref.path, "sha256": "1" * 64}}) + assert spec.sha256() != other.sha256() != AlgorithmSpec().sha256() + + +def test_manifest_section_records_group_mapping_and_anchor(): + spec = decoupled() + section = oc.run_correction_manifest(spec, {"0": "miles", "1": "miles"}) + record = section["islands"]["0"] + assert record["group"] == spec.correction.to_dict() + assert record["mapped"] == {"--use-tis": True, "--tis-clip": "2.0", "--tis-clip-low": "0.0"} + assert record["prox_anchor"] == "per_step" + assert section["warnings"] == [] + + +def test_backend_without_mapping_is_recorded_and_warned(): + from yeto.rl.engine import backends + + backends.register(backends.BackendEntry("fake", "x", {})) + try: + section = oc.run_correction_manifest(decoupled(), {"0": "fake"}) + assert section["islands"]["0"]["mapped"] is None + assert [w["backend"] for w in section["warnings"]] == ["fake"] + # a default group needs no mapping: no warning + assert oc.run_correction_manifest(AlgorithmSpec(), {"0": "fake"})["warnings"] == [] + finally: + backends.unregister("fake") + + +def test_island_tape_gets_record_and_warnings(tmp_path): + from yeto.rl.adapters.miles import island_entry + + events = [] + island_entry_append = island_entry._append_ports_event + island_entry._append_ports_event = lambda _a, _m, e: events.append(e) + try: + island_entry._append_correction_record(SimpleNamespace(), SimpleNamespace(), + AlgorithmSpec()) + assert events == [] # default tapes unchanged + spec = decoupled(level="sequence") + island_entry._append_correction_record(SimpleNamespace(), SimpleNamespace(), spec) + finally: + island_entry._append_ports_event = island_entry_append + assert events[0]["event"] == oc.CORRECTION_RECORD_EVENT + assert events[0]["group"]["level"] == "sequence" and events[0]["prox_anchor"] == "per_step" + assert [(e["event"], e["field"], e["backend"]) for e in events[1:]] == [ + (oc.CORRECTION_CONFIG_WARNING, "level", "miles")] + + +# -- L6.3 Miles mapping ------------------------------------------------------- + + +def test_miles_tis_row(): + assert algorithm_argv(decoupled()) == ["--use-tis", "--tis-clip", "2.0", "--tis-clip-low", "0.0"] + + +def test_miles_icepop_row(): + spec = AlgorithmSpec(correction=CorrectionSpec(method="tis", lower=0.5, upper=5.0, + ratio_mask="icepop")) + assert algorithm_argv(spec) == ["--use-tis", "--tis-clip", "5.0", "--tis-clip-low", "0.5", + "--custom-tis-function-path", mc.ICEPOP_PATH] + assert mc.selected_corrections(spec) == ("icepop",) + assert spec.to_legacy_runtime_attrs() == AlgorithmSpec().to_legacy_runtime_attrs() + + +def test_miles_mis_row(): + spec = AlgorithmSpec(correction=CorrectionSpec(method="tis", lower=0.5, upper=2.0, + ratio_mask="mis", level="sequence", + veto=1e-4, batch_normalize=True)) + assert algorithm_argv(spec) == ["--use-tis", "--tis-clip", "2.0", "--tis-clip-low", "0.5", + "--custom-tis-function-path", mc.MIS_PATH] + attrs = {k: v for k, v in spec.to_legacy_runtime_attrs().items() if not k.startswith("yeto_")} + assert attrs == { + "rs_level": "sequence", "rs_lower_bound": 0.5, "rs_upper_bound": 2.0, + "rs_veto_threshold": 1e-4, "tis_batch_normalize": True, "tis_level": "sequence", + "tis_lower_bound": 0.5, "tis_mode": "mask", "tis_upper_bound": 2.0, "use_rs": True} + assert mc.selected_corrections(spec) == ("mis_mask",) + assert mo.correction_mapping(spec).warnings == () + + +def test_miles_opsm_row(): + spec = AlgorithmSpec(correction=CorrectionSpec(method="opsm", opsm_delta=1e-4)) + assert algorithm_argv(spec) == ["--use-opsm", "--opsm-delta", "0.0001"] + + +@pytest.mark.parametrize("kw, field", [ + (dict(level="sequence"), "level"), + (dict(veto=1e-4), "veto"), + (dict(batch_normalize=True), "batch_normalize"), +]) +def test_miles_unmappable_fields_warn(kw, field): + mapping = mo.correction_mapping(decoupled(**kw)) + assert [w.field for w in mapping.warnings] == [field] + assert mapping.warnings[0].backend == "miles" + event = mapping.warnings[0].event() + assert event["event"] == "correction_config_warning" and "ran_with" in event + + +# -- L6.7 decoupled PPO: refused Miles flags ----------------------------------- + + +@pytest.mark.parametrize("flag", ["--use-rollout-logprobs", "--skip-actor-forward-only", + "--mask-offpolicy-in-partial-rollout"]) +def test_forbidden_miles_flag_is_refused_by_name(flag): + with pytest.raises(MilesConfigError, match=flag): + check_extra_argv([flag], decoupled()) + with pytest.raises(oc.CorrectionConfigError, match=flag): + mo.refuse_forbidden_argv(["--use-tis", flag]) + + +def test_rollout_logprobs_from_the_spec_is_refused_in_the_final_argv(): + spec = AlgorithmSpec(correction=CorrectionSpec(use_rollout_logprobs=True)) + assert "--use-rollout-logprobs" in mo.correction_mapping(spec).argv + with pytest.raises(oc.CorrectionConfigError, match="--use-rollout-logprobs"): + mo.refuse_forbidden_argv(algorithm_argv(spec)) + + +def test_decoupled_ppo_records_per_step_anchor(): + mapping = mo.correction_mapping(decoupled()) + assert mapping.prox_anchor == oc.PROX_PER_STEP + assert mapping.argv == ("--use-tis", "--tis-clip", "2.0", "--tis-clip-low", "0.0") + + +# -- L6.5 cross-island comparison --------------------------------------------- + + +def _record(upper: float, anchor: str) -> dict: + mapping = oc.CorrectionMapping(backend="x", prox_anchor=anchor) + return oc.correction_record(decoupled(upper).correction, mapping, backend="x") + + +def test_islands_with_different_upper_give_one_warning(): + event = oc.compare_island_records({"0": _record(2.0, "per_step"), "1": _record(5.0, "per_step")}) + assert event["event"] == "correction_config_warning" + assert event["differences"] == {"upper": {"0": 2.0, "1": 5.0}} + + +def test_different_prox_anchor_is_listed(): + event = oc.compare_island_records({"miles": _record(2.0, "per_step"), + "verl": _record(2.0, "first_local_step")}) + assert event["differences"] == {"prox_anchor": {"miles": "per_step", + "verl": "first_local_step"}} + assert oc.compare_island_records({"0": _record(2.0, "per_step"), + "1": _record(2.0, "per_step")}) is None + + +def test_launcher_manifest_section(capsys): + from tests.test_rl_launcher import _args + from yeto.launch.rl_ports import _offpolicy_correction_manifest + + args = _args([]) + section = _offpolicy_correction_manifest(args, decoupled(level="sequence")) + assert "0" in section["islands"] + assert {w["field"] for w in section["warnings"]} == {"level"} + assert "correction_config_warning" in capsys.readouterr().out + + +# -- L6.8 offpolicy_gate -------------------------------------------------------- + + +def test_dropped_length_histogram_fixed_buckets(): + hist = oc.length_histogram([300, 5000]) + assert hist == {"0-255": 0, "256-1023": 1, "1024-4095": 0, "4096-16383": 1, "16384+": 0} + + +def test_gate_event_granularity(): + event = oc.offpolicy_gate_event(backend="verl", counts={"kept": 3}, dropped_lengths={}, + truncation_by_turn={0: (1, 2)}, granularity="trajectory") + assert event["granularity"] == "trajectory" + assert event["truncation_by_turn"] == {"0": {"truncated": 1, "total": 2, "ratio": 0.5}} + with pytest.raises(ValueError): + oc.offpolicy_gate_event(backend="x", counts={}, dropped_lengths={}, + truncation_by_turn=None, granularity="turn") + + +class _Span: + def __init__(self, version, start, end): + self.version, self.abs_start, self.abs_end = str(version), start, end + + +class _Sample: + def __init__(self, group, spans, length, prompt=3, **kw): + self.group_index, self.response_length = group, length + self.tokens = [0] * (prompt + length) + self.weight_versions = [SimpleNamespace(spans=[_Span(v, prompt + a, prompt + b) + for v, a, b in spans])] + self.__dict__.update(kw) + + +def test_miles_buffer_filter_drops_feed_the_gate(): + args = SimpleNamespace() + old_a = [_Sample(0, [(1, 0, 300)], 300)] + old_b = [_Sample(1, [(1, 0, 5000)], 5000)] + fresh = [_Sample(2, [(4, 0, 10)], 10)] + kept = carry_over.carry_over_buffer_filter(args, [old_a, old_b, fresh], 4, + current_version=4, limit=2) + assert kept == [fresh] + fields = carry_over.take_carry_stats(args, 4) + gate = oc.gate_from_carry_over(fields, kept=1, backend="miles") + assert oc.GATE_DROPPED_LENGTHS not in fields + assert gate["counts"] == {"group_dropped": 2, "kept": 1, "unknown_version_dropped": 0} + hist = gate["dropped_length_histogram"]["group_dropped"] + assert (hist["256-1023"], hist["4096-16383"]) == (1, 1) + assert gate["granularity"] == "token" + + +def test_miles_truncation_by_turn_is_token_level(): + # two model turns (loss_mask runs) of 2 tokens, separated by one tool token + sample = _Sample(0, [(3, 0, 5)], 5, loss_mask=[1, 1, 0, 1, 1], + rollout_log_probs=[-0.5] * 5) + current = [0.5, -0.5, -0.5, 0.4, -0.5] # ratios e^1, 1, (tool), e^0.9, 1 + args = SimpleNamespace(tis_clip=2.0, tis_clip_low=0.0) + out = carry_over.estimate_cross_version_truncation(args, [sample], 4, lambda _s: current) + assert out[carry_over.TRUNCATED_BY_TURN] == {0: (1, 2), 1: (1, 2)} + + +def test_driver_tapes_offpolicy_gate_as_its_own_event(): + from yeto.rl.engine.driver import IslandDriver + + emitted = [] + fake = SimpleNamespace(emit=lambda name, **kw: emitted.append((name, kw)), + clock=lambda: 0.0, _labels=lambda: {}) + carry = {"max_policy_age": 2, "over_age_discarded_groups": 1, + oc.GATE_DROPPED_LENGTHS: {"group_dropped": [300]}, + oc.TRUNCATED_BY_TURN: {0: (1, 2)}} + batch = SimpleNamespace(carry_over=carry, policy_version=3, groups=(object(),)) + IslandDriver._emit_carry_over(fake, 3, batch) + assert [name for name, _ in emitted] == [oc.OFFPOLICY_GATE_EVENT, "rl_rollout_carry_over"] + assert emitted[0][1]["counts"]["kept"] == 1 + assert emitted[0][1]["counts"]["group_dropped"] == 1 + assert emitted[0][1]["granularity"] == "token" + assert oc.GATE_DROPPED_LENGTHS not in emitted[1][1] + assert oc.TRUNCATED_BY_TURN not in emitted[1][1] + + + +@pytest.mark.parametrize("backend, label", [("verl", "verl"), ("miles", "miles"), + (None, "unknown")]) +def test_driver_labels_offpolicy_gate_with_the_batch_backend(backend, label): + from yeto.rl.engine.driver import IslandDriver + from yeto.rl.engine.ports import RolloutBatchHandle + + emitted = [] + fake = SimpleNamespace(emit=lambda name, **kw: emitted.append((name, kw)), + clock=lambda: 0.0, _labels=lambda: {}) + batch = RolloutBatchHandle( + rollout_id=1, policy_version=1, policy_hash="h", groups=(), completed=0, aborted=0, + carry_over={oc.GATE_DROPPED_LENGTHS: {}, oc.GATE_GRANULARITY: oc.GRANULARITY_TRAJECTORY}, + backend=backend) + IslandDriver._emit_carry_over(fake, 1, batch) + assert emitted[0][0] == oc.OFFPOLICY_GATE_EVENT + assert emitted[0][1]["backend"] == label + assert emitted[0][1]["granularity"] == oc.GRANULARITY_TRAJECTORY + +# -- L6.9 A_hard >= 2 ---------------------------------------------------------- + + +def test_a_hard_two_passes_support_and_maps_to_both_backends(): + from yeto.rl.adapters.miles.policy_age import SUPPORT as MILES, policy_age_argv + from yeto.rl.adapters.verl.policy_age import SUPPORT as VERL, policy_age_overrides + + MILES.check(2) + VERL.check(2) + assert policy_age_argv(2, overlap=True) == ("--fully-async", "--max-weight-staleness", "2") + assert "async_training.staleness_threshold=2.0" in policy_age_overrides(2, groups_per_round=4) + with pytest.raises(PolicyAgeError): + MILES.check(MILES.max_policy_age + 1) + + +def test_miles_admission_mapping(): + assert mo.admission_argv(0) == () + assert mo.admission_argv(2) == ("--max-weight-staleness", "2") + assert mo.admission_argv(2, admission_limit=1) == ( + "--max-weight-staleness", "2", "--async-data-buffer-capacity-factor", "1") + with pytest.raises(ValueError, match="k=3"): + mo.admission_argv(2, admission_limit=3) + + +# -- L6.12 start version -------------------------------------------------------- + + +def test_continued_trajectory_gap_from_start_version(): + provenance = TokenProvenance((3, 3), (-0.1, -0.1)).extend(4, [-0.2]).extend(5, [-0.3, -0.3]) + assert oc.start_version_gap(6, provenance.versions) == 3 + assert oc.oldest_is_start(provenance.versions) is True + assert 6 - provenance.oldest == oc.start_version_gap(6, provenance.versions) + + +def test_tape_segments_check(): + records = [{"policy_versions": [[3, 0, 2], [4, 2, 3], [5, 3, 5]]}, + {"policy_versions": [[5, 0, 2], [4, 2, 3]]}, # not a continuation order + {"reward": 1.0}] + assert oc.segments_start_check(records) == { + "checked": 2, "oldest_is_start": 1, "oldest_before_start": 1} + + +# -- L6.13 cross_version_is: observe-only, from the spec --------------------------- + + +def test_cross_version_bounds_follow_the_spec(): + assert oc.cross_version_bounds(decoupled(5.0).correction) == (0.0, 5.0) + icepop = CorrectionSpec(method="tis", lower=0.5, upper=5.0, ratio_mask="icepop") + assert oc.cross_version_bounds(icepop) == (0.5, 5.0) + assert oc.cross_version_bounds(CorrectionSpec()) == oc.OBSERVE_DEFAULT_BOUNDS + + +def test_staleness_policy_correction_is_derived(): + assert StalenessPolicy.for_algorithm(decoupled()).correction == "tis" + icepop = AlgorithmSpec(correction=CorrectionSpec(method="tis", lower=0.5, upper=5.0, + ratio_mask="icepop")) + assert StalenessPolicy.for_algorithm(icepop).correction == "icepop" + assert StalenessPolicy.for_algorithm(AlgorithmSpec()).correction == "tis" + + +def test_miles_rollout_handle_keeps_the_gate_inputs(): + # L6.10 GPU: the Miles handle dropped the gate keys, so the driver taped no offpolicy_gate + from yeto.rl.adapters.miles.rollout import _carry_over_fields + + payload = {"max_policy_age": 1, "over_age_discarded_groups": 1, + oc.GATE_DROPPED_LENGTHS: {"group_dropped": [300], "unknown_version_dropped": []}, + oc.TRUNCATED_BY_TURN: {"0": [1, 2]}} + fields = _carry_over_fields(payload) + gate = oc.gate_from_carry_over(fields, kept=3, backend="miles") + assert gate["counts"] == {"group_dropped": 1, "kept": 3, "unknown_version_dropped": 0} + assert gate["truncation_by_turn"]["0"]["truncated"] == 1 + assert oc.GATE_DROPPED_LENGTHS not in fields + + +@pytest.mark.parametrize("stale, expected", [(2, 2), (0, 0), (None, None)]) +def test_miles_buffer_stale_drops_count_apart_from_the_gate(stale, expected): + # GOAL §6: Miles' own --fully-async buffer drops get their own reason code + from yeto.rl.adapters.miles.rollout import _carry_over_fields + + payload = {"max_policy_age": 1, "over_age_discarded_groups": 1, "stale_dropped_groups": stale, + oc.GATE_DROPPED_LENGTHS: {"group_dropped": [300], "unknown_version_dropped": []}} + gate = oc.gate_from_carry_over(_carry_over_fields(payload), kept=3, backend="miles") + assert gate["counts"]["group_dropped"] == 1 + assert gate["counts"].get(oc.GATE_BACKEND_STALE_DROPPED) == expected + assert oc.GATE_BACKEND_STALE_DROPPED not in gate["dropped_length_histogram"] diff --git a/tests/test_rl_policy_age.py b/tests/test_rl_policy_age.py index 9e1b3f780..3f82926f6 100644 --- a/tests/test_rl_policy_age.py +++ b/tests/test_rl_policy_age.py @@ -57,7 +57,7 @@ def test_elastic_join_carries_the_bound_identity(monkeypatch): assert a != b -@pytest.mark.parametrize("backend,limit,stage", [("miles", 2, 3), ("verl", 2, 2)]) +@pytest.mark.parametrize("backend,limit,stage", [("miles", 5, 3), ("verl", 5, 2)]) def test_launcher_refuses_an_unsupported_limit_before_launch(backend, limit, stage): from tests.test_rl_launcher import _args from yeto.launcher import _prepare_rl_args diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index dfd7931d4..15585c2c3 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -70,7 +70,7 @@ def test_build_inputs_match_the_pins(): def test_commits_and_digest_are_pinned_together_with_the_build_record(): """Commit pins, tag, Dockerfile ARGs and the image digest move together: the latest build record must name exactly the pinned commits and digest.""" - record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801" + record_dir = REPO / "openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-44d41ef-2fa8801" record = json.loads((record_dir / "build-record.json").read_text()) assert record["miles_commit"] == rl.MILES_NEXT_COMMIT assert record["sglang_commit"] == rl.SGLANG_NEXT_COMMIT @@ -113,28 +113,32 @@ def test_ports_setup_reuses_the_image_forks_and_still_verifies(): assert f'm["commit"] == "{c}"' in miles reuse, fallback = miles[decision:].split("\nelse\n", 1) assert f"[yeto-setup] image provides miles {c}" in reuse - assert f'peft.__version__ == "{rl.MILES_PEFT_VERSION}"' in reuse + assert f"{rl.MILES_NEXT_IMAGE_MANIFEST} peft version {rl.MILES_PEFT_VERSION}" in reuse + assert f'peft.__version__ == "{rl.MILES_PEFT_VERSION}"' in reuse # fallback only assert f"pip install -q 'peft=={rl.MILES_PEFT_VERSION}'" in reuse assert "pip install -q --no-deps -e ~/miles" not in reuse assert fallback == ( f"python3 -m pip install -q --no-deps -e ~/miles 'peft=={rl.MILES_PEFT_VERSION}'\nfi" ) # SGLang: reuse /sgl-workspace/sglang only if it is exactly the pinned, - # clean fork checkout that `import sglang` resolves to; else clone. + # clean fork checkout and the image manifest names it there (no import: + # L2.3); without a manifest entry the import check is the fallback. s = rl.SGLANG_NEXT_COMMIT + assert sglang.startswith("SGLANG_IN_IMAGE=0\nif ") guard, rest = sglang.split("; then\n", 1) - assert guard.startswith("if ") for needle in ( f'[ "$(git -C /sgl-workspace/sglang rev-parse HEAD 2>/dev/null)" = {s} ]', f'= {rl.SGLANG_NEXT_REPOSITORY} ]', "git -C /sgl-workspace/sglang status --porcelain --untracked-files=all", - "import os, sys, sglang", "[ ! -e ~/sglang ]", ): assert needle in guard + check, rest = rest.split('\nfi\nif [ "$SGLANG_IN_IMAGE" = 1 ]; then\n', 1) + assert f"{rl.MILES_NEXT_IMAGE_MANIFEST} sglang commit {s} /sgl-workspace/sglang" in check + match, absent = check.split("\n0) ", 1)[1].split("\n3)\n", 1) + assert "import sglang" not in match and "SGLANG_IN_IMAGE=1" in match + assert "import os, sys, sglang" in absent and "image_manifest_fallback" in absent reuse, fallback = rest.split("\nelse\n", 1) - assert "ln -sfn /sgl-workspace/sglang ~/sglang" in reuse - assert f"git -C ~/sglang fetch --depth 1 origin {s}" in fallback assert fallback.endswith("pip install -q --no-deps -e ~/sglang/python\nfi") @@ -261,6 +265,74 @@ def test_pure_python_check_rejects_build_inputs(tmp_path): ml.check_pure_python(repo, base, head, pkgs) +def test_layer_drops_megatron_bridges_eager_modelopt_import(): + """L2.2: the trainer imports megatron.bridge without paying for modelopt (21-30 s).""" + import ast + import importlib.util + + spec = importlib.util.spec_from_file_location("make_layer", REPO / "docker/miles-ports/make_layer.py") + ml = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ml) + tree = ast.parse(ml.MEGATRON_BRIDGE_INIT.read_text()) + imported = set() + for node in ast.walk(tree): # module level and inside try/if: all run at import time + if isinstance(node, ast.Import): + imported |= {a.name for a in node.names} + elif isinstance(node, ast.ImportFrom): + imported.add(node.module) + assert not any(name.split(".")[0] == "modelopt" for name in imported) + # the rest of the package init is untouched: bridges still register + assert {"megatron.bridge.models", "megatron.bridge.models.conversion.auto_bridge"} <= imported + script = (REPO / "scripts/build_miles_ports_image.sh").read_text() + assert ml.MEGATRON_BRIDGE_COMMIT in script and "lazy-modelopt" in script + + +def _make_layer_module(): + import importlib.util + + spec = importlib.util.spec_from_file_location("make_layer", REPO / "docker/miles-ports/make_layer.py") + ml = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ml) + return ml + + +def test_patched_files_applies_a_patch_to_the_commits_files(tmp_path): + ml = _make_layer_module() + repo = tmp_path / "repo" + run = lambda *a: subprocess.run(["git", "-C", str(repo), *a], check=True, capture_output=True, text=True) + (repo / "pkg").mkdir(parents=True) + run("init", "-q") + (repo / "pkg/a.py").write_text("import heavy\n\n\ndef f():\n return heavy.x\n") + (repo / "pkg/b.py").write_text("untouched\n") + run("add", "-A") + run("-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "-m", "base") + head = run("rev-parse", "HEAD").stdout.strip() + (repo / "pkg/a.py").write_text("def f():\n import heavy\n\n return heavy.x\n") + patch = tmp_path / "lazy.patch" + patch.write_text(run("diff").stdout) + (repo / "pkg/a.py").write_text("worktree edits are ignored\n") + + assert ml.patched_files(repo, head, patch) == {"pkg/a.py": b"def f():\n import heavy\n\n return heavy.x\n"} + + +def test_megatron_lm_patch_drops_the_module_level_modelopt_imports(): + """L2.2: the three Megatron-LM imports a trainer process ran (traced in the image) become lazy.""" + ml = _make_layer_module() + patch = ml.MEGATRON_LM_PATCH.read_text() + files = [line[len("--- a/"):] for line in patch.splitlines() if line.startswith("--- a/")] + assert files == ["megatron/post_training/checkpointing.py", "megatron/training/checkpointing.py", + "megatron/training/training.py"] + removed = {line[1:] for line in patch.splitlines() if line.startswith("-") and not line.startswith("---")} + assert {"import modelopt", " from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state", + " from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation", + } <= removed + added = [line[1:] for line in patch.splitlines() if line.startswith("+") and not line.startswith("+++")] + # every modelopt import the patch adds sits in a function body (indented) + assert all(line.startswith(" ") for line in added if re.match(r"\s*(import|from) modelopt", line)) + script = (REPO / "scripts/build_miles_ports_image.sh").read_text() + assert ml.MEGATRON_LM_COMMIT in script and "megatron_lm_lazy_modelopt.patch" in script + + def test_gpu_exact_needs_modal_islands(): from yeto.gpu_spec import parse_gpu_spec diff --git a/tests/test_rl_recovery_abstraction.py b/tests/test_rl_recovery_abstraction.py index b70736b66..6064e8e1a 100644 --- a/tests/test_rl_recovery_abstraction.py +++ b/tests/test_rl_recovery_abstraction.py @@ -242,3 +242,48 @@ def test_inflight_saver_commit_semantics_per_backend(tmp_path, monkeypatch): out = modal.rehearse(1) assert out["commit_error"] is None and out["commit_s"] is not None assert log.read_text().splitlines() == ["commit tape-vol"] + + +# ---------------------------------------------------------------- rl-trainer-start-stop L3.2 +def test_scheduled_cut_saves_and_resumes_like_a_manual_cut(tmp_path): + """D3: a stop the trainer scheduler orders writes a SCHEDULED cut; the next + start restores it the same way as a MANUAL cut.""" + assert rc.SCHEDULED in rc.TRIGGERS and rc.SCHEDULED in rc.FINAL_TRIGGERS + store = _store(tmp_path / "store") + rank_a, ctl_a, driver_a, cp_a = _trained_island(tmp_path, store, 2) + driver_a.sync = SimpleNamespace(outer_version=2) + with pytest.raises(Exception, match="final=True"): + cp_a.save(driver_a, rollout_id=2, trigger=rc.SCHEDULED) + with pytest.raises(Exception, match="final=True"): + cp_a.save(driver_a, rollout_id=2, final=True, trigger=rc.PERIODIC) + saved = cp_a.save(driver_a, rollout_id=2, final=True, trigger=rc.SCHEDULED) + cut = rc.ConsistentCut.from_record(saved[rc.CUT_KEY]) + assert cut.trigger == rc.SCHEDULED and saved["final"] is True + + rank_b = make_rank(9) + ctl_b = _rctl(tmp_path, "b", store) + driver_b = _Driver(rank_b) + got = _cp(ctl_b, driver_b, _trainer(rank_b), _Ledger(2)).resume(driver_b) + assert got["next_rollout_id"] == 2 and driver_b.local_step == cut.island_steps[ISLAND] == 2 + assert rc.ConsistentCut.from_record(got[rc.CUT_KEY]).trigger == rc.SCHEDULED + for n, v in params(rank_a).items(): + assert torch.equal(v, params(rank_b)[n]) + + +def test_local_only_sync_passes_the_final_trigger_to_the_round_cut(): + from yeto.rl.engine.bridges import LocalOnlySync + + calls = [] + + class Cuts: + def save(self, driver, **kw): + calls.append(kw) + + sync = LocalOnlySync(3) + sync.round_cuts = Cuts() + driver = SimpleNamespace(published_version=1, at_safe_point=False, emit=lambda *a, **k: None) + sync.finish(driver) + sync.final_trigger = rc.SCHEDULED + sync.finish(driver) + assert calls == [{"rollout_id": 1, "final": True}, + {"rollout_id": 1, "final": True, "trigger": rc.SCHEDULED}] diff --git a/tests/test_rl_sample_buffer.py b/tests/test_rl_sample_buffer.py new file mode 100644 index 000000000..0af088342 --- /dev/null +++ b/tests/test_rl_sample_buffer.py @@ -0,0 +1,213 @@ +"""rl-trainer-start-stop L4 (CPU, fakes): the durable sample buffer, its watermark +in the cut, arrival polling, and samplers that follow the trainer.""" + +from __future__ import annotations + +import subprocess +import sys +import threading +import time +from pathlib import Path + +import pytest +import torch +from test_rl_engine_driver import NAME, PROFILE, _driver + +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.fake_lifecycle import ( + BufferFedFakeRollout, + FakeMembership, + FakeRoundCuts, + FakeSampler, + cuttable_fake_engine, +) +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer, SampleBufferError +from yeto.rl.engine.trainer_lifecycle import ( + DriverLifecycle, + DurableSampleBuffer, + LifecycleError, + LifecycleState, + SampleIndexEntry, + StopCondition, +) + +REPO = Path(__file__).resolve().parents[1] +BATCH = 4 # fake engine: 2 groups x 2 samples + + +def _buffer(root, **kw): + return FileSampleBuffer(CheckpointStore(root), batch_size=BATCH, **kw) + + +def _add(buf, sample_id, version): + ref = buf.write_payload(sample_id, sample_id.encode()) + buf.append(SampleIndexEntry(sample_id, version, ref, tokens=3)) + + +# ---------------------------------------------------------------- L4.1 index +def test_append_select_commit_and_backlog(tmp_path): + buf = _buffer(tmp_path) + assert isinstance(buf, DurableSampleBuffer) + for i, v in enumerate([0, 0, 1, 1, 1, 2]): + _add(buf, f"s{i}", v) + stats = buf.backlog(current_version=3) + assert (stats.samples, stats.batches, stats.oldest_version, stats.max_version_gap) == (6, 1.5, 0, 3) + assert stats.version_histogram == {0: 2, 1: 3, 2: 1} and stats.arrival_rate > 0 + got = buf.select("trainer", 4) + assert [s.position for s in got] == [0, 1, 2, 3] and got[2].entry.sample_id == "s2" + assert buf.read_payload(got[0].entry.payload_ref) == b"s0" + assert buf.backlog(3).samples == 2 and buf.watermark == 0 + buf.commit_consumed(4) + assert buf.watermark == 4 + with pytest.raises(SampleBufferError): + buf.commit_consumed(5) # beyond what was selected + with pytest.raises(SampleBufferError): + buf.select("other", 1) + with pytest.raises(SampleBufferError, match="payload"): + buf.append(SampleIndexEntry("x", 0, "payloads/missing.bin")) + + +def test_reopen_after_process_exit(tmp_path): + """Samples a sampler process appended survive that process; a torn last line + is ignored and cut off by the next append.""" + script = ( + "import sys\n" + "from yeto.rl.engine.resume import CheckpointStore\n" + "from yeto.rl.engine.sample_buffer import FileSampleBuffer\n" + "from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry\n" + "buf = FileSampleBuffer(CheckpointStore(sys.argv[1]), batch_size=4)\n" + "for i in range(5):\n" + " ref = buf.write_payload(f'p{i}', b'x' * i)\n" + " buf.append(SampleIndexEntry(f'p{i}', i // 2, ref))\n" + ) + subprocess.run([sys.executable, "-c", script, str(tmp_path)], check=True, cwd=REPO) + index = tmp_path / "sample-buffer" / "index.jsonl" + with open(index, "ab") as fh: + fh.write(b'{"schema": "yeto.rl.sample-buf') # crash in the middle of a line + buf = _buffer(tmp_path) + assert len(buf) == 5 and buf.read_payload(buf.select("t", 5)[4].entry.payload_ref) == b"xxxx" + _add(buf, "p5", 3) + again = _buffer(tmp_path) + assert len(again) == 6 and [s.entry.sample_id for s in again.select("t", 9)][-1] == "p5" + again.mark_data_end() + assert _buffer(tmp_path).backlog(0).data_ended + with pytest.raises(SampleBufferError, match="data end"): + _add(again, "late", 3) + + +# ---------------------------------------------------------------- L4.3 arrival +def test_arrival_callback_fires_within_one_interval(tmp_path): + reader = _buffer(tmp_path) + writer = _buffer(tmp_path) # another process in production + fired = threading.Event() + reader.on_arrival(fired.set) + interval = 0.5 + reader.start_polling(interval) + try: + assert not fired.wait(interval * 1.5) # nothing appended yet + appended = time.monotonic() + _add(writer, "a0", 0) + assert fired.wait(interval * 3) + # one interval plus scheduling slack + assert time.monotonic() - appended <= interval + 0.25 + finally: + reader.stop_polling() + assert reader.poll() == 0 # the polled sample is not reported twice + + +# ---------------------------------------------------------------- L4.2 / L4.4 lifecycle +class _Island: + def __init__(self, tmp_path, *, fail_rounds=None): + self.tmp_path = tmp_path + self.buffer = _buffer(tmp_path / "store") + self.cuts = FakeRoundCuts(tmp_path / "cuts", watermark=lambda: self.buffer.watermark) + self.membership = FakeMembership() + self.fail_rounds = dict(fail_rounds or {}) # incarnation -> rounds whose step fails + self.rollouts = [] + self.engine = None + + def build(self, handle): + engine = cuttable_fake_engine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0, + capability_profile=PROFILE, + failed_step_rounds=set(self.fail_rounds.get(handle.incarnation, ()))) + resume = handle.resume_from + self.buffer.reset_to(0 if resume is None else resume.watermark) + engine.rollout = BufferFedFakeRollout(engine, self.buffer) + self.rollouts.append(engine.rollout) + self.engine = engine + sync = LocalOnlySync(1000) + sync.round_cuts = self.cuts + return _driver(engine, sync, self.tmp_path, name=f"events-{handle.incarnation}.jsonl") + + def lifecycle(self): + return DriverLifecycle(self.build, membership=self.membership, cuts=self.cuts) + + def published(self): + return None if self.engine is None or self.engine.published is None else self.engine.published[0] + + +def test_crash_after_consuming_trains_the_samples_again(tmp_path): + """Spec "Trainer crash after consuming": the trainer consumes samples and fails + before the next cut; the next start resumes at the watermark of the last cut.""" + island = _Island(tmp_path, fail_rounds={1: {2}}) + sampler = FakeSampler(island.buffer, island.published) + lc = island.lifecycle() + sampler.produce(3 * BATCH) + + h = lc.start(lc.plan_start({}), resume_from=None) + lc.run(h, island.buffer, StopCondition(2)) + cut = lc.drain_and_cut(h, LeaveReason.STOPPED) + lc.stop(h, LeaveReason.STOPPED) + assert (cut.trigger, cut.policy_version, cut.watermark) == (recovery.SCHEDULED, 2, 2 * BATCH) + record = island.cuts.pointer()[recovery.CUT_KEY] + assert record["buffer_watermark"] == 2 * BATCH + + h = lc.start(lc.plan_start({}), resume_from=cut) + with pytest.raises(LifecycleError): + lc.run(h, island.buffer, StopCondition(1)) # round 2 consumes, then its step fails + assert h.state is LifecycleState.STOPPED and island.cuts.latest() == cut + crashed = [sid for sid, _ in island.rollouts[1].trained[2]] + + h = lc.start(lc.plan_start({}), resume_from=island.cuts.latest()) + lc.run(h, island.buffer, StopCondition(1)) + lc.stop(h, LeaveReason.STOPPED) + again = [sid for sid, _ in island.rollouts[2].trained[2]] + assert again == crashed == [f"s0-{i:05d}" for i in range(2 * BATCH, 3 * BATCH)] + first = [sid for r in (0, 1) for sid, _ in island.rollouts[0].trained[r]] + assert not set(first) & set(again) # nothing before the watermark trained again + assert island.cuts.latest().watermark == 3 * BATCH and island.cuts.latest().policy_version == 3 + assert [t[0] for t in island.membership.tape] == ["JOIN", "LEAVE"] * 3 + assert island.membership.tape[3] == ("LEAVE", 1, LeaveReason.TRAINING_FAILED) + + +def test_samplers_follow_the_trainer_and_keep_the_last_version(tmp_path): + """D5: while training, each published version reaches the sampler; samples + record the version the sampler loaded; after a stop the sampler keeps it.""" + island = _Island(tmp_path) + sampler = FakeSampler(island.buffer, island.published) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + assert lc.ready(h, time.monotonic() + 30) + for _ in range(3): + sampler.load() + sampler.produce(BATCH) + receipt = lc.run(h, island.buffer, StopCondition(1)) + assert h.published_version == receipt.trained_version # published before the wait + assert sampler.load() == 3 + lc.stop(h, LeaveReason.STOPPED) + island.engine = None # no trainer: nothing new to load + assert sampler.load() == 3 + sampler.produce(BATCH) + versions = [e.behavior_version for e in (s.entry for s in _buffer(tmp_path / "store").select("x", 99))] + assert versions == [0] * BATCH + [1] * BATCH + [2] * BATCH + [3] * BATCH + trained = island.rollouts[0].trained + assert [sorted({v for _, v in trained[r]}) for r in range(3)] == [[0], [1], [2]] + # the next start trains the backlog the sampler wrote while the trainer was off + h = lc.start(lc.plan_start({}), resume_from=island.cuts.latest()) + lc.run(h, island.buffer, StopCondition(1)) + lc.stop(h, LeaveReason.STOPPED) + assert {v for _, v in island.rollouts[1].trained[3]} == {3} + assert island.buffer.backlog(4).samples == 0 and island.cuts.latest().watermark == 4 * BATCH diff --git a/tests/test_rl_sampler_island.py b/tests/test_rl_sampler_island.py new file mode 100644 index 000000000..315047604 --- /dev/null +++ b/tests/test_rl_sampler_island.py @@ -0,0 +1,483 @@ +"""rl-trainer-start-stop L4.6 (CPU, no Ray, no Miles): sampler / trainer islands over +the durable sample buffer -- sample roles and launch flags, policy releases through +the store, the sampler loop (release loading, backlog cap, trainer stopped), the +engine weight load order, the releasing publisher and the group renumbering.""" + +from __future__ import annotations + +import asyncio +import dataclasses +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tests.buffer_rollout_procs import sample +from yeto.rl.adapters.miles import buffer_rollout as br +from yeto.rl.adapters.miles import sampler_island as si +from yeto.rl.engine import sample_roles +from yeto.rl.engine.policy_release import ( + LATEST, + RELEASE_DIR, + PolicyFollower, + PolicyReleaseStore, +) +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.sampler_loop import SamplerLoop + + +def _group(g, version): + return [sample(2 * g, g, version), sample(2 * g + 1, g, version)] + + +def _write_weights(directory: Path) -> None: + (directory / "model.safetensors").write_text("w") + + +# ---------------------------------------------------------------- roles and launch +def _launch_args(**kw): + base = dict(rl_sample_role="sampler", rl_single_island_no_sync=True, + rl_checkpoint_store="s3://b/run", rl_elastic=False, rl_max_policy_age=0) + return SimpleNamespace(**{**base, **kw}) + + +def test_role_both_is_the_default_and_adds_nothing(): + args = SimpleNamespace() + assert sample_roles.sample_role(args) == "both" + sample_roles.check_launch(args) + assert sample_roles.island_flags(args) == "" + assert sample_roles.resolve(args, groups_per_round=4) is None + with pytest.raises(ValueError, match="needs --rl-sample-role sampler or trainer"): + sample_roles.check_launch(SimpleNamespace(rl_sample_wait_s=5.0)) + + +@pytest.mark.parametrize("kw, message", [ + ({"rl_single_island_no_sync": False}, "--rl-single-island-no-sync"), + ({"rl_checkpoint_store": None}, "--rl-checkpoint-store"), + ({"rl_elastic": True}, "--rl-elastic"), + ({"rl_sample_role": "trainer"}, "--rl-max-policy-age >= 1"), + ({"rl_sampler_max_backlog_rounds": 0}, "a value >= 1"), + ({"rl_sample_wait_s": 3.0}, "--rl-sample-role trainer"), + ({"rl_sampler_min_round_seconds": 0.0}, "a value > 0"), + ({"rl_sample_role": "trainer", "rl_max_policy_age": 1, "rl_sampler_min_round_seconds": 5.0}, + "--rl-sampler-min-round-seconds needs --rl-sample-role sampler"), + ({"rl_sample_role": "x"}, "expected one of"), +]) +def test_launch_refuses_unsupported_role_combinations(kw, message): + with pytest.raises(ValueError, match=message): + sample_roles.check_launch(_launch_args(**kw)) + + +def test_launch_forwards_the_role_flags_and_the_island_resolves_them(): + from yeto import launcher + + args = _launch_args(rl_sampler_max_backlog_rounds=2) + sample_roles.check_launch(args) + _, flags = launcher._ports_infra_flags(args) + assert " --rl-sample-role sampler --rl-sampler-max-backlog-rounds 2" in flags + _, default = launcher._ports_infra_flags(SimpleNamespace()) + assert "--rl-sample-role" not in default + island = SimpleNamespace(rl_sample_role="sampler", rl_resume_store="/mnt/store", + rl_sampler_max_backlog_rounds=2) + assert sample_roles.resolve(island, groups_per_round=8) == sample_roles.SampleBufferConfig( + role="sampler", store="/mnt/store", max_backlog_groups=16) + trainer = SimpleNamespace(rl_sample_role="trainer", rl_resume_store="/mnt/store", + rl_sample_wait_s=30.0) + assert sample_roles.resolve(trainer, groups_per_round=8).wait_s == 30.0 + with pytest.raises(ValueError, match="--rl-resume-store"): + sample_roles.resolve(SimpleNamespace(rl_sample_role="trainer"), groups_per_round=8) + + +def test_sampler_min_round_seconds_is_forwarded_and_resolved_for_the_sampler_only(): + args = _launch_args(rl_sampler_min_round_seconds=240.0) + sample_roles.check_launch(args) + assert " --rl-sampler-min-round-seconds 240.0" in sample_roles.island_flags(args) + island = SimpleNamespace(rl_sample_role="sampler", rl_resume_store="/mnt/store", + rl_sampler_min_round_seconds=240.0) + assert sample_roles.resolve(island, groups_per_round=8).min_round_s == 240.0 + pair = SimpleNamespace(rl_island_sample_role=["0:sampler", "1:trainer"], + rl_sampler_min_round_seconds=240.0, rl_sample_wait_s=9.0, + rl_sampler_max_backlog_rounds=None) + assert sample_roles.island_args(pair, 1, 2).rl_sampler_min_round_seconds is None + assert sample_roles.island_args(pair, 0, 2).rl_sampler_min_round_seconds == 240.0 + + +def test_role_flags_exist_on_both_layers(): + import argparse + + from yeto.rl.engine.run_flags import CLI, ISLAND, flags_for + + for layer in (CLI, ISLAND): + parser = argparse.ArgumentParser() + for flag in flags_for(layer): + if flag.options[0] in ("--rl-sample-role", "--rl-sample-wait-s", + "--rl-sampler-max-backlog-rounds"): + parser.add_argument(*flag.options, **dict(flag.kwargs)) + ns = parser.parse_args(["--rl-sample-role", "trainer", "--rl-sample-wait-s", "7.5"]) + assert (ns.rl_sample_role, ns.rl_sample_wait_s, ns.rl_sampler_max_backlog_rounds) == ( + "trainer", 7.5, None) + assert parser.parse_args([]).rl_sample_role == "both" + + +def test_apply_sample_role_sets_the_buffer_config_and_the_island_environment(): + from yeto.rl.adapters.miles.island_entry import apply_sample_role + + miles_args, environ = SimpleNamespace(), {} + apply_sample_role(SimpleNamespace(groups_per_round=4, learner_id=0), miles_args, environ) + assert environ == {} and not hasattr(miles_args, "yeto_rl_sample_buffer") + args = SimpleNamespace(groups_per_round=4, learner_id=3, rl_sample_role="sampler", + rl_resume_store="/mnt/store") + apply_sample_role(args, miles_args, environ) + assert miles_args.yeto_rl_sample_buffer.max_backlog_groups == 16 + assert environ == {br.BUFFER_ROOT_ENV: "/mnt/store", br.SAMPLER_ID_ENV: "island3"} + + +def test_ray_workers_get_the_buffer_environment(): + from yeto.rl.adapters.miles.entry import connect_island_ray + + seen = {} + ray = SimpleNamespace(init=lambda **kw: seen.update(kw), is_initialized=lambda: False) + env = {br.BUFFER_ROOT_ENV: "/mnt/store", br.MAX_POLICY_AGE_ENV: "2"} + connect_island_ray(environ={"RAY_ADDRESS": "10.0.0.1:6379", **env}, ray_module=ray) + workers = seen["runtime_env"]["env_vars"] + assert {k: workers[k] for k in env} == env + assert br.SAMPLER_ID_ENV not in workers + + +# ---------------------------------------------------------------- Miles argv +def _full_config(role, **kw): + from tests.test_rl_miles_adapter_config import make_config, sub + + cfg = sub(make_config(), "trainable", parameter_mode="full", target_modules=()) + return dataclasses.replace(cfg, sample_buffer=sample_roles.SampleBufferConfig( + role=role, store="/mnt/store", **kw)) + + +def _tolerant(): + from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec + + return AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)) + + +def test_miles_argv_of_the_sampler_and_the_trainer(): + from tests.test_rl_miles_adapter_config import flag_value, make_config, sub + from yeto.rl.adapters.miles import config as mc + + base = mc.translate_run_config(sub(make_config(), "trainable", parameter_mode="full", + target_modules=()), _tolerant()) + sampler = mc.translate_run_config(_full_config("sampler", max_backlog_groups=4), _tolerant()) + assert [a for a in sampler.argv if a not in base.argv] == [ + "--rollout-function-path", br.ROLLOUT_FUNCTION_PATHS["sampler"], "--debug-rollout-only"] + trainer = mc.translate_run_config(_full_config("trainer"), _tolerant(), max_policy_age=1) + assert flag_value(trainer.argv, "--rollout-function-path") == br.ROLLOUT_FUNCTION_PATHS["trainer"] + assert flag_value(trainer.argv, "--save-hf") == mc.RELEASE_EXPORTER_SAVE_HF + assert "--debug-rollout-only" not in trainer.argv and not trainer.rollout_overlap + # partitioned + policy age: the buffer rollout replaces Miles' fully_async + partitioned = sub(_full_config("trainer"), "serving", offload_train=False) + partitioned = dataclasses.replace(partitioned, parallel=dataclasses.replace( + make_config(colocated=False).parallel)) + assert "--fully-async" not in mc.translate_run_config(partitioned, _tolerant(), max_policy_age=1).argv + + +def test_miles_argv_refuses_a_lora_buffer_island(): + from tests.test_rl_miles_adapter_config import make_config + from yeto.rl.adapters.miles import config as mc + + cfg = dataclasses.replace(make_config(), sample_buffer=sample_roles.SampleBufferConfig( + role="sampler", store="/mnt/store", max_backlog_groups=4)) + with pytest.raises(mc.MilesConfigError, match="full-parameter mode only"): + mc.translate_run_config(cfg, _tolerant()) + + +# ---------------------------------------------------------------- releases +def test_release_is_committed_by_latest_and_republish_is_a_no_op(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path), keep=2, clock=lambda: 1.0) + assert releases.latest() is None + with pytest.raises(RuntimeError, match="export failed"): + releases.release(1, "h1", lambda d: (_ for _ in ()).throw(RuntimeError("export failed"))) + assert releases.latest() is None # a failed export never becomes LATEST + first = releases.release(1, "h1", _write_weights, buffer_selected=4) + assert (first.seq, first.path, first.buffer_selected) == (1, "r1-v1", 4) + assert releases.resolve(first) == tmp_path / RELEASE_DIR / "r1-v1" + assert releases.release(1, "h1", _write_weights) is None # a restarted trainer republishes + releases.release(2, "h2", _write_weights) + rewind = releases.release(1, "h1b", _write_weights) # restored from an older cut + assert (rewind.seq, rewind.version) == (3, 1) + assert sorted(p.name for p in (tmp_path / RELEASE_DIR).iterdir()) == [LATEST, "r2-v2", "r3-v1"] + + +def test_follower_loads_each_new_release_once_and_keeps_the_version_on_failure(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + loaded, fail = [], [False] + + def load(release, directory): + if fail[0]: + raise RuntimeError("engine down") + loaded.append((release.version, directory.name)) + + follower = PolicyFollower(releases, load) + assert follower.refresh() is None and follower.version == 0 + releases.release(1, "h1", _write_weights) + assert follower.refresh().version == 1 and follower.refresh() is None + releases.release(2, "h2", _write_weights) + fail[0] = True + with pytest.raises(RuntimeError, match="engine down"): + follower.refresh() + assert follower.version == 1 # still serving v1 + fail[0] = False + follower.refresh() + assert loaded == [(1, "r1-v1"), (2, "r2-v2")] and follower.version == 2 + + +# ---------------------------------------------------------------- sampler loop +class _Sampler: + """A fake sampler island: each round writes 2 groups at the served version.""" + + def __init__(self, root, releases, max_backlog_groups=4, rounds=6): + self.events = [] + self.served = [] + self.follower = PolicyFollower(releases, lambda r, d: self.served.append(r.version)) + writer_buffer = FileSampleBuffer(CheckpointStore(root), batch_size=2) + self.writer = br.GroupWriter(writer_buffer, "island0-a") + self.sleeps = 0 + self.on_sleep = lambda: None + self.loop = SamplerLoop( + follower=self.follower, buffer=FileSampleBuffer(CheckpointStore(root), batch_size=2), + generate=self.generate, num_rollouts=rounds, max_backlog_groups=max_backlog_groups, + emit=lambda event, **f: self.events.append((event, f)), poll_s=0.0, + sleep=self.sleep, should_stop=lambda: self.sleeps > 20) + + def generate(self, rid): + for g in range(2): # Miles numbers groups from 0 in every round + self.writer.write(_group(g, self.follower.version), group_key=f"r{rid}-g{g}") + + def sleep(self, _dt): + self.sleeps += 1 + self.on_sleep() + + def names(self): + return [e for e, _ in self.events] + + +def test_sampler_loads_releases_and_waits_at_the_backlog_cap(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + sampler = _Sampler(tmp_path, releases, max_backlog_groups=4, rounds=5) + trainer_reads = iter([(1, 4), (2, 8)]) # (version, groups read) per release + + def trainer_runs(): # the trainer reads the backlog and releases a new version + version, read = next(trainer_reads) + releases.release(version, f"h{version}", _write_weights, buffer_selected=read) + + sampler.on_sleep = trainer_runs + assert sampler.loop.run() == 5 + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=2) + assert [e.behavior_version for e in buffer._entries] == [0, 0, 0, 0, 1, 1, 1, 1, 2, 2] + assert sampler.served == [1, 2] + assert sampler.names().count("rl_sampler_backlog_full") == 2 + assert buffer.data_ended and sampler.names()[-1] == "rl_sampler_data_end" + + +def test_sampler_throttle_keeps_min_round_seconds_between_round_starts(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + sampler = _Sampler(tmp_path, releases, max_backlog_groups=100, rounds=3) + now = [0.0] + slept = [] + sampler.loop.clock = lambda: now[0] + sampler.loop.min_round_s = 10.0 + sampler.loop.sleep = slept.append + + def generate(rid, inner=sampler.generate): + now[0] += 4.0 # a round takes 4 s + inner(rid) + + sampler.loop.generate = generate + assert sampler.loop.run() == 3 + assert slept == [6.0, 6.0] # not after the last round + assert sampler.names().count("rl_sampler_throttle") == 2 + + +def test_sampler_keeps_the_last_version_while_the_trainer_is_stopped(tmp_path): + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + releases.release(3, "h3", _write_weights, buffer_selected=0) + sampler = _Sampler(tmp_path, releases, max_backlog_groups=4, rounds=5) + assert sampler.loop.run() == 2 # the cap holds; stopped while waiting + buffer = FileSampleBuffer(CheckpointStore(tmp_path), batch_size=2) + assert {e.behavior_version for e in buffer._entries} == {3} and len(buffer) == 4 + assert sampler.names()[-1] == "rl_sampler_stopped" and not buffer.data_ended + + +# ---------------------------------------------------------------- engine load +class _Engine: + def __init__(self, log, name, version_after=None): + self.log, self.server_url, self.version, self.version_after = log, name, None, version_after + + async def update_weights_from_disk(self, model_path, weight_version=None): + self.log.append(("load", self.server_url, model_path, weight_version)) + + async def update_weight_version(self, version): + self.log.append(("stamp", self.server_url, version)) + self.version = self.version_after or version + + async def get_weight_version(self): + return self.version + + +class _Controller: + def __init__(self, log, engines): + self.log, self.engines = log, engines + + async def start_update_weights(self): + self.log.append(("lock",)) + return SimpleNamespace(rollout_engines=self.engines, snapshot_cell_id_to_hashes={"c0": "x"}) + + async def end_update_weights(self, snapshot_cell_id_to_hashes): + self.log.append(("end", snapshot_cell_id_to_hashes)) + + async def abort_update_weights(self): + self.log.append(("abort",)) + + +def _loader(log, engines): + executor = SimpleNamespace(set_weight_version=lambda v: log.append(("executor", v))) + return si.EngineReleaseLoader(_Controller(log, engines), executor, asyncio.run) + + +def test_engine_load_runs_under_the_update_lock_in_publisher_order(tmp_path): + log = [] + loader = _loader(log, [_Engine(log, "e0"), _Engine(log, "e1")]) + releases = PolicyReleaseStore(CheckpointStore(tmp_path)) + release = releases.release(2, "h2", _write_weights) + loader(release, releases.resolve(release)) + path = str(releases.resolve(release)) + assert log == [("lock",), ("load", "e0", path, "2"), ("load", "e1", path, "2"), + ("stamp", "e0", "2"), ("stamp", "e1", "2"), ("end", {"c0": "x"}), + ("executor", 2)] + log.clear() + loader.stamp(0) # base weights: version only + assert [step[0] for step in log] == ["lock", "stamp", "stamp", "end", "executor"] + + +def test_engine_load_aborts_the_lock_when_an_engine_disagrees(tmp_path): + log = [] + loader = _loader(log, [_Engine(log, "e0"), _Engine(log, "e1", version_after="9")]) + with pytest.raises(br.SampleBufferRolloutError, match="1/2 rollout engines"): + loader.stamp(1) + assert log[-1] == ("abort",) and ("executor", 1) not in log + with pytest.raises(br.SampleBufferRolloutError, match="no rollout engines"): + _loader([], []).stamp(1) + + +# ---------------------------------------------------------------- trainer releases +class _State: + def __init__(self, version, digest): + self.policy_version, self._digest = version, digest + + def policy_tensor_hash(self): + return self._digest + + +def test_releasing_publisher_releases_every_publish(tmp_path): + published, exported, events = [], [], [] + inner = SimpleNamespace(publish=lambda state: published.append(state.policy_version) or "result", + verify_serving_policy="inner-attr") + positions = iter([None, 6, None]) + publisher = si.ReleasingPublisher( + inner, PolicyReleaseStore(CheckpointStore(tmp_path)), + export=lambda v, d: exported.append((v, d.name)) or _write_weights(d), + read_position=lambda: next(positions), emit=lambda e, **f: events.append((e, f))) + assert publisher.publish(_State(0, "h0")) == "result" + publisher.publish(_State(1, "h1")) + publisher.publish(_State(1, "h1")) # republished after a restart: no new release + latest = PolicyReleaseStore(CheckpointStore(tmp_path)).latest() + assert published == [0, 1, 1] and exported == [(0, "r1-v0"), (1, "r2-v1")] + assert (latest.version, latest.buffer_selected) == (1, 6) + assert [f["seq"] for e, f in events if e == "rl_policy_released"] == [1, 2] + assert publisher.verify_serving_policy == "inner-attr" + + +def test_wire_policy_releases_only_on_a_trainer_island(tmp_path): + driver = SimpleNamespace(publisher=SimpleNamespace(publish=lambda s: None), + emit=lambda event, **fields: None, + rollout=SimpleNamespace(_in_executor=lambda fn, *a: fn(None, *a))) + assert si.wire_policy_releases(driver, miles_args=SimpleNamespace(), actor_model=None, + runner=None) is None + args = SimpleNamespace(yeto_rl_sample_buffer=sample_roles.SampleBufferConfig( + role="trainer", store=str(tmp_path))) + calls = [] + actor = SimpleNamespace(export_hf=lambda rollout_id, path: calls.append((rollout_id, path))) + runner = SimpleNamespace(run=lambda value: value) + wired = si.wire_policy_releases(driver, miles_args=args, actor_model=actor, runner=runner) + assert driver.publisher is wired + br._ACTIVE.pop(br.TRAINER, None) + wired.publish(_State(0, "h0")) + assert calls == [(0, str(tmp_path / RELEASE_DIR / "r1-v0"))] + + +# ---------------------------------------------------------------- group identity +def test_groups_from_restarted_samplers_get_unique_buffer_identities(tmp_path): + """Two sampler processes (a restart) both number their groups from 0; the + trainer's batch has unique group and sample indexes, the same on a re-read.""" + for sampler_id in ("island0-a", "island0-b"): + writer = br.GroupWriter(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=2), sampler_id) + for g in range(2): + writer.write(_group(g, 1), group_key=f"r0-g{g}") + reader = br.BatchReader(FileSampleBuffer(CheckpointStore(tmp_path), batch_size=4), + max_policy_age=1, wait_s=0.0, poll_s=0.0) + taken, _ = reader.take(0, 4, current_version=1) + groups = [br.renumber_group(g, position=pos, samples_per_group=2) for pos, g in taken] + assert [s["group_index"] for g in groups for s in g] == [0, 0, 1, 1, 2, 2, 3, 3] + assert [s["index"] for g in groups for s in g] == list(range(8)) + assert groups[2][0]["metadata"]["yeto_sampler_identity"] == {"group_index": 0, "index": 0} + reader.reset(0) + again, _ = reader.take(1, 4, current_version=1) + assert [pos for pos, _ in again] == [pos for pos, _ in taken] + with pytest.raises(br.SampleBufferRolloutError, match="samples per group 1"): + br.renumber_group(taken[0][1], position=0, samples_per_group=1) + + +def test_sampler_debug_rollout_only_is_not_a_placement_rewrite(): + """L5.4 GPU run s24-l54-pair-20261014e: a fixed-partition sampler island asked + for --debug-rollout-only and was refused as a rewritten placement.""" + import dataclasses + + from yeto.rl.adapters.miles import config as mc + from yeto.rl.adapters.miles.placement import PlacementRewriteError + + launch = mc.translate_run_config(_make_miles_config(), _algorithm_spec()) + ns = SimpleNamespace(use_session_server=False, partial_rollout=False, indep_dp=False, + debug_rollout_only=True, colocate=launch.placement.kind == "colocated", + actor_num_nodes=1, actor_num_gpus_per_node=launch.placement.trainer_gpus, + rollout_num_gpus=launch.placement.rollout_gpus) + with pytest.raises(PlacementRewriteError, match="debug_rollout_only=True"): + mc.validate_parsed_args(ns, launch, num_cells=lambda _a: 1) + sampler = dataclasses.replace(launch, argv=(*launch.argv, mc.SAMPLER_ONLY_FLAG)) + mc.validate_parsed_args(ns, sampler, num_cells=lambda _a: 1) + + +def _make_miles_config(): + from tests.test_rl_miles_adapter_config import make_config + + return make_config() + + +def _algorithm_spec(): + from yeto.rl.engine.algorithm import AlgorithmSpec + + return AlgorithmSpec() + + +def test_ray_metadata_sink_discard_drops_every_record(): + from yeto.rl.adapters.miles.rollout import RayMetadataSink + + class _Actor: + class take_all: + @staticmethod + def remote(): + return ["a", "b"] + + sink = RayMetadataSink.__new__(RayMetadataSink) + sink._ray = SimpleNamespace(get=lambda x: x) + sink._actor = _Actor() + assert sink.discard() == 2 diff --git a/tests/test_rl_sampler_snapshot.py b/tests/test_rl_sampler_snapshot.py new file mode 100644 index 000000000..4ab5bff0e --- /dev/null +++ b/tests/test_rl_sampler_snapshot.py @@ -0,0 +1,106 @@ +"""CPU tests for the gated sampler snapshot mode (rl-trainer-cold-start-cost L8.2).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from yeto.rl.adapters.miles import config as mc +from yeto.rl.adapters.miles import sampler_snapshot as ss +from yeto.rl.adapters.miles.pins import SGLANG_NEXT_COMMIT +from yeto.rl.engine.algorithm import AlgorithmSpec + +from test_rl_miles_adapter_config import make_config, sub + +DISTRIBUTED_PASS = ss.CompatRecord( + sglang_commit=SGLANG_NEXT_COMMIT, update_path=ss.UPDATE_PATH_DISTRIBUTED, + result="pass", check="G-L8.2", evidence="test", +) + + +def snapshot_config(*, colocated: bool = False): + cfg = make_config(colocated=colocated) + return sub(cfg, "serving", offload_train=False, sampler_snapshot=True) + + +def test_shipped_records_hold_the_l8_1_tensor_pass(): + record = ss.require_recorded_pass(ss.UPDATE_PATH_TENSOR) + assert record.sglang_commit == SGLANG_NEXT_COMMIT + assert record.check == "G-L8.1" + + +def test_refused_without_recorded_pass_names_the_missing_check(): + with pytest.raises(ss.SamplerSnapshotRefused) as exc: + ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED, records=()) + message = str(exc.value) + assert "no recorded pass" in message + assert "snapshot compatibility check" in message + assert SGLANG_NEXT_COMMIT[:9] in message and "'distributed'" in message + + +def test_record_for_another_commit_does_not_count(): + other = ss.CompatRecord("0" * 40, ss.UPDATE_PATH_DISTRIBUTED, "pass", "G-L8.2", "") + with pytest.raises(ss.SamplerSnapshotRefused, match="no recorded pass"): + ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED, records=(other,)) + + +def test_recorded_fail_is_refused(): + failed = ss.CompatRecord(SGLANG_NEXT_COMMIT, ss.UPDATE_PATH_DISTRIBUTED, "fail", "G-L8.2", "") + with pytest.raises(ss.SamplerSnapshotRefused, match="recorded 'fail'"): + ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED, records=(failed,)) + + +def test_default_off_adds_no_sglang_config(): + argv = mc.translate_run_config( + sub(make_config(colocated=False), "serving", offload_train=False), AlgorithmSpec() + ).argv + assert "--sglang-config" not in argv + + +def test_translation_refuses_without_record(monkeypatch): + monkeypatch.setattr(ss, "load_compat_records", lambda text=None: ()) + with pytest.raises(mc.MilesConfigError, match="no recorded pass"): + mc.translate_run_config(snapshot_config(), AlgorithmSpec()) + + +def test_translation_with_record_sets_sampler_overrides_only(monkeypatch): + monkeypatch.setattr(ss, "load_compat_records", lambda text=None: (DISTRIBUTED_PASS,)) + launch = mc.translate_run_config(snapshot_config(), AlgorithmSpec()) + argv = list(launch.argv) + value = argv[argv.index("--sglang-config") + 1] + decoded = ss.decode_sglang_config_arg(value) + (model,) = decoded["sglang"] + (group,) = model["server_groups"] + assert model["name"] == "default" + assert group == {"worker_type": "regular", "num_gpus": launch.placement.rollout_gpus, + "overrides": {"enable_memory_saver": True, "enable_weights_cpu_backup": True}} + assert "--offload-rollout" not in argv # the trainer side is unchanged + + +def test_translation_refuses_colocated(monkeypatch): + monkeypatch.setattr(ss, "load_compat_records", lambda text=None: (DISTRIBUTED_PASS,)) + with pytest.raises(mc.MilesConfigError, match="dedicated sampler engines"): + mc.translate_run_config(snapshot_config(colocated=True), AlgorithmSpec()) + + +def test_translation_refuses_extra_sglang_config(monkeypatch): + monkeypatch.setattr(ss, "load_compat_records", lambda text=None: (DISTRIBUTED_PASS,)) + request = SimpleNamespace(kind="fixed-partition", placement_map_arg=None, rollout_gpus=2) + with pytest.raises(ss.SamplerSnapshotRefused, match="drop --sglang-config"): + ss.snapshot_argv(request, extra_argv=("--sglang-config", "x.yaml")) + + +def test_cli_flag_is_off_by_default_and_parses(): + from yeto import cli + + base = ["launch", "--training-mode", "rl", "--gpu", "modal:2xh100", "--model", "x", "--data", "y"] + assert cli.build_parser().parse_args(base).rl_sampler_snapshot is False + assert cli.build_parser().parse_args([*base, "--rl-sampler-snapshot"]).rl_sampler_snapshot is True + + +def test_shipped_records_hold_the_l8_2_distributed_pass(): + record = ss.require_recorded_pass(ss.UPDATE_PATH_DISTRIBUTED) + assert record.check == "G-L8.2" + argv = mc.translate_run_config(snapshot_config(), AlgorithmSpec()).argv + assert "--sglang-config" in argv diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py new file mode 100644 index 000000000..aaa291503 --- /dev/null +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -0,0 +1,192 @@ +"""L5.4 launch composition: one launch of a sampler island and a trainer island +over one checkpoint store (``--rl-island-sample-role``), the trainer island held +back for the head's scheduler. Dry-run / plan level only: no cloud, no Ray.""" + +from __future__ import annotations + +import shlex +import threading +import types + +import pytest + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 + +PAIR_GPU = "modal:1xh100,modal:1xh100" +STORE = "modal-volume://yeto-l54/run-a" +SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" +PAIR = ("--rl-single-island-no-sync", "--tuning", "full", + "--rl-island-sample-role", "0:sampler", "--rl-island-sample-role", "1:trainer", + "--rl-max-policy-age", "2", "--rl-algorithm-spec", SPEC, + "--rl-checkpoint-store", STORE, "--rl-cut-every", "1", + "--rl-sampler-max-backlog-rounds", "3", "--rl-sample-wait-s", "900", + "--rl-trainer-scheduler", "s1", "--rl-trainer-island", "1", + "--cluster-prefix", "l54-pair") + + +def _args(extra=(), gpu=PAIR_GPU, base=PAIR): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", tuple(base) + tuple(extra), gpu=gpu) + args.controller = "local" + _prepare_rl_args(args) + return args + + +def _plan(extra=()): + from yeto.launcher import dry_run_plan + + args = _args(extra) + return args, dry_run_plan(args) + + +def _island_tokens(command: str) -> list[str]: + tail = command.split("-m yeto.rl.adapters.miles.island_entry", 1)[1] + return [{"$LEARNER_ID": "0"}.get(t, t) for t in shlex.split(tail)] + + +def _flag(tokens: list[str], flag: str) -> str: + return tokens[tokens.index(flag) + 1] + + +def test_pair_plan_has_no_syncer_and_the_trainer_island_is_held(): + _, plan = _plan() + assert (plan["syncer"], plan["outer_sync"], plan["role_syncers"]) == (None, False, []) + assert plan["trainer_scheduler"] == {"mode": "s1", "island": 1, "interval_s": 30.0} + sampler, trainer = plan["island_requests"] + assert (sampler["sample_role"], trainer["sample_role"]) == ("sampler", "trainer") + assert "held_by_trainer_scheduler" not in sampler + assert trainer["held_by_trainer_scheduler"] is True + + +def test_pair_island_commands_carry_their_own_role_and_one_learner(): + _, plan = _plan() + sampler, trainer = (_island_tokens(i["learner_command"]) for i in plan["island_requests"]) + for tokens in (sampler, trainer): + assert "--rl-single-island-no-sync" in tokens and "--syncer" not in tokens + assert _flag(tokens, "--learner-id") == "0" + assert "--num-learners" not in tokens # one learner: no sweep roster + assert _flag(tokens, "--parameter-mode") == "full" and "--lora-r" not in tokens + assert _flag(tokens, "--rl-resume-store") == "/root/yeto-checkpoint-store/run-a" + assert _flag(tokens, "--rl-max-policy-age") == "2" + assert _flag(sampler, "--rl-sample-role") == "sampler" + assert _flag(sampler, "--rl-sampler-max-backlog-rounds") == "3" + assert "--rl-sample-wait-s" not in sampler + assert _flag(trainer, "--rl-sample-role") == "trainer" + assert _flag(trainer, "--rl-sample-wait-s") == "900.0" + assert "--rl-sampler-max-backlog-rounds" not in trainer + assert _flag(trainer, "--rl-cut-every") == "1" + + +def test_pair_modal_islands_keep_their_island_index_but_run_as_learner_zero(): + """Function name, tape subdir: island index. LEARNER_ID / NUM_LEARNERS: 0 / 1.""" + from yeto.gpu_spec import parse_gpu_spec + from yeto.launch.modal_islands import build_modal_island_config + from yeto.launcher import _island_task, make_miles_island_task + + args = _args(("--modal-tape-volume", "tapes")) + specs = parse_gpu_spec(args.gpu) + for m, spec in enumerate(specs): + task = _island_task(make_miles_island_task, args, spec, m, 2, "none", {}) + cfg = build_modal_island_config(args, spec, m, task, "none") + assert cfg.learner_id == m and cfg.tape_subdir.endswith(f"/l{m}") + assert cfg.envs["LEARNER_ID"] == "0" + assert cfg.checkpoint_store_volume_name == "yeto-l54" # one store for both + + +def test_pair_island_commands_pass_the_island_parser_and_no_sync_checks(): + from yeto.rl.adapters.miles import island_entry + from yeto.rl.engine import sample_roles + + _, plan = _plan() + for island, role in zip(plan["island_requests"], ("sampler", "trainer")): + parsed = island_entry.parse_args(_island_tokens(island["learner_command"])) + assert (parsed.parameter_mode, parsed.rl_sample_role) == ("full", role) + island_entry._check_single_island_no_sync(parsed) + assert island_entry._full_parameter_mode(parsed, parsed.sync_preset) is True + config = sample_roles.resolve(parsed, groups_per_round=parsed.groups_per_round) + assert (config.role, config.store) == (role, "/root/yeto-checkpoint-store/run-a") + + +def test_full_parameter_without_a_buffer_role_is_still_refused_on_the_island(): + from yeto.rl.adapters.miles import island_entry + + args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=True, + rl_sample_role="both") + with pytest.raises(ValueError, match="full-parameter sync preset"): + island_entry._full_parameter_mode(args, "strict-avg") + + +def _replace(base, old: tuple, new: tuple) -> list[str]: + """``base`` with the consecutive tokens ``old`` replaced by ``new``.""" + base = list(base) + for at in range(len(base) - len(old) + 1): + if tuple(base[at:at + len(old)]) == old: + return base[:at] + list(new) + base[at + len(old):] + raise AssertionError(f"{old} not in {base}") + + +NO_SCHEDULER = _replace(PAIR, ("--rl-trainer-scheduler", "s1", "--rl-trainer-island", "1"), ()) + + +@pytest.mark.parametrize("old, new, message", [ + (("--rl-single-island-no-sync",), (), "no-sync|single-island"), + (("--tuning", "full"), ("--tuning", "full", "--rl-sample-role", "trainer"), + "replaces --rl-sample-role"), + (("1:trainer",), ("1:sampler",), "one sampler and one trainer"), + (("--rl-max-policy-age", "2"), ("--rl-max-policy-age", "0"), "max-policy-age"), + (("--rl-checkpoint-store", STORE, "--rl-cut-every", "1"), (), "checkpoint-store"), +]) +def test_pair_refusals(old, new, message): + with pytest.raises((ValueError, SystemExit), match=message): + _args(base=_replace(NO_SCHEDULER, old, new)) + + +def test_pair_needs_two_islands(): + with pytest.raises(ValueError, match="island"): + _args(gpu="modal:1xh100") + + +def test_tuning_full_outside_buffer_islands_is_refused(): + with pytest.raises(ValueError, match="requires --tuning lora"): + _args(("--tuning", "full"), gpu="modal:1xh100", base=()) + + +def test_held_trainer_drains_the_buffer_after_the_fleet_ended(tmp_path, monkeypatch): + """When run() returns (the sampler ended), the loop sees data_ended and keeps + ticking until it finished; then run_with_trainer_loop returns.""" + from yeto.launch import trainer_loop as launch_loop + + monkeypatch.setenv("YETO_RUNS_DIR", str(tmp_path / "runs")) + seen: list[bool] = [] + finished = threading.Event() + + class _Loop: + finished = False + + def restore(self, records): + return False + + def tick(self): + ended = binding_box["binding"].data_ended() + seen.append(ended) + if ended: + self.finished = True + finished.set() + + def close(self, reason): + pass + + binding_box: dict = {} + + def fake_build(args, binding, journal): + binding_box["binding"] = binding + return _Loop() + + monkeypatch.setattr(launch_loop, "build", fake_build) + args = types.SimpleNamespace(rl_trainer_scheduler="s1", rl_trainer_scheduler_interval=0.01, + rl_trainer_drain_seconds=30.0, cluster_prefix="drain-a") + binding = launch_loop.CloudTrainerBinding(lifecycle=lambda first: None, buffer=None, cuts=None) + assert launch_loop.run_with_trainer_loop(args, lambda: "codes", + factory=lambda a, i: binding) == "codes" + assert finished.is_set() and seen[-1] is True diff --git a/tests/test_rl_sync_timing.py b/tests/test_rl_sync_timing.py new file mode 100644 index 000000000..a54220c6b --- /dev/null +++ b/tests/test_rl_sync_timing.py @@ -0,0 +1,100 @@ +"""rl-trainer-start-stop L5.6: island-side sync timing. A fake elastic syncer +delays its ELASTIC_BASE reply; the island must report that delay as +``sync/sync_wait_seconds``. CPU, no Ray.""" + +from __future__ import annotations + +import json +import socket +import threading +import time + +import pytest +import torch + +from yeto.protocol import read_frame, write_frame +from yeto.rl.elastic_client import ( + DeltaTensor, ElasticBase, ElasticClientConfig, ElasticInit, ElasticIslandClient, Join, JoinAck, + Leave, decode, encode, +) +from yeto.rl.engine.bridges import DecoupledSync, ElasticAvgSync +from yeto.rl.engine.sync_timing import SyncTiming + +from test_rl_engine_driver import _driver, _engine, _strict_config + +KEY = b"k" +DELAY_S = 0.4 + + +class DelayingSyncer: + """Answers each DELTA_TENSOR with base v+1 after ``delay_s``.""" + + def __init__(self, delay_s: float): + self.a, self.b = socket.socketpair() + self.delay_s = delay_s + self.version, self.base = 0, None + threading.Thread(target=self._serve, daemon=True).start() + + def _base(self): + write_frame(self.b, *encode(ElasticBase(0, self.version, tuple(self.base)), KEY)) + + def _serve(self): + while True: + try: + t, p = read_frame(self.b, 1 << 20) + except Exception: + return + msg = decode(KEY, t, p) + if isinstance(msg, Join): + write_frame(self.b, *encode(JoinAck(0, msg.island_id, 1, 0, bytes(32), False), KEY)) + elif isinstance(msg, ElasticInit): + self.base = list(msg.params) + self._base() + elif isinstance(msg, DeltaTensor): + time.sleep(self.delay_s) # the other islands / the merge + self.version += 1 + self.base = [b + u for b, u in zip(self.base, msg.update)] + self._base() + elif isinstance(msg, Leave): + self.b.close() + return + + def connect(self, addr): + return self.a + + +def test_sync_timing_split(): + t = SyncTiming(submit_ts=100.0, global_ready_ts=103.0, recv_seconds=0.5) + assert t.sync_wait_seconds == 2.5 + assert SyncTiming(100.0, 103.0).sync_wait_seconds == 3.0 # receive time not separated + assert SyncTiming(100.0, 99.0).sync_wait_seconds == 0.0 # base was there before the submit + assert SyncTiming(100.0, None).fields() == { + "sync/submit_ts": 100.0, "sync/global_ready_ts": None, "sync/recv_seconds": None, + "sync/sync_wait_seconds": None} + + +def test_elastic_island_measures_the_delayed_reply(tmp_path): + fake = DelayingSyncer(DELAY_S) + engine = _engine(torch.tensor([1.0, 3.0])) + client = ElasticIslandClient( + ElasticClientConfig(("x", 0), 0, lease_s=30.0, join_timeout_s=2, heartbeat_connection=False), + KEY, connect=fake.connect) + sync = ElasticAvgSync(_strict_config(tmp_path, engine, learner_id=0, rounds=2), client=client, + base_wait_s=5.0) + _driver(engine, sync, tmp_path).run() + client.close() + tape = [json.loads(l) for l in (tmp_path / "events.jsonl").read_text().splitlines()] + timing = [e for e in tape if e["event"] == "rl_sync_timing"] + assert [e["rollout_id"] for e in timing] == [0, 1] + for e in timing: + assert e["outer_sync"] == "elastic" + assert e["sync/global_ready_ts"] >= e["sync/submit_ts"] + assert 0.0 <= e["sync/recv_seconds"] < 0.2 + assert e["sync/sync_wait_seconds"] == pytest.approx(DELAY_S, abs=0.15) + + +def test_decoupled_remote_wait_is_the_blocked_time(): + sync = DecoupledSync(args=None) + assert sync._remote_wait(time.sleep, 0.2) is None + sync._remote_wait(lambda: None) + assert 0.2 <= sync._remote_wait_seconds < 0.4 diff --git a/tests/test_rl_trainer_lifecycle.py b/tests/test_rl_trainer_lifecycle.py new file mode 100644 index 000000000..50a566e6e --- /dev/null +++ b/tests/test_rl_trainer_lifecycle.py @@ -0,0 +1,211 @@ +"""rl-trainer-start-stop L3 (CPU, fakes): lifecycle states, SCHEDULED cuts, +JOIN/LEAVE per incarnation, and repeated starts over IslandDriver.""" + +from __future__ import annotations + +import time + +import pytest +import torch +from test_rl_engine_driver import NAME, PROFILE, _driver + +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.driver import RoundFailedError +from yeto.rl.engine.fake_lifecycle import FakeMembership, FakeRoundCuts, cuttable_fake_engine +from yeto.rl.engine.trainer_lifecycle import ( + AlwaysOnScheduler, + BacklogStats, + CutRef, + DriverLifecycle, + Hold, + LifecycleError, + LifecycleState, + Start, + Stop, + StopCondition, + TrainerLifecycle, + allowed_transition, +) + +FULL = [s.value for s in LifecycleState] + + +class _Island: + """Builds a fresh fake trainer + LocalOnlySync per incarnation (acquire).""" + + def __init__(self, tmp_path, *, num_rollout=100, **engine_kwargs): + self.tmp_path = tmp_path + self.num_rollout = num_rollout + self.engine_kwargs = engine_kwargs + self.cuts = FakeRoundCuts(tmp_path / "cuts") + self.membership = FakeMembership() + self.engines = [] + self.released = [] + + def build(self, handle): + engine = cuttable_fake_engine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0, + capability_profile=PROFILE, **self.engine_kwargs) + self.engines.append(engine) + sync = LocalOnlySync(self.num_rollout) + sync.round_cuts = self.cuts + return _driver(engine, sync, self.tmp_path, name=f"events-{handle.incarnation}.jsonl") + + def lifecycle(self): + return DriverLifecycle(self.build, membership=self.membership, cuts=self.cuts, + release=lambda h: self.released.append(h.incarnation)) + + +def _deadline(s=30.0): + return time.monotonic() + s + + +def _publishes(engine): + return [c[1] for c in engine.calls if c[0] == "publish"] + + +def test_transitions_move_forward_or_to_stopped(): + s = LifecycleState + assert allowed_transition(s.COLD, s.ACQUIRING) and allowed_transition(s.TRAINING, s.DRAINING) + assert allowed_transition(s.RESTORING, s.STOPPED) + assert not allowed_transition(s.COLD, s.TRAINING) + assert not allowed_transition(s.TRAINING, s.WARMING) + assert not allowed_transition(s.STOPPED, s.STOPPED) + + +def test_three_start_stop_cycles_continue_the_policy(tmp_path): + """L3.1 + L3.4: start -> train -> stop three times on fresh fake trainers; each + start restores the SCHEDULED cut of the previous stop and goes on at v+1.""" + island = _Island(tmp_path) + lc = island.lifecycle() + assert isinstance(lc, TrainerLifecycle) + cuts = [] + for cycle in range(3): + resume = island.cuts.latest() + h = lc.start(lc.plan_start({"cycle": cycle}), resume_from=resume) + assert lc.ready(h, _deadline()) + receipt = lc.run(h, None, StopCondition(rounds=2)) + assert h.state is LifecycleState.TRAINING + assert (receipt.rounds, receipt.first_version, receipt.trained_version) == \ + (2, 2 * cycle, 2 * cycle + 2) + cut = lc.drain_and_cut(h, LeaveReason.STOPPED) + lc.stop(h, LeaveReason.STOPPED) + assert h.states == FULL and h.error is None + assert cut == CutRef(cut_id=cut.cut_id, trigger=recovery.SCHEDULED, policy_version=2 * cycle + 2) + cuts.append(cut) + engine = island.engines[cycle] + # the restored version is published again first, then v+1, v+2 + assert _publishes(engine) == [2 * cycle, 2 * cycle + 1, 2 * cycle + 2] + restored = [c for c in engine.calls if c[0] == "restore_cut"] + assert restored == ([] if resume is None else [("restore_cut", resume.cut_id)]) + assert torch.equal(engine.tensors[NAME], torch.full((1, 2), float(2 * cycle + 2))) + assert len({id(e) for e in island.engines}) == 3 # a fresh trainer per start + assert island.membership.tape == [ + ("JOIN", 0), ("LEAVE", 0, LeaveReason.STOPPED), + ("JOIN", 1), ("LEAVE", 1, LeaveReason.STOPPED), + ("JOIN", 2), ("LEAVE", 2, LeaveReason.STOPPED), + ] + assert island.released == [0, 1, 2] + assert [c.trigger for c in cuts] == [recovery.SCHEDULED] * 3 + + +def test_run_twice_before_the_stop_trains_on(tmp_path): + island = _Island(tmp_path) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + assert lc.run(h, None, StopCondition(1)).trained_version == 1 + assert lc.run(h, None, StopCondition(2)).trained_version == 3 + lc.stop(h, LeaveReason.STOPPED) # drains first + assert h.cut.policy_version == 3 and h.cut.trigger == recovery.SCHEDULED + assert island.membership.tape[-1] == ("LEAVE", 0, LeaveReason.STOPPED) + + +def test_data_end_drains_on_its_own_with_finished(tmp_path): + island = _Island(tmp_path, num_rollout=3) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + receipt = lc.run(h, None, StopCondition()) + assert receipt.data_ended and receipt.trained_version == 3 + assert h.state is LifecycleState.DRAINING + lc.stop(h, LeaveReason.FINISHED) + assert h.states == FULL + assert h.cut.trigger == recovery.MANUAL and h.cut.policy_version == 3 + assert island.membership.tape == [("JOIN", 0), ("LEAVE", 0, LeaveReason.FINISHED)] + + +def test_reclaim_cuts_with_preempt_and_leaves_preempted(tmp_path): + island = _Island(tmp_path) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + lc.run(h, None, StopCondition(1)) + cut = lc.drain_and_cut(h, LeaveReason.PREEMPTED) + lc.stop(h, LeaveReason.PREEMPTED) + assert cut.trigger == recovery.PREEMPT and cut.policy_version == 1 + assert island.membership.tape[-1] == ("LEAVE", 0, LeaveReason.PREEMPTED) + # the next start restores the PREEMPT cut like any other + h2 = lc.start(lc.plan_start({}), resume_from=island.cuts.latest()) + assert lc.run(h2, None, StopCondition(1)).trained_version == 2 + lc.stop(h2, LeaveReason.STOPPED) + + +def test_stop_before_any_round_keeps_the_previous_cut(tmp_path): + island = _Island(tmp_path) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + lc.run(h, None, StopCondition(1)) + lc.stop(h, LeaveReason.STOPPED) + first = island.cuts.latest() + h2 = lc.start(lc.plan_start({}), resume_from=first) + assert lc.ready(h2, _deadline()) + lc.stop(h2, LeaveReason.STOPPED) + assert h2.states == FULL and island.cuts.latest() == first + assert island.membership.tape[-2:] == [("JOIN", 1), ("LEAVE", 1, LeaveReason.STOPPED)] + + +def test_training_failure_stops_with_the_p5_code(tmp_path): + island = _Island(tmp_path, failed_step_rounds={1}) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + with pytest.raises(LifecycleError) as raised: + lc.run(h, None, StopCondition(3)) + assert isinstance(raised.value.__cause__, RoundFailedError) + assert h.state is LifecycleState.STOPPED and h.states[-2] == "training" + assert island.membership.tape == [("JOIN", 0), ("LEAVE", 0, LeaveReason.TRAINING_FAILED)] + assert island.released == [0] + lc.stop(h, LeaveReason.STOPPED) # already stopped: no second LEAVE + assert len(island.membership.tape) == 2 + + +def test_start_refuses_a_restore_at_another_version(tmp_path): + island = _Island(tmp_path) + lc = island.lifecycle() + h = lc.start(lc.plan_start({}), resume_from=None) + lc.run(h, None, StopCondition(1)) + lc.stop(h, LeaveReason.STOPPED) + stale = CutRef(cut_id="other", trigger=recovery.SCHEDULED, policy_version=7) + h2 = lc.start(lc.plan_start({}), resume_from=stale) + assert not lc.ready(h2, _deadline()) + assert h2.state is LifecycleState.STOPPED and isinstance(h2.error, LifecycleError) + assert island.membership.tape[-1] == ("LEAVE", 1, LeaveReason.INFRA_FAILED) + + +def test_acquire_failure_never_joins(tmp_path): + membership = FakeMembership() + + def build(handle): + raise RuntimeError("no capacity") + + lc = DriverLifecycle(build, membership=membership) + with pytest.raises(RuntimeError, match="no capacity"): + lc.start(lc.plan_start({}), resume_from=None) + assert membership.tape == [] + + +def test_always_on_scheduler_is_start_once(): + s = AlwaysOnScheduler() + busy = BacklogStats(batches=2.0, samples=8) + assert s.decide(busy, LifecycleState.COLD) == Start() + assert s.decide(busy, LifecycleState.TRAINING, {"h100": 3.0}) == Hold() + ended = BacklogStats(batches=0.0, samples=0, data_ended=True) + assert s.decide(ended, LifecycleState.TRAINING) == Stop(LeaveReason.FINISHED) diff --git a/tests/test_rl_trainer_loop.py b/tests/test_rl_trainer_loop.py new file mode 100644 index 000000000..92a79b45b --- /dev/null +++ b/tests/test_rl_trainer_loop.py @@ -0,0 +1,304 @@ +"""rl-trainer-start-stop L5.3b (CPU, fakes): the scheduler in the controller loop. + +A fake cloud (the island builder: a fresh fake trainer per incarnation, released +on stop) and the real file sample buffer, scheduler S1 and controller journal. +The head restarts in the middle: a new loop over the same journal continues.""" + +from __future__ import annotations + +import argparse +import time + +import pytest +import torch +from test_rl_engine_driver import NAME, PROFILE, _driver + +from yeto import runs +from yeto.launch import trainer_loop as launch_loop +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.fake_lifecycle import ( + BufferFedFakeRollout, + FakeMembership, + FakeRoundCuts, + FakeSampler, + cuttable_fake_engine, +) +from yeto.rl.engine.journal import Journal +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import FileSampleBuffer +from yeto.rl.engine.trainer_lifecycle import DriverLifecycle, Hold, LifecycleState, Start, Stop +from yeto.rl.engine.trainer_loop import LOOP_KIND, TrainerLoop, TrainerLoopThread, next_incarnation +from yeto.rl.engine.trainer_scheduler import SCHEDULER_STATE_KIND, S1Config, S1Scheduler + +BATCH = 4 # fake engine: 2 groups x 2 samples +T_S = 5.0 + + +class _Clock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +class _FakeCloud: + """Acquire = build a fresh fake trainer island; release = give the cards back.""" + + def __init__(self, tmp_path) -> None: + self.tmp_path = tmp_path + self.cuts = FakeRoundCuts(tmp_path / "cuts", watermark=lambda: self.trainer_buffer.watermark) + self.membership = FakeMembership() + self.buffer = None # the head's buffer + self.trainer_buffer = None # the buffer of the running trainer + self.rollouts = [] + self.engine = None + self.acquired: list[int] = [] + self.released: list[int] = [] + + def new_buffer(self) -> FileSampleBuffer: + # every head opens the buffer in the shared store; this in-process trainer shares it + self.buffer = FileSampleBuffer(CheckpointStore(self.tmp_path / "store"), batch_size=BATCH) + return self.buffer + + def build(self, handle): + self.acquired.append(handle.incarnation) + engine = cuttable_fake_engine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0, + capability_profile=PROFILE) + resume = handle.resume_from + self.trainer_buffer = self.buffer + self.buffer.reset_to(0 if resume is None else resume.watermark) + engine.rollout = BufferFedFakeRollout(engine, self.buffer) + self.rollouts.append(engine.rollout) + self.engine = engine + sync = LocalOnlySync(1000) + sync.round_cuts = self.cuts + return _driver(engine, sync, self.tmp_path, name=f"events-{handle.incarnation}.jsonl") + + def release(self, handle) -> None: + self.released.append(handle.incarnation) + + def published(self): + return None if self.engine is None or self.engine.published is None else self.engine.published[0] + + +class _Head: + """One head process: journal, scheduler, lifecycle, loop.""" + + def __init__(self, tmp_path, cloud: _FakeCloud, clock: _Clock, *, release_orphan=None, + data_ended=lambda: False) -> None: + self.journal = Journal(tmp_path / "head-journal") + records = self.journal.records + cloud.new_buffer() + self.scheduler = S1Scheduler(S1Config(a_hard=0, start_seconds=T_S), + record=self.journal.append, clock=clock) + self.lifecycle = DriverLifecycle(cloud.build, membership=cloud.membership, cuts=cloud.cuts, + release=cloud.release, + first_incarnation=next_incarnation(records)) + self.loop = TrainerLoop(scheduler=self.scheduler, lifecycle=self.lifecycle, + buffer=cloud.buffer, cuts=cloud.cuts, record=self.journal.append, + data_ended=data_ended, release_orphan=release_orphan, + ready_timeout_s=30.0) + self.restored = self.loop.restore(records) + # after the orphan is released: the backlog continues at the newest cut + cut = cloud.cuts.latest() + cloud.buffer.reset_to(0 if cut is None else cut.watermark) + + def crash(self) -> None: + self.journal.close() # the process dies: the journal lock goes with it + + +def _ticks(loop, n): + return [type(loop.tick()) for _ in range(n)] + + +def _idle_stop(loop, clock): + """Backlog empty: the first tick starts the idle time, T_s later the loop stops.""" + assert _ticks(loop, 1) == [Hold] and loop.state is LifecycleState.TRAINING + clock.now += T_S + assert _ticks(loop, 1) == [Stop] and loop.state is LifecycleState.STOPPED + + +def test_start_stop_cycles_across_a_head_restart(tmp_path): + clock = _Clock() + cloud = _FakeCloud(tmp_path) + data_end = {"ended": False} + head = _Head(tmp_path, cloud, clock, data_ended=lambda: data_end["ended"]) + assert head.restored is False and head.loop.state is LifecycleState.COLD + sampler = FakeSampler(cloud.buffer, cloud.published) + + # cycle 0: Q_on = A_hard + 1 = 1 batch; 2 batches -> start, train 2 rounds, idle stop + assert _ticks(head.loop, 1) == [Hold] # empty backlog: cold trainer holds + sampler.produce(2 * BATCH) + assert _ticks(head.loop, 3) == [Start, Hold, Hold] + _idle_stop(head.loop, clock) + cut0 = cloud.cuts.latest() + assert (cut0.trigger, cut0.policy_version, cut0.watermark) == (recovery.SCHEDULED, 2, 2 * BATCH) + + # cycle 1: one more batch; the head dies while the trainer runs and is idle + sampler.load() + sampler.produce(BATCH) + assert _ticks(head.loop, 3) == [Start, Hold, Hold] + assert head.scheduler.idle_since == clock.now + old_lifecycle, old_handle = head.lifecycle, head.loop.handle + head.crash() + + # the cloud reclaims the orphaned trainer (spot): PREEMPT cut, LEAVE PREEMPTED + orphans = [] + + def release_orphan(incarnation): + orphans.append(incarnation) + old_lifecycle.stop(old_handle, LeaveReason.PREEMPTED) + + head = _Head(tmp_path, cloud, clock, release_orphan=release_orphan, + data_ended=lambda: data_end["ended"]) + assert head.restored and orphans == [1] + assert head.loop.state is LifecycleState.STOPPED and next_incarnation(head.journal.records) == 2 + cut1 = cloud.cuts.latest() + assert (cut1.trigger, cut1.policy_version, cut1.watermark) == (recovery.PREEMPT, 3, 3 * BATCH) + sampler = FakeSampler(cloud.buffer, cloud.published, name="s1") + + # cycle 2 after the restart: continues at v3, idle stop -> SCHEDULED + assert _ticks(head.loop, 1) == [Hold] + sampler.load() + sampler.produce(BATCH) + assert _ticks(head.loop, 2) == [Start, Hold] + assert head.loop.handle.start_version == 3 and head.loop.current_version() == 4 + _idle_stop(head.loop, clock) + assert cloud.cuts.latest().trigger == recovery.SCHEDULED + + # cycle 3: data end -> FINISHED, the loop is done + sampler.load() + sampler.produce(BATCH) + assert _ticks(head.loop, 2) == [Start, Hold] + data_end["ended"] = True + assert _ticks(head.loop, 1) == [Stop] and head.loop.finished + assert _ticks(head.loop, 1) == [Hold] + head.crash() + + assert cloud.membership.tape == [ + ("JOIN", 0), ("LEAVE", 0, LeaveReason.STOPPED), + ("JOIN", 1), ("LEAVE", 1, LeaveReason.PREEMPTED), + ("JOIN", 2), ("LEAVE", 2, LeaveReason.STOPPED), + ("JOIN", 3), ("LEAVE", 3, LeaveReason.FINISHED), + ] + assert cloud.acquired == [0, 1, 2, 3] and cloud.released == [0, 1, 2, 3] + # versions continue across starts; no sample is trained twice + trained = [sid for r in cloud.rollouts for batch in r.trained.values() for sid, _ in batch] + assert len(trained) == 5 * BATCH and len(set(trained)) == len(trained) + assert sorted(v for r in cloud.rollouts for v in r.trained) == [0, 1, 2, 3, 4] + + # the journal holds the loop and the scheduler records of both heads + head = Journal(tmp_path / "head-journal") + loop_events = [(r["event"], r["incarnation"]) for r in head.records if r["kind"] == LOOP_KIND] + assert loop_events == [ + ("started", 0), ("ready", 0), ("stopped", 0), + ("started", 1), ("ready", 1), ("orphaned", 1), + ("started", 2), ("ready", 2), ("stopped", 2), + ("started", 3), ("ready", 3), ("stopped", 3), + ] + stops = [r for r in head.records if r["kind"] == LOOP_KIND and r["event"] == "stopped"] + assert [(r["reason"], r["trigger"]) for r in stops] == [ + ("STOPPED", recovery.SCHEDULED), ("STOPPED", recovery.SCHEDULED), ("FINISHED", recovery.MANUAL)] + assert any(r["kind"] == SCHEDULER_STATE_KIND for r in head.records) + head.close() + + +def test_restart_continues_the_idle_time(tmp_path): + """D8: the idle time spent before the head restart counts after it.""" + clock = _Clock() + cloud = _FakeCloud(tmp_path) + head = _Head(tmp_path, cloud, clock) + FakeSampler(cloud.buffer, cloud.published).produce(BATCH) + assert _ticks(head.loop, 3) == [Start, Hold, Hold] # train, then idle starts + idle_since = head.scheduler.idle_since + handle, lifecycle = head.loop.handle, head.lifecycle + head.crash() + + clock.now += T_S - 1 + head = _Head(tmp_path, cloud, clock, release_orphan=lambda inc: None) + assert head.scheduler.idle_since == idle_since + # the orphan is released by the cloud layer; this fake keeps it, end it here + lifecycle.stop(handle, LeaveReason.STOPPED) + head.crash() + + +def test_thread_ticks_until_finished(tmp_path): + clock = _Clock() + cloud = _FakeCloud(tmp_path) + head = _Head(tmp_path, cloud, clock, data_ended=lambda: True) + FakeSampler(cloud.buffer, cloud.published).produce(BATCH) + runner = TrainerLoopThread(head.loop, interval_s=0.01).start() + deadline = time.monotonic() + 30 + while not head.loop.finished and time.monotonic() < deadline: + time.sleep(0.01) + runner.stop() + assert runner.error is None and head.loop.finished + assert cloud.membership.tape == [("JOIN", 0), ("LEAVE", 0, LeaveReason.FINISHED)] + head.crash() + + +def test_new_journal_starts_at_incarnation_zero(tmp_path): + assert next_incarnation([]) == 0 + assert next_incarnation([{"kind": LOOP_KIND, "incarnation": 4}, {"kind": "other", "incarnation": 9}]) == 5 + + +# ---------------------------------------------------------------- launcher switch +def _args(tmp_path, monkeypatch, mode): + monkeypatch.setattr(runs, "RUNS_DIR", tmp_path / "runs") + parser = argparse.ArgumentParser() + launch_loop.add_arguments(parser) + args = parser.parse_args([] if mode is None else ["--rl-trainer-scheduler", mode, + "--rl-trainer-scheduler-interval", "0.01", + "--rl-trainer-start-seconds", "0"]) + args.cluster_prefix, args.rl_max_policy_age = "run-a", 0 + return args + + +def test_switch_off_by_default_runs_unchanged(tmp_path, monkeypatch): + args = _args(tmp_path, monkeypatch, None) + assert args.rl_trainer_scheduler == "off" and not launch_loop.enabled(args) + launch_loop.check(args) # no binding needed + assert launch_loop.run_with_trainer_loop(args, lambda: "codes") == "codes" + assert not (tmp_path / "runs").exists() # no journal written + + +def test_switch_on_without_a_binding_refuses(tmp_path, monkeypatch): + args = _args(tmp_path, monkeypatch, "s1") + monkeypatch.setattr(launch_loop, "binding_factory", None) + with pytest.raises(SystemExit, match="no trainer island binding"): + launch_loop.check(args) + + +def test_switch_on_runs_the_loop_beside_the_controller(tmp_path, monkeypatch): + """The fleet controller's run() returns after the samplers finished; the loop + started and stopped the trainer meanwhile, and stops it at the end.""" + args = _args(tmp_path, monkeypatch, "s1") + cloud = _FakeCloud(tmp_path) + cloud.new_buffer() + FakeSampler(cloud.buffer, cloud.published).produce(2 * BATCH) + + def binding(_args, _island): + return launch_loop.CloudTrainerBinding( + lifecycle=lambda first: DriverLifecycle(cloud.build, membership=cloud.membership, + cuts=cloud.cuts, release=cloud.release, + first_incarnation=first), + buffer=cloud.buffer, cuts=cloud.cuts) + + def controller_run(): + deadline = time.monotonic() + 30 + while (cloud.membership.tape[-1:] != [("LEAVE", 0, LeaveReason.STOPPED)] + and time.monotonic() < deadline): + time.sleep(0.01) + return "codes" + + assert launch_loop.run_with_trainer_loop(args, controller_run, factory=binding) == "codes" + assert cloud.membership.tape == [("JOIN", 0), ("LEAVE", 0, LeaveReason.STOPPED)] + assert cloud.cuts.latest().policy_version == 2 + journal = Journal(launch_loop.journal_dir(args)) + # after run() returned (the samplers ended), the loop drained: data ended -> finished + assert [r["event"] for r in journal.records if r["kind"] == LOOP_KIND] == [ + "started", "ready", "stopped", "finished"] + journal.close() diff --git a/tests/test_rl_trainer_loop_failures.py b/tests/test_rl_trainer_loop_failures.py new file mode 100644 index 000000000..d6ffa0980 --- /dev/null +++ b/tests/test_rl_trainer_loop_failures.py @@ -0,0 +1,235 @@ +"""T-FPUB: the trainer loop stops after too many failed incarnations in a row. + +GPU run s24-l54-pair-20261014i: every trainer incarnation failed at its first +publish and the scheduler started a new one about 15 s later, with no limit. +Now ``max_consecutive_failures`` (default 3) failed incarnations with no trained +round between them stop the loop: it releases the trainer, journals +``gave_up`` and raises; a trained round resets the count. +""" + +from __future__ import annotations + +import argparse +import threading +from types import SimpleNamespace + +import pytest + +from yeto import runs +from yeto.launch import trainer_loop as launch_loop +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine.journal import Journal +from yeto.rl.engine.trainer_lifecycle import ( + BacklogStats, + Hold, + LifecycleError, + LifecycleState, + RunReceipt, + Start, +) +from yeto.rl.engine.trainer_loop import ( + DEFAULT_MAX_CONSECUTIVE_FAILURES, + GAVE_UP, + LOOP_KIND, + STARTED, + TrainerLoop, + TrainerLoopGaveUp, + TrainerLoopThread, +) + + +class _Scheduler: + """Start whenever no trainer runs; else hold (the loop trains on a hold).""" + + def __init__(self, *_a, **_k) -> None: + pass + + def decide(self, backlog, state): + return Hold() if state is LifecycleState.TRAINING else Start() + + +class _Buffer: + def backlog(self, _version): + return BacklogStats(batches=5.0, samples=20) + + +class _Cuts: + def latest(self): + return None + + +class _Lifecycle: + """``outcomes``: per incarnation, "start_error" (start raises), "train_error" + (the first run fails, as the L5.4 publish), or "ok" (each run trains a round).""" + + def __init__(self, outcomes) -> None: + self.outcomes = list(outcomes) + self.started: list[int] = [] + self.stopped: list[tuple[int, LeaveReason]] = [] + self._next = 0 + + def plan_start(self, spec): + plan = SimpleNamespace(incarnation=self._next) + self._next += 1 + return plan + + def start(self, plan, resume_from=None): + self.started.append(plan.incarnation) + outcome = self.outcomes[plan.incarnation] if plan.incarnation < len(self.outcomes) else "train_error" + if outcome == "start_error": + raise RuntimeError("island did not come up") + return SimpleNamespace(incarnation=plan.incarnation, state=LifecycleState.TRAINING, + outcome=outcome, start_version=0, trained_version=None, + error=None, leave_reason=None, cut=None, released=None) + + def ready(self, handle, deadline): + return True + + def run(self, handle, buffer, condition): + if handle.outcome == "ok": + handle.outcome = "train_error" # one good round, then it breaks + handle.trained_version = (handle.trained_version or 0) + 1 + return RunReceipt(rounds=1, first_version=0, trained_version=handle.trained_version, + data_ended=False) + handle.state = LifecycleState.STOPPED + handle.error = RuntimeError("Megatron produced an empty or duplicate LoRA mapping") + handle.leave_reason = LeaveReason.INFRA_FAILED + raise LifecycleError("trainer failed") + + def stop(self, handle, reason): + self.stopped.append((handle.incarnation, reason)) + handle.state = LifecycleState.STOPPED + + +def _loop(outcomes, **kwargs): + records: list[dict] = [] + lifecycle = _Lifecycle(outcomes) + loop = TrainerLoop(scheduler=_Scheduler(), lifecycle=lifecycle, buffer=_Buffer(), cuts=_Cuts(), + record=lambda kind, **f: records.append({"kind": kind, **f}), **kwargs) + return loop, lifecycle, records + + +def _tick_until_gave_up(loop, limit=50): + for _ in range(limit): + try: + loop.tick() + except TrainerLoopGaveUp as error: + return error + raise AssertionError("the loop never gave up") + + +def test_three_failed_incarnations_stop_the_loop(): + loop, lifecycle, records = _loop([]) + assert loop.max_consecutive_failures == DEFAULT_MAX_CONSECUTIVE_FAILURES == 3 + error = _tick_until_gave_up(loop) + assert "3 trainer incarnations failed in a row" in str(error) + assert lifecycle.started == [0, 1, 2] # no fourth Start + events = [r["event"] for r in records if r["kind"] == LOOP_KIND] + assert events.count(STARTED) == 3 and events[-1] == GAVE_UP + assert records[-1]["failures"] == 3 and records[-1]["limit"] == 3 + assert loop.finished and isinstance(loop.tick(), Hold) # nothing more happens + assert lifecycle.started == [0, 1, 2] + + +def test_failed_starts_count_and_the_limit_is_configurable(): + loop, lifecycle, records = _loop(["start_error"] * 5, max_consecutive_failures=2) + _tick_until_gave_up(loop) + assert lifecycle.started == [0, 1] + with pytest.raises(ValueError, match="max_consecutive_failures"): + _loop([], max_consecutive_failures=0) + + +def test_a_trained_round_resets_the_count(): + # 0, 1 fail; 2 trains a round (count back to 0) then fails; 3, 4 fail -> give up + loop, lifecycle, _records = _loop(["train_error", "start_error", "ok"]) + _tick_until_gave_up(loop) + assert lifecycle.started == [0, 1, 2, 3, 4] + + +def test_thread_reports_the_give_up_to_the_launcher(): + loop, _lifecycle, _records = _loop([], max_consecutive_failures=1) + told: list[BaseException] = [] + runner = TrainerLoopThread(loop, 0.001, on_error=told.append).start() + assert runner.join(10) + assert isinstance(runner.error, TrainerLoopGaveUp) and told == [runner.error] + + +def _args(tmp_path, monkeypatch, *extra): + monkeypatch.setattr(runs, "RUNS_DIR", tmp_path / "runs") + parser = argparse.ArgumentParser() + launch_loop.add_arguments(parser) + args = parser.parse_args(["--rl-trainer-scheduler", "s1", "--rl-trainer-scheduler-interval", + "0.001", *extra]) + args.cluster_prefix, args.rl_max_policy_age = "run-a", 0 + return args + + +def _binding(lifecycle): + def factory(_args, _island): + return launch_loop.CloudTrainerBinding(lifecycle=lambda first: lifecycle, + buffer=_Buffer(), cuts=_Cuts()) + return factory + + +def test_launcher_aborts_the_fleet_and_fails_the_run(tmp_path, monkeypatch): + monkeypatch.setattr(launch_loop, "S1Scheduler", _Scheduler) + args = _args(tmp_path, monkeypatch, "--rl-trainer-max-failures", "2") + lifecycle = _Lifecycle([]) + aborted = threading.Event() + reasons: list[BaseException] = [] + + def abort(error): + reasons.append(error) + aborted.set() + + def controller_run(): # the fleet controller: ends with the abort error + assert aborted.wait(10) + raise reasons[0] + + with pytest.raises(TrainerLoopGaveUp): + launch_loop.run_with_trainer_loop(args, controller_run, factory=_binding(lifecycle), + abort=abort) + assert lifecycle.started == [0, 1] + journal = Journal(launch_loop.journal_dir(args)) + assert [r["event"] for r in journal.records if r["kind"] == LOOP_KIND][-1] == GAVE_UP + journal.close() + + +def test_launcher_without_abort_raises_after_the_fleet(tmp_path, monkeypatch): + monkeypatch.setattr(launch_loop, "S1Scheduler", _Scheduler) + args = _args(tmp_path, monkeypatch, "--rl-trainer-max-failures", "1") + ended = threading.Event() + + class _Thread(TrainerLoopThread): + def _run(self) -> None: + super()._run() + ended.set() + + monkeypatch.setattr(launch_loop, "TrainerLoopThread", _Thread) + + def controller_run(): + assert ended.wait(10) + return "codes" + + with pytest.raises(TrainerLoopGaveUp): + launch_loop.run_with_trainer_loop(args, controller_run, factory=_binding(_Lifecycle([]))) + + +def test_bad_limit_is_refused_before_launch(tmp_path, monkeypatch): + args = _args(tmp_path, monkeypatch, "--rl-trainer-max-failures", "0") + monkeypatch.setattr(launch_loop, "binding_factory", lambda a, i: None) + with pytest.raises(SystemExit, match="max-failures"): + launch_loop.check(args) + + +def test_fleet_controller_run_ends_with_the_abort_error(): + from yeto.launch.fleet import FleetController + + controller = FleetController.__new__(FleetController) + controller.learners = {} + controller._abort_error = None + controller._fleet = lambda *a, **k: None + error = TrainerLoopGaveUp("gave up") + controller.abort(error) + with pytest.raises(TrainerLoopGaveUp): + controller.run() diff --git a/tests/test_rl_trainer_remote.py b/tests/test_rl_trainer_remote.py new file mode 100644 index 000000000..e8d13f136 --- /dev/null +++ b/tests/test_rl_trainer_remote.py @@ -0,0 +1,435 @@ +"""rl-trainer-start-stop L5.3c (CPU, fakes): the trainer on its own island. + +The head runs the scheduler loop with :class:`RemoteTrainerLifecycle`; each +incarnation is a fake "island" (a thread running :class:`IslandTrainerAgent` over a +fresh fake driver). Head and island share only the store directory (commands, cut +pointer, sample buffer) and the island's event lines (``YETO_RL_EVENT`` echo +format). No cloud resource is started.""" + +from __future__ import annotations + +import argparse +import json +import threading +from pathlib import Path + +import pytest +import torch +from test_rl_engine_driver import NAME, PROFILE, _driver + +from yeto import runs +from yeto.launch import trainer_loop as launch_loop +from yeto.launch.trainer_island import ( + IslandEvents, + MirrorStore, + ModalTrainerIslands, + ModalVolumeFiles, + TrainerIsland, + modal_binding, +) +from yeto.rl import event_echo +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine import recovery +from yeto.rl.engine.bridges import LocalOnlySync +from yeto.rl.engine.fake_lifecycle import BufferFedFakeRollout, FakeRoundCuts, FakeSampler, cuttable_fake_engine +from yeto.rl.engine.journal import Journal +from yeto.rl.engine.resume import FLAT_MANIFEST, CheckpointStore, write_snapshot +from yeto.rl.engine.sample_buffer import BUFFER_DIR, INDEX, CommittedBacklog, FileSampleBuffer +from yeto.rl.engine.start_phases import SUBMITTED_ENV +from yeto.rl.engine.trainer_lifecycle import BacklogStats, Hold, LifecycleState, Start, Stop +from yeto.rl.engine.trainer_loop import FINISHED, LOOP_KIND, TrainerLoop +from yeto.rl.engine.trainer_remote import ( + CommandReader, + CommandWriter, + IslandTrainerAgent, + LocalStoreFiles, + RemoteTrainerLifecycle, + StoreCutLog, + command_path, + consumed_watermark, + island_env, +) +from yeto.rl.engine.trainer_scheduler import S1Config, S1Scheduler + +BATCH = 4 # fake engine: 2 groups x 2 samples +T_S = 5.0 +POLL = 0.005 + + +class _Clock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +class _StoreRoundCuts(FakeRoundCuts): + """Fake round cuts with the pointer where the real one lands (flat layout).""" + + @property + def pointer_path(self) -> Path: + return self.root / "round-cut.json" + + +class _FakeIslands: + """The cloud: ``launch`` starts an island thread; its events travel as echo + lines. ``die_at`` makes one incarnation's island exit without a stopped event.""" + + def __init__(self, store_root: Path, sampler_buffer: FileSampleBuffer) -> None: + self.store_root = store_root + self.sampler_buffer = sampler_buffer + self.engine = None + self.rollouts: list = [] + self.launched: list[tuple[int, str | None]] = [] + self.released: list[int] = [] + self.agents: list = [] + self.die_at: int | None = None + + def published(self): + return None if self.engine is None or self.engine.published is None else self.engine.published[0] + + def launch(self, incarnation, resume_from): + env = island_env(incarnation, resume_from) # what the island would read + self.launched.append((incarnation, resume_from and resume_from.cut_id)) + store = CheckpointStore(self.store_root) + buffer = FileSampleBuffer(store, batch_size=BATCH) # the trainer's own reader + cuts = _StoreRoundCuts(self.store_root, watermark=lambda: buffer.watermark) + buffer.reset_to(0 if resume_from is None else resume_from.watermark) + engine = cuttable_fake_engine(tensors={NAME: torch.zeros(1, 2)}, step_delta=1.0, + capability_profile=PROFILE) + engine.rollout = BufferFedFakeRollout(engine, buffer) + self.engine = engine + self.rollouts.append(engine.rollout) + sync = LocalOnlySync(1000) + sync.round_cuts = cuts + driver = _driver(engine, sync, self.store_root.parent, name=f"tape-{incarnation}.jsonl") + job = {"incarnation": incarnation, "events": IslandEvents(), "alive": True, "env": env} + + def emit(event, **fields): + job["events"].feed(event_echo.format_record( + {"island_id": 0, "time_unix": 0.0, "event": event, **fields})) + + files = LocalStoreFiles(store) + agent = IslandTrainerAgent(driver, commands=CommandReader(files, incarnation), + incarnation=int(env["YETO_TRAINER_INCARNATION"]), + resume_from=resume_from, emit=emit, cuts=StoreCutLog(files), + watermark=lambda version: buffer.watermark, poll_s=POLL) + self.agents.append(agent) + + def island(): + try: + if self.die_at == incarnation: + return # crashed before any event + agent.run() + except Exception: # noqa: BLE001 - the stopped event carries it + pass + finally: + job["alive"] = False + + job["thread"] = threading.Thread(target=island, daemon=True) + job["thread"].start() + return job + + def events(self, job): + return job["events"].take() + + def alive(self, job): + return job["alive"] + + def release(self, job): + job["thread"].join(timeout=30) + self.released.append(job["incarnation"]) + return not job["thread"].is_alive() + + def release_incarnation(self, incarnation): + self.released.append(incarnation) + return True + + +def _setup(tmp_path): + store_root = tmp_path / "store" + (store_root).mkdir() + (store_root / FLAT_MANIFEST).write_text("{}") # flat layout: pointer in the store root + sampler_buffer = FileSampleBuffer(CheckpointStore(store_root), batch_size=BATCH) + islands = _FakeIslands(store_root, sampler_buffer) + files = LocalStoreFiles(CheckpointStore(store_root)) + cuts = StoreCutLog(files) + lifecycle = RemoteTrainerLifecycle(islands, files, poll_s=POLL, exit_grace_s=0.2) + mirror = MirrorStore(tmp_path / "head" / "mirror", files, [f"{BUFFER_DIR}/{INDEX}"]) + head_buffer = CommittedBacklog(FileSampleBuffer(mirror, batch_size=BATCH), + lambda: consumed_watermark(cuts, lifecycle)) + return islands, lifecycle, cuts, head_buffer + + +def _loop(islands, lifecycle, cuts, head_buffer, clock, journal): + scheduler = S1Scheduler(S1Config(a_hard=0, start_seconds=T_S), record=journal.append, clock=clock) + return TrainerLoop(scheduler=scheduler, lifecycle=lifecycle, buffer=head_buffer, cuts=cuts, + record=journal.append, release_orphan=islands.release_incarnation, + ready_timeout_s=30.0) + + +def _ticks(loop, n): + return [type(loop.tick()) for _ in range(n)] + + +def test_remote_start_stop_to_the_data_end(tmp_path): + clock = _Clock() + islands, lifecycle, cuts, head_buffer = _setup(tmp_path) + journal = Journal(tmp_path / "head" / "journal") + loop = _loop(islands, lifecycle, cuts, head_buffer, clock, journal) + sampler = FakeSampler(islands.sampler_buffer, islands.published) + + # incarnation 0: two batches -> start, train 2 rounds, idle stop (SCHEDULED) + assert _ticks(loop, 1) == [Hold] + sampler.produce(2 * BATCH) + assert head_buffer.backlog(0).batches == 2 + assert _ticks(loop, 3) == [Start, Hold, Hold] + assert loop.state is LifecycleState.TRAINING and loop.current_version() == 2 + assert head_buffer.backlog(2).samples == 0 # the position the running trainer reported + assert _ticks(loop, 1) == [Hold] + clock.now += T_S + assert _ticks(loop, 1) == [Stop] and loop.state is LifecycleState.STOPPED + handle = loop.handle + assert handle.leave_reason is LeaveReason.STOPPED and handle.released is True + cut0 = cuts.latest() + assert (cut0.trigger, cut0.policy_version, cut0.watermark) == (recovery.SCHEDULED, 2, 2 * BATCH) + assert head_buffer.backlog(2).samples == 0 # the trainer's committed watermark + + # incarnation 1: one batch plus a remainder of 2, then the data ends + sampler.load() + sampler.produce(BATCH + 2) + islands.sampler_buffer.mark_data_end() + assert _ticks(loop, 3) == [Start, Hold, Stop] # below Q_on, but the data ended + assert loop.finished and loop.handle.leave_reason is LeaveReason.FINISHED, ( + loop.handle.leave_reason, loop.handle.error, loop.handle.states) + assert loop.handle.start_version == 2 and loop.current_version() == 3 + assert _ticks(loop, 1) == [Hold] + + assert islands.launched == [(0, None), (1, cut0.cut_id)] + assert islands.released == [0, 1] + trained = [sid for r in islands.rollouts for batch in r.trained.values() for sid, _ in batch] + assert len(trained) == 3 * BATCH and len(set(trained)) == len(trained) + assert cuts.latest().trigger == recovery.MANUAL and head_buffer.backlog(3).samples == 2 + + stops = [r for r in journal.records if r["kind"] == LOOP_KIND and r["event"] == "stopped"] + assert [(r["reason"], r["trigger"], r["released"]) for r in stops] == [ + ("STOPPED", recovery.SCHEDULED, True), ("FINISHED", recovery.MANUAL, True)] + # every command reached the island in order + sent = json.loads((tmp_path / "store" / command_path(0)).read_text().splitlines()[-1]) + assert sent == {"seq": 3, "verb": "stop", "rounds": None, "reason": "STOPPED"} + journal.close() + + +def test_island_that_dies_is_failed_and_released(tmp_path): + clock = _Clock() + islands, lifecycle, cuts, head_buffer = _setup(tmp_path) + journal = Journal(tmp_path / "head" / "journal") + loop = _loop(islands, lifecycle, cuts, head_buffer, clock, journal) + FakeSampler(islands.sampler_buffer, islands.published).produce(BATCH) + islands.die_at = 0 + assert _ticks(loop, 1) == [Start] + assert loop.state is LifecycleState.STOPPED + assert loop.handle.leave_reason is LeaveReason.INFRA_FAILED and "without a stopped event" in repr(loop.handle.error) + assert islands.released == [0] + islands.die_at = None + assert _ticks(loop, 2) == [Start, Hold] # the next incarnation starts again + assert loop.handle.incarnation == 1 + loop.close() + assert islands.released == [0, 1] + journal.close() + + +def test_data_end_below_one_batch_finishes_without_a_trainer(tmp_path): + """S1 boundary: less than one batch at the data end -> stop, the loop ends.""" + clock = _Clock() + islands, lifecycle, cuts, head_buffer = _setup(tmp_path) + journal = Journal(tmp_path / "head" / "journal") + loop = _loop(islands, lifecycle, cuts, head_buffer, clock, journal) + FakeSampler(islands.sampler_buffer, islands.published).produce(BATCH - 1) + islands.sampler_buffer.mark_data_end() + assert _ticks(loop, 1) == [Stop] and loop.finished and islands.launched == [] + records = list(journal.records) + assert [(r["event"], r["remainder"]) for r in records if r["kind"] == LOOP_KIND] == [(FINISHED, BATCH - 1)] + journal.close() + + # a restarted head sees the run finished + journal = Journal(tmp_path / "head" / "journal") + loop = _loop(islands, lifecycle, cuts, head_buffer, clock, journal) + loop.restore(journal.records) + assert loop.finished and _ticks(loop, 1) == [Hold] + journal.close() + + +@pytest.mark.parametrize("state", [LifecycleState.COLD, LifecycleState.STOPPED, LifecycleState.TRAINING]) +def test_s1_stops_finished_below_one_batch_at_the_data_end(state): + s1 = S1Scheduler(S1Config(a_hard=3, start_seconds=T_S), clock=_Clock()) + short = BacklogStats(batches=0.5, samples=2, data_ended=True) + assert s1.decide(short, state) == Stop(LeaveReason.FINISHED) + assert s1.decide(short, LifecycleState.WARMING) == Hold() + # a full batch at the data end starts a cold trainer below Q_on = 4 + full = BacklogStats(batches=1.5, samples=6, data_ended=True) + assert s1.decide(full, LifecycleState.COLD) == Start() + + +def test_s1_counts_a_partial_batch_as_idle(): + clock = _Clock() + s1 = S1Scheduler(S1Config(a_hard=0, start_seconds=T_S), clock=clock) + partial = BacklogStats(batches=0.5, samples=2) + assert s1.decide(partial, LifecycleState.TRAINING) == Hold() + clock.now += T_S + assert s1.decide(partial, LifecycleState.TRAINING) == Stop(LeaveReason.STOPPED) + + +# ---------------------------------------------------------------- store pieces +def test_command_channel_round_trip(tmp_path): + files = LocalStoreFiles(CheckpointStore(tmp_path)) + writer, reader = CommandWriter(files, 7), CommandReader(files, 7) + assert reader.poll() == [] + writer.send("run", rounds=2) + assert [(c.seq, c.verb, c.rounds) for c in reader.poll()] == [(1, "run", 2)] + writer.send("stop", reason=LeaveReason.STOPPED) + assert [(c.seq, c.reason) for c in reader.poll()] == [(2, "STOPPED")] + assert reader.poll() == [] + with pytest.raises(ValueError): + writer.send("stop") + + +def test_store_cut_log_reads_the_snapshot_layout(tmp_path): + store = CheckpointStore(tmp_path / "store") + state = tmp_path / "state" + state.mkdir() + cut = recovery.ConsistentCut(checkpoint_id="r000004-i1", trigger=recovery.SCHEDULED, outer_step=None, + island_steps={"local": 4}, buffer_watermark=12) + (state / "round-cut.json").write_text(json.dumps( + {"cut_id": "r000004-i1", "next_rollout_id": 4, recovery.CUT_KEY: cut.to_record()})) + log = StoreCutLog(LocalStoreFiles(store)) + assert log.latest() is None and log.watermark() == 0 + write_snapshot(store, state, {"cut_id": "r000004-i1"}) + ref = log.latest() + assert (ref.cut_id, ref.trigger, ref.policy_version, ref.watermark) == ("r000004-i1", recovery.SCHEDULED, 4, 12) + assert log.watermark() == 12 + + +class _Volume: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + + def read_file(self, path): + if path not in self.files: + raise FileNotFoundError(path) + yield self.files[path] + + def batch_upload(self, force=False): + volume = self + + class _Batch: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def put_file(self, fh, path): + volume.files[path] = fh.read() + + return _Batch() + + +def test_modal_volume_files_use_the_store_prefix(): + volume = _Volume() + files = ModalVolumeFiles("vol", "/runs/a/", volume_factory=lambda name: volume) + assert files.read("x/y.json") is None + files.write("x/y.json", b"1") + assert volume.files == {"runs/a/x/y.json": b"1"} and files.read("x/y.json") == b"1" + + +class _ModalOps: + def __init__(self) -> None: + self.spawned: list = [] + self.cancelled: list[str] = [] + self.stuck: set[str] = set() + + def spawn(self, cfg): + self.spawned.append(cfg) + return f"fc-{len(self.spawned)}" + + def status(self, call_id): + return "RUNNING" if call_id not in self.cancelled or call_id in self.stuck else "FAILED" + + def cancel(self, call_id): + self.cancelled.append(call_id) + + +class _Cfg: + def __init__(self) -> None: + self.envs = {"BASE": "1"} + + +def test_modal_islands_spawn_per_incarnation_and_release_an_orphan(tmp_path, monkeypatch): + import yeto.launch.trainer_island as ti + + monkeypatch.setattr(ti, "replace", lambda cfg, **kw: type("Cfg", (), kw)()) + ops, tails = _ModalOps(), [] + islands = ModalTrainerIslands(ops, _Cfg(), name="run-a-l1", jobs_path=tmp_path / "jobs.json", + tail=lambda call_id, label, events: tails.append((call_id, label)), + sleep=lambda s: None, clock=lambda: 1700000000.25) + job = islands.launch(3, None) + # phase ``submitted`` of the incarnation is its spawn, not the build of the held task + assert ops.spawned[0].envs == {"BASE": "1", **island_env(3, None), + SUBMITTED_ENV: "1700000000.250"} + assert tails == [("fc-1", "run-a-l1-i3")] + job.events.feed("[run-a-l1-i3] " + event_echo.format_record( + {"island_id": 1, "time_unix": 0.0, "event": "rl_trainer_lifecycle", "incarnation": 3})) + job.events.feed("noise\n") + assert [e["incarnation"] for e in islands.events(job)] == [3] and islands.events(job) == [] + assert islands.alive(job) + + # a new head (fresh instance) releases incarnation 3 by its recorded call id + head2 = ModalTrainerIslands(ops, _Cfg(), name="run-a-l1", jobs_path=tmp_path / "jobs.json", + tail=lambda *a: None, sleep=lambda s: None) + assert head2.release_incarnation(3) is True and ops.cancelled == ["fc-1"] + assert head2.release_incarnation(9) is False + job2 = islands.launch(4, None) + ops.stuck.add(job2.call_id) + assert islands.release(job2) is False # the cloud never confirmed + + +# ---------------------------------------------------------------- launcher switch +def _args(tmp_path, monkeypatch, extra=()): + monkeypatch.setattr(runs, "RUNS_DIR", tmp_path / "runs") + parser = argparse.ArgumentParser() + launch_loop.add_arguments(parser) + args = parser.parse_args(["--rl-trainer-scheduler", "s1", *extra]) + args.cluster_prefix, args.rl_max_policy_age, args.rollout_batch_size = "run-a", 0, BATCH + args.gpu, args.rl_checkpoint_store = "modal:1xh100,nebius:1xh100", "modal-volume://vol/runs/a" + return args + + +def test_check_refuses_without_a_modal_trainer_island(tmp_path, monkeypatch): + with pytest.raises(SystemExit, match="needs --rl-trainer-island"): + launch_loop.check(_args(tmp_path, monkeypatch)) + with pytest.raises(SystemExit, match="only a Modal"): + launch_loop.check(_args(tmp_path, monkeypatch, ["--rl-trainer-island", "1"])) + args = _args(tmp_path, monkeypatch, ["--rl-trainer-island", "0"]) + launch_loop.check(args) + assert launch_loop.held_island(args) == 0 + args.rl_checkpoint_store = "s3://bucket" + with pytest.raises(SystemExit, match="modal-volume"): + launch_loop.check(args) + + +def test_modal_binding_reads_the_volume_and_backlog(tmp_path, monkeypatch): + args = _args(tmp_path, monkeypatch, ["--rl-trainer-island", "0"]) + store_root = tmp_path / "store" + (store_root / FLAT_MANIFEST).parent.mkdir(parents=True) + (store_root / FLAT_MANIFEST).write_text("{}") + FakeSampler(FileSampleBuffer(CheckpointStore(store_root), batch_size=BATCH), lambda: None).produce(6) + islands = ModalTrainerIslands(_ModalOps(), _Cfg(), name="l0", jobs_path=tmp_path / "j.json", + tail=lambda *a: None) + binding = modal_binding(args, TrainerIsland("l0", "modal"), journal_dir=launch_loop.journal_dir(args), + files=LocalStoreFiles(CheckpointStore(store_root)), islands=islands) + assert binding.buffer.backlog(0).samples == 6 and binding.cuts.latest() is None + assert binding.release_orphan == islands.release_incarnation + assert isinstance(binding.lifecycle(5), RemoteTrainerLifecycle) + assert binding.lifecycle(5).plan_start({}).incarnation == 5 diff --git a/tests/test_rl_trainer_scheduler.py b/tests/test_rl_trainer_scheduler.py new file mode 100644 index 000000000..f85a89310 --- /dev/null +++ b/tests/test_rl_trainer_scheduler.py @@ -0,0 +1,196 @@ +"""Scheduler S1 (rl-trainer-start-stop L5.1-L5.3): every scheduler scenario of the spec.""" + +from __future__ import annotations + +import pytest + +from yeto.rl.elastic_client import LeaveReason +from yeto.rl.engine.journal import Journal +from yeto.rl.engine.trainer_lifecycle import ( + AlwaysOnScheduler, + BacklogStats, + Hold, + LifecycleState, + Start, + Stop, + TrainerScheduler, +) +from yeto.rl.engine.trainer_scheduler import ( + SCHEDULER_STATE_KIND, + S1Config, + S1Scheduler, + last_start_seconds, +) + +COLD, TRAINING = LifecycleState.COLD, LifecycleState.TRAINING +EMPTY = BacklogStats(batches=0.0, samples=0) + + +def _q(batches: float) -> BacklogStats: + return BacklogStats(batches=batches, samples=int(batches * 4)) + + +class _Clock: + def __init__(self, t: float = 1000.0) -> None: + self.t = t + + def __call__(self) -> float: + return self.t + + +def test_s1_implements_the_port(): + assert isinstance(S1Scheduler(S1Config(a_hard=2, start_seconds=500.0)), TrainerScheduler) + + +def test_cap_by_version_limit(): + # T_s=500, T_b=44, alpha=0.2 -> Q_s1 = ceil(500*0.8/(0.2*44)) = 46 >= 45; cap A_hard+1 = 3 + s = S1Scheduler(S1Config(a_hard=2, start_seconds=500.0, alpha=0.2, batch_seconds=44.0)) + assert s.q_on(500.0) == 3 + assert s.decide(_q(3), COLD) == Start() + + +def test_alpha_below_cap_starts_at_q_s1(): + # Q_s1 = ceil(10*0.5/(0.5*10)) = 1 < A_hard + 1 + s = S1Scheduler(S1Config(a_hard=4, start_seconds=10.0, alpha=0.5, batch_seconds=10.0)) + assert s.decide(_q(0.5), COLD) == Hold() + assert s.decide(_q(1), COLD) == Start() + + +def test_default_threshold_without_alpha(): + s = S1Scheduler(S1Config(a_hard=2, start_seconds=500.0)) + assert s.decide(_q(2), COLD) == Hold() + assert s.decide(_q(3), COLD) == Start() + + +def test_alpha_has_no_default_and_needs_batch_seconds(): + assert S1Config(a_hard=2, start_seconds=1.0).alpha is None + with pytest.raises(ValueError, match="batch_seconds"): + S1Config(a_hard=2, start_seconds=1.0, alpha=0.2) + with pytest.raises(ValueError, match="alpha"): + S1Config(a_hard=2, start_seconds=1.0, alpha=0.0, batch_seconds=1.0) + + +def test_ski_rental_stop(): + clock = _Clock() + s = S1Scheduler(S1Config(a_hard=2, start_seconds=300.0), clock=clock) + assert s.decide(EMPTY, TRAINING) == Hold() # idle starts here + clock.t += 299.0 + assert s.decide(EMPTY, TRAINING) == Hold() + clock.t += 1.0 + assert s.decide(EMPTY, TRAINING) == Stop(LeaveReason.STOPPED) + + +def test_new_samples_reset_the_idle_time(): + clock = _Clock() + s = S1Scheduler(S1Config(a_hard=2, start_seconds=300.0), clock=clock) + s.decide(EMPTY, TRAINING) + clock.t += 200.0 + assert s.decide(_q(1), TRAINING) == Hold() + clock.t += 200.0 + assert s.decide(EMPTY, TRAINING) == Hold() # idle again from here, not from t=0 + clock.t += 299.0 + assert s.decide(EMPTY, TRAINING) == Hold() + + +def test_data_end_finishes_and_trains_the_remainder(): + s = S1Scheduler(S1Config(a_hard=2, start_seconds=300.0)) + ended = BacklogStats(batches=1.0, samples=4, data_ended=True) + assert s.decide(ended, COLD) == Start() # below Q_on, but nothing more will come + assert s.decide(ended, TRAINING) == Hold() + done = BacklogStats(batches=0.0, samples=0, data_ended=True) + assert s.decide(done, TRAINING) == Stop(LeaveReason.FINISHED) + # nothing left: a stopped (or never started) trainer also finishes the run (L5.3c) + assert s.decide(done, COLD) == Stop(LeaveReason.FINISHED) + + +def test_transition_states_hold_and_prices_are_ignored(): + s = S1Scheduler(S1Config(a_hard=0, start_seconds=0.0)) + for state in (LifecycleState.ACQUIRING, LifecycleState.JOINING, LifecycleState.RESTORING, + LifecycleState.WARMING, LifecycleState.DRAINING): + assert s.decide(EMPTY, state, {"h100": 99.0}) == Hold() + assert s.decide(_q(1), COLD, {"h100": 99.0}) == Start() + + +def test_start_once(): + s = AlwaysOnScheduler() + assert s.decide(EMPTY, COLD) == Start() + assert s.decide(EMPTY, TRAINING) == Hold() # an empty backlog does not stop it + done = BacklogStats(batches=0.0, samples=0, data_ended=True) + assert s.decide(done, TRAINING) == Stop(LeaveReason.FINISHED) + + +# ---------------------------------------------------------------- L5.2: T_s +def _summary(total_s, island=1): + return {"event": "trainer_start_summary", "island_id": island, "total_s": total_s} + + +def test_t_s_from_last_start_summary(): + events = [{"event": "trainer_start_phase", "phase": "submitted"}, _summary(120.0), + _summary(None), {"event": "trainer_start_phase", "phase": "first_step_done"}, + _summary(80.5)] + assert last_start_seconds(events) == 80.5 + clock = _Clock() + s = S1Scheduler(S1Config(a_hard=2, start_seconds=500.0), + start_events=lambda: events, clock=clock) + assert s.start_seconds() == 80.5 + s.decide(EMPTY, TRAINING) + clock.t += 80.5 + assert s.decide(EMPTY, TRAINING) == Stop(LeaveReason.STOPPED) + + +def test_t_s_falls_back_to_configuration(): + s = S1Scheduler(S1Config(a_hard=2, start_seconds=500.0), + start_events=lambda: [_summary(None)]) + assert s.start_seconds() == 500.0 + assert S1Scheduler(S1Config(a_hard=2, start_seconds=7.0)).start_seconds() == 7.0 + + +def test_fixed_t_s_overrides_the_measured_start_and_journals_both(): + records = [] + s = S1Scheduler(S1Config(a_hard=2, start_seconds=30.0, fixed_start_seconds=True), + start_events=lambda: [_summary(650.0)], + record=lambda kind, tx_id=None, **f: records.append(f), clock=_Clock()) + assert s.start_seconds() == 30.0 + s.decide(EMPTY, TRAINING) + assert records[-1]["inputs"]["start_seconds"] == 30.0 + assert records[-1]["inputs"]["measured_start_seconds"] == 650.0 + + +def test_measured_t_s_changes_q_s1(): + events: list[dict] = [] + s = S1Scheduler(S1Config(a_hard=10, start_seconds=500.0, alpha=0.5, batch_seconds=10.0), + start_events=lambda: events) + assert s.q_on(s.start_seconds()) == 11 # Q_s1 = 50 -> cap 11 + events.append(_summary(20.0)) + assert s.q_on(s.start_seconds()) == 2 + + +# ---------------------------------------------------------------- L5.3: journal +def test_head_restart_while_idle_continues_the_idle_time(tmp_path): + clock = _Clock() + cfg = S1Config(a_hard=2, start_seconds=300.0) + with Journal(tmp_path) as journal: + s = S1Scheduler(cfg, record=journal.append, clock=clock) + assert s.decide(EMPTY, TRAINING) == Hold() + clock.t += 100.0 + s.decide(EMPTY, TRAINING) + [rec] = [r for r in journal.records if r["kind"] == SCHEDULER_STATE_KIND] + assert rec["state"] == {"idle_since": 1000.0} + assert rec["inputs"]["start_seconds"] == 300.0 and rec["decision"] == {"action": "hold"} + + clock.t += 150.0 # the head is down; idle time t = 250 s at restart + with Journal(tmp_path) as journal: + s = S1Scheduler(cfg, record=journal.append, clock=clock) + assert s.restore(journal.records) + assert s.decide(EMPTY, TRAINING) == Hold() + clock.t += 50.0 # 300 s since idle began, not 50 s + assert s.decide(EMPTY, TRAINING) == Stop(LeaveReason.STOPPED) + last = [r for r in journal.records if r["kind"] == SCHEDULER_STATE_KIND][-1] + assert last["decision"] == {"action": "stop", "reason": "STOPPED"} + assert last["state"] == {"idle_since": None} + + +def test_restore_without_records_starts_fresh(): + s = S1Scheduler(S1Config(a_hard=2, start_seconds=300.0)) + assert not s.restore([{"kind": "auto_state", "state": {}}]) + assert s.idle_since is None diff --git a/tests/test_rl_verl_fully_async_64b.py b/tests/test_rl_verl_fully_async_64b.py index 986f64eed..145cd8a2f 100644 --- a/tests/test_rl_verl_fully_async_64b.py +++ b/tests/test_rl_verl_fully_async_64b.py @@ -63,17 +63,19 @@ def test_limit_drives_staleness_threshold_and_partial_rollout(): assert not any("staleness" in o for o in vconf.build_overrides(run)) -def test_verl_declares_stage_two_up_to_limit_one(): +def test_verl_declares_stage_two_up_to_limit_four(): from yeto.rl.adapters.verl.entry import verl_capabilities from yeto.rl.adapters.verl.policy_age import SUPPORT, policy_age_overrides - assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 1) + assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 4) assert policy_age_overrides(0) == () assert policy_age_overrides(1, groups_per_round=8) == ( "async_training.staleness_threshold=1.0", "async_training.partial_rollout=True", "async_training.trigger_parameter_sync_step=1", "async_training.require_batches=1") with pytest.raises(PolicyAgeError, match="supports up to stage 2"): - policy_age_overrides(2, groups_per_round=8) + policy_age_overrides(5, groups_per_round=8) + # L6.9: A_hard = 2 maps to staleness_threshold 2 (decoupled PPO needs >= 2) + assert "async_training.staleness_threshold=2.0" in policy_age_overrides(2, groups_per_round=8) sync_caps = verl_capabilities("sha256:" + "0" * 64) assert sync_caps.execution.max_policy_staleness == 0 and sync_caps.placements == {"colocated"} fa = verl_capabilities("sha256:" + "0" * 64, fully_async_limit=1) @@ -325,6 +327,9 @@ def test_driver_runs_fully_async_rounds_with_carry_over_and_discard(tmp_path): assert carry[2]["over_age_discarded_groups"] == 1 and carry[2]["over_age_discarded_tokens"] == 10 cutoff = [e for e in events if e["event"] == "rl_rollout_cutoff"] assert [e["rollout_id"] for e in cutoff] == [2] + gates = [e for e in events if e["event"] == "offpolicy_gate"] + assert [e["rollout_id"] for e in gates] == [0, 1, 2] + assert {(e["backend"], e["granularity"]) for e in gates} == {("verl", "trajectory")} trained = [e for e in events if e["event"] == "rl_round_trained"] assert [e["aborted_in_flight_groups"] for e in trained] == [0, 0, 1] # never null (async10) checks = [e for e in events if e["event"] == "rl_publication_check"] diff --git a/tests/test_rl_verl_fully_async_translate.py b/tests/test_rl_verl_fully_async_translate.py index 65d64557b..5b5e9699b 100644 --- a/tests/test_rl_verl_fully_async_translate.py +++ b/tests/test_rl_verl_fully_async_translate.py @@ -77,4 +77,4 @@ def test_resumed_calls_give_per_token_versions_and_logprobs(): def test_verl_declares_stage_two_after_6_4b(): from yeto.rl.adapters.verl.policy_age import SUPPORT - assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 1) + assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 4) diff --git a/tests/test_rl_verl_offpolicy_correction.py b/tests/test_rl_verl_offpolicy_correction.py new file mode 100644 index 000000000..e5ae7fb1a --- /dev/null +++ b/tests/test_rl_verl_offpolicy_correction.py @@ -0,0 +1,344 @@ +"""rl-offpolicy-correction-passthrough, verl side (L6.4, L6.6, L6.11 CPU part). + +CPU only: no verl install, no Ray. The parity tests (L6.6) load the correction +functions straight from the verl fork and the Miles checkout on this machine +and skip when either is absent. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from yeto.rl.adapters.verl import fully_async_round as far +from yeto.rl.adapters.verl import offpolicy_correction as voc +from yeto.rl.adapters.verl.fully_async_translate import VersionMap +from yeto.rl.engine import offpolicy_correction as oc +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec + +K = voc.PREFIX +DECOUPLED = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=2), + correction=CorrectionSpec(method="tis", lower=0.0, upper=2.0)) + + +def _spec(**correction) -> AlgorithmSpec: + return AlgorithmSpec(correction=CorrectionSpec(**correction)) + + +def _fields(mapping) -> set[str]: + return {w.field for w in mapping.warnings} + + +# ------------------------------------------------------------------ L6.4 / L6.7 mapping + + +def test_decoupled_ppo_writes_all_three_values_over_the_fully_async_defaults(): + m = voc.correction_mapping(DECOUPLED) + assert m.backend == "verl" and m.prox_anchor == oc.PROX_FIRST_LOCAL_STEP + assert m.values[K + "bypass_mode"] is False + assert m.values[K + "rollout_is"] == "token" + assert m.values[K + "rollout_is_threshold"] == 2.0 + assert {K + "bypass_mode=False", K + "rollout_is=token", + K + "rollout_is_threshold=2.0", K + "rollout_rs=null"} <= set(m.overrides) + assert m.warnings == () + + +def test_tis_lower_warns_and_runs_without_a_lower_bound(): + m = voc.correction_mapping(_spec(method="tis", lower=0.5, upper=2.0)) + (w,) = m.warnings + assert (w.field, w.backend, w.ran_with) == ("lower", "verl", 0.0) + assert m.values[K + "rollout_is_threshold"] == 2.0 + + +def test_icepop_maps_to_a_lower_upper_string_without_denominator_warning(): + m = voc.correction_mapping(_spec(method="tis", lower=0.5, upper=5.0, ratio_mask="icepop")) + assert m.values[K + "rollout_is_threshold"] == "0.5_5.0" + assert K + "rollout_is_threshold='0.5_5.0'" in m.overrides + assert m.warnings == () + zero = voc.correction_mapping(_spec(method="tis", lower=0.0, upper=5.0, ratio_mask="icepop")) + assert _fields(zero) == {"lower"} and zero.warnings[0].ran_with == voc.ICEPOP_MIN_LOWER + + +def test_mis_maps_to_rejection_sampling_and_warns_on_the_denominator(): + m = voc.correction_mapping(_spec(method="tis", lower=0.5, upper=2.0, ratio_mask="mis")) + assert m.values[K + "rollout_rs"] == "token_k1" + assert m.values[K + "rollout_rs_threshold"] == "0.5_2.0" + assert m.values[K + "rollout_is"] is None + assert _fields(m) == {"ratio_mask"} + geo = voc.correction_mapping(_spec(method="tis", lower=0.5, upper=2.0, ratio_mask="mis", + level="geometric")) + assert geo.values[K + "rollout_rs"] == "seq_mean_k1" and _fields(geo) == {"ratio_mask"} + seq = voc.correction_mapping(_spec(method="tis", lower=0.5, upper=2.0, ratio_mask="mis", + level="sequence")) + assert _fields(seq) == {"ratio_mask", "level"} + + +def test_unmappable_fields_warn(): + assert _fields(voc.correction_mapping(_spec(method="tis", lower=0.0, upper=2.0, veto=1e-4))) == {"veto"} + assert _fields(voc.correction_mapping(_spec(method="opsm", opsm_delta=0.1))) == {"method"} + norm = voc.correction_mapping(_spec(method="tis", lower=0.0, upper=2.0, batch_normalize=True)) + assert norm.values[K + "rollout_is_batch_normalize"] is True and norm.warnings == () + seq = voc.correction_mapping(_spec(method="tis", lower=0.0, upper=2.0, level="sequence")) + assert seq.values[K + "rollout_is"] == "sequence" and seq.warnings == () + + +def test_bypass_mode_true_is_refused_with_the_setting_named(): + with pytest.raises(oc.CorrectionConfigError, match="bypass_mode=True"): + voc.refuse_forbidden_overrides([K + "bypass_mode=true"]) + m = voc.correction_mapping(_spec(method="tis", lower=0.0, upper=2.0, use_rollout_logprobs=True)) + with pytest.raises(oc.CorrectionConfigError, match="pi_prox"): + voc.apply_to_overrides([], m) + voc.refuse_forbidden_overrides([K + "bypass_mode=False"]) + + +def test_apply_replaces_every_rollout_correction_entry_once(): + out = voc.apply_to_overrides(["a=1", K + "bypass_mode=True", K + "rollout_is=sequence", + "+" + K + "rollout_is_batch_normalize=True"], + voc.correction_mapping(DECOUPLED)) + keys = [o.split("=", 1)[0] for o in out] + assert len(keys) == len(set(keys)) and "a=1" in out + assert K + "bypass_mode=False" in out and K + "rollout_is=token" in out + + +def test_backend_role_feeds_the_run_manifest_and_the_cross_island_warning(): + assert oc.backend_mapping(DECOUPLED, "verl").prox_anchor == oc.PROX_FIRST_LOCAL_STEP + manifest = oc.run_correction_manifest(DECOUPLED, {"i0": "miles", "i1": "verl"}) + assert manifest["islands"]["i1"]["mapped"][K + "bypass_mode"] is False + (cross,) = [w for w in manifest["warnings"] if w.get("scope") == "cross_island"] + assert cross["differences"]["prox_anchor"] == {"i0": oc.PROX_PER_STEP, + "i1": oc.PROX_FIRST_LOCAL_STEP} + + +# ------------------------------------------------------------------ L6.11 fully_async path + + +@pytest.mark.parametrize("limit", [2, 4]) +def test_fully_async_plan_runs_the_correction_mapping(limit, tmp_path): + import json + + from yeto.rl.adapters.verl.island_entry import append_correction_record, fully_async_plan + + import yeto.rl.adapters.verl.config as vconf + + run = vconf.VerlRunConfig(model_path="/m", train_file="/t", val_file="/v", out_dir="/o", + groups_per_round=8, samples_per_group=4) + spec = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=limit), + correction=CorrectionSpec(method="tis", lower=0.5, upper=2.0)) + tape = tmp_path / "tape.jsonl" + args = SimpleNamespace(groups_per_round=8, samples_per_group=4, global_rounds=3, + event_tape=str(tape), learner_id=0) + out, keys = fully_async_plan(limit, vconf.build_overrides(run), spec.to_dict(), args, "none") + append_correction_record(args, spec.to_dict()) + append_correction_record(args, AlgorithmSpec().to_dict()) # default: tape unchanged + # D8: the island tape records the group, the verl values and prox_anchor + events = [json.loads(line) for line in tape.read_text().splitlines()] + assert events[0]["event"] == oc.CORRECTION_RECORD_EVENT and events[0]["backend"] == "verl" + assert events[0]["prox_anchor"] == oc.PROX_FIRST_LOCAL_STEP + assert events[0]["group"]["upper"] == 2.0 + assert [(e["event"], e["field"]) for e in events[1:]] == [(oc.CORRECTION_CONFIG_WARNING, "lower")] + assert f"async_training.staleness_threshold={float(limit)}" in out + assert "async_training.partial_rollout=True" in out + assert [o for o in out if o.startswith(K + "bypass_mode=")] == [K + "bypass_mode=False"] + assert K + "rollout_is=token" in out and K + "rollout_is_threshold=2.0" in out + assert K + "bypass_mode" in keys + + +def _vmap(n: int) -> VersionMap: + vm = VersionMap() + for v in range(n): + vm.record(v, v) + return vm + + +def _meta(uid, lo, hi, lengths): + n = len(lengths) + return far.sample_meta({"uid": [uid] * n, "min_global_steps": [lo] * n, + "max_global_steps": [hi] * n}, lengths) + + +def test_round_collection_drops_whole_groups_over_a_hard_and_feeds_the_gate(): + rc = far.RoundCollector(_vmap(6), current_outer=5, limit=2, required=2) + assert rc.offer(_meta("fresh", 5, 5, [10, 20])) + assert not rc.offer(_meta("old", 2, 5, [300, 5000])) # oldest v2: gap 3 > A_hard 2 + assert not rc.offer(_meta("unk", None, None, [7, 7])) + assert rc.offer(_meta("edge", 3, 4, [1, 1])) # gap 2 == A_hard: kept + fields = rc.carry_over_fields() + assert fields["over_age_discarded_groups"] == 1 + # the driver builds offpolicy_gate from these keys and removes them + gate = oc.gate_from_carry_over(fields, kept=len(rc.kept), backend="verl") + assert oc.GATE_DROPPED_LENGTHS not in fields and oc.GATE_GRANULARITY not in fields + assert gate["backend"] == "verl" and gate["granularity"] == oc.GRANULARITY_TRAJECTORY + assert gate["counts"] == {oc.GATE_KEPT: 2, oc.GATE_GROUP_DROPPED: 1, + oc.GATE_UNKNOWN_VERSION_DROPPED: 1} + assert gate["dropped_length_histogram"][oc.GATE_GROUP_DROPPED] == \ + oc.length_histogram([300, 5000]) + assert gate["truncation_by_turn"] is None + + +# ------------------------------------------------------------------ L6.6 CPU parity + + +VERL_SRC = Path(os.environ.get("YETO_VERL_SRC", "/home/michael/work/verl-fork-acad9875")) +MILES_SRC = Path(os.environ.get("YETO_MILES_SRC", "/home/michael/work/miles-next")) +VERL_HELPER = VERL_SRC / "verl/trainer/ppo/rollout_corr_helper.py" +MILES_CORRECTIONS = MILES_SRC / "miles/backends/training_utils/loss_hub/corrections.py" + + +def _load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def engines(): + if not (VERL_HELPER.is_file() and MILES_CORRECTIONS.is_file()): + pytest.skip("verl fork or Miles checkout not on this machine") + # rollout_corr_helper imports verl_F (masked_sum / masked_mean: copied below + # from verl/utils/torch_functional.py at acad9875), DataProto and two config + # classes (type hints only). Stubs replace the heavy package imports. + stub = types.ModuleType("verl.utils.torch_functional") + + def masked_sum(values, mask, axis=None): + return (torch.where(mask.bool(), values, 0.0) * mask).sum(axis=axis) + + def masked_mean(values, mask, axis=None): + return masked_sum(values, mask, axis) / (mask.sum(axis=axis) + 1e-8) + + stub.masked_sum, stub.masked_mean = masked_sum, masked_mean + + class DataProto: + @classmethod + def from_dict(cls, tensors): + obj = cls() + obj.batch = tensors + return obj + + names = {"verl": types.ModuleType("verl"), "verl.utils": types.ModuleType("verl.utils"), + "verl.utils.torch_functional": stub, + "verl.protocol": types.SimpleNamespace(DataProto=DataProto), + "verl.trainer.config.algorithm": types.SimpleNamespace(RolloutCorrectionConfig=object), + "verl.workers.config.actor": types.SimpleNamespace(PolicyLossConfig=object)} + names["verl"].utils = names["verl.utils"] + names["verl.utils"].torch_functional = stub + saved = {k: sys.modules.get(k) for k in names} + sys.modules.update(names) + try: + verl = _load("_yeto_parity_verl_rollout_corr", VERL_HELPER) + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + miles = _load("_yeto_parity_miles_corrections", MILES_CORRECTIONS) + return verl, miles + + +def _logprobs(seed: int = 3): + g = torch.Generator().manual_seed(seed) + rollout = -torch.rand(4, 16, generator=g, dtype=torch.float64) * 3 + train = rollout + torch.randn(4, 16, generator=g, dtype=torch.float64) * 0.8 + mask = torch.ones(4, 16, dtype=torch.float64) + mask[1, 12:] = 0 + mask[3, 5:] = 0 + return train, rollout, mask + + +def _miles_weights(miles, fn, train, rollout, mask, lower, upper): + rows = [m.bool() for m in mask] + flat_t = [t[r] for t, r in zip(train, rows)] + flat_r = [t[r] for t, r in zip(rollout, rows)] + n = sum(int(r.sum()) for r in rows) + args = SimpleNamespace(tis_clip=upper, tis_clip_low=lower) + pg, masks, _ = getattr(miles, fn)(args, pg_loss=torch.ones(n, dtype=torch.float64), + train_log_probs=flat_t, rollout_log_probs=flat_r, + loss_masks=[r.double() for r in rows]) + return pg, masks + + +def _verl(verl, train, rollout, mask, **kw): + weights, new_mask, _ = verl.compute_rollout_correction_and_rejection_mask( + old_log_prob=train, rollout_log_prob=rollout, response_mask=mask, **kw) + return (None if weights is None else weights.batch["rollout_is_weights"][mask.bool()]), new_mask + + +def test_parity_tis_with_lower_zero_matches(engines): + verl, miles = engines + train, rollout, mask = _logprobs() + miles_w, _ = _miles_weights(miles, "vanilla_tis_function", train, rollout, mask, 0.0, 2.0) + verl_w, new_mask = _verl(verl, train, rollout, mask, rollout_is="token", rollout_is_threshold=2.0) + torch.testing.assert_close(verl_w, miles_w) + assert torch.equal(new_mask, mask) + + +def test_parity_tis_lower_differs_as_warned(engines): + verl, miles = engines + train, rollout, mask = _logprobs() + miles_w, _ = _miles_weights(miles, "vanilla_tis_function", train, rollout, mask, 0.5, 2.0) + verl_w, _ = _verl(verl, train, rollout, mask, rollout_is="token", rollout_is_threshold=2.0) + low = miles_w > verl_w + assert low.any() and torch.all(miles_w[low] == 0.5) + torch.testing.assert_close(verl_w[~low], miles_w[~low]) + + +def test_parity_icepop_matches_and_keeps_the_denominator(engines): + verl, miles = engines + train, rollout, mask = _logprobs() + miles_w, miles_masks = _miles_weights(miles, "icepop_function", train, rollout, mask, 0.5, 2.0) + verl_w, new_mask = _verl(verl, train, rollout, mask, rollout_is="token", + rollout_is_threshold="0.5_2.0") + torch.testing.assert_close(verl_w, miles_w) + assert (verl_w == 0).any() + assert torch.equal(new_mask, mask) # verl IcePop zeroes weights; loss_mask unchanged + assert sum(int(m.sum()) for m in miles_masks) == int(mask.sum()) + + +def test_verl_rejection_sampling_shrinks_the_denominator(engines): + verl, _ = engines + train, rollout, mask = _logprobs() + _, new_mask = _verl(verl, train, rollout, mask, rollout_rs="token_k1", + rollout_rs_threshold="0.5_2.0") + assert 0 < new_mask.sum() < mask.sum() + ratio = torch.exp(train - rollout) + kept = mask.bool() & (ratio >= 0.5) & (ratio <= 2.0) + assert torch.equal(new_mask.bool(), kept) + + +def test_verl_use_rollout_log_probs_is_not_read_by_any_code(): + """L6.6: ``actor.use_rollout_log_probs`` (fully_async yaml: True) has no + reader in the fork, so with ``bypass_mode=False`` pi_prox is recomputed.""" + if not VERL_SRC.is_dir(): + pytest.skip("verl fork not on this machine") + readers = [p for p in (VERL_SRC / "verl").rglob("*.py") + if "use_rollout_log_probs" in p.read_text(errors="ignore") + and p.name != "actor.py"] + assert readers == [] + + +# ------------------------------------------------------------------ example configs + + +@pytest.mark.parametrize("name,a_hard", [("decoupled-ppo-c2", 2), ("decoupled-ppo-c2-agentic", 4)]) +def test_decoupled_ppo_examples_map_on_both_backends(name, a_hard): + from yeto.rl.adapters.miles.offpolicy_correction import admission_argv + from yeto.rl.adapters.miles.offpolicy_correction import correction_mapping as miles_mapping + from yeto.rl.adapters.verl.policy_age import policy_age_overrides + + examples = Path(__file__).resolve().parents[1] / "examples" / "rl_algorithms" + spec = AlgorithmSpec.from_json_file(str(examples / f"{name}.json")) + assert spec.rejections() == [] and spec.execution.max_policy_staleness == a_hard + assert miles_mapping(spec).argv == ("--use-tis", "--tis-clip", "2.0", "--tis-clip-low", "0.0") + assert admission_argv(a_hard) == ("--max-weight-staleness", str(a_hard)) + verl = voc.correction_mapping(spec) + assert verl.values[K + "rollout_is_threshold"] == 2.0 and verl.warnings == () + assert f"async_training.staleness_threshold={float(a_hard)}" in \ + policy_age_overrides(a_hard, groups_per_round=8) diff --git a/tests/test_rl_verl_policy_age.py b/tests/test_rl_verl_policy_age.py index 055c66913..9f02ad767 100644 --- a/tests/test_rl_verl_policy_age.py +++ b/tests/test_rl_verl_policy_age.py @@ -59,11 +59,11 @@ def test_verl_declares_stage_two_and_limit_zero_overrides_are_empty(): from yeto.rl.adapters.verl.entry import verl_capabilities from yeto.rl.adapters.verl.policy_age import SUPPORT, policy_age_overrides - assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 1) # 6.4b: fully_async path + assert (SUPPORT.stage, SUPPORT.max_policy_age) == (2, 4) # 6.4b; L6.9 raised the cap assert verl_capabilities("sha256:" + "0" * 64).execution.max_policy_staleness == 0 assert policy_age_overrides(0) == () with pytest.raises(PolicyAgeError, match="backend 'verl' supports up to stage 2"): - policy_age_overrides(2, groups_per_round=4) + policy_age_overrides(5, groups_per_round=4) def test_launcher_accepts_verl_limit_one(): diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index e5c46d185..ccdefa57b 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -264,7 +264,7 @@ def test_forward_tables_name_flags_of_both_layers(): "--password-stdin", "--policy-sweep-fragments", "--porcelain", "--port", "--profile", "--prompt", "--q-min", "--query-gpu", "--quiet", "--quorum-theta", "--quorum-timeout-s", "--raid-devices", - "--release", "--require-profile-binding", "--resume", "--revision", "--rl-eval", "--rl-test", "--rm", "--run", + "--release", "--require-profile-binding", "--resume", "--rl-eval", "--rl-test", "--rm", "--run", "--runtime", "--seed-column", "--shm-size", "--soft-deadline-s", "--syncer-epoch", "--temp-dir", "--ulimit", "--untracked-files", "--wandb-group", "--yes", }) diff --git a/tests/test_trainer_cold_start.py b/tests/test_trainer_cold_start.py new file mode 100644 index 000000000..be261ab5a --- /dev/null +++ b/tests/test_trainer_cold_start.py @@ -0,0 +1,343 @@ +"""rl-trainer-cold-start-cost L0 (start phases), L2.3/L2.4 (image manifest +instead of imports; peft in the image) and L2.5 (model prefetch). CPU only.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time +import zipfile +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from yeto import rl +from yeto.launch import images, manifest as launch_manifest +from yeto.launch.remote_setup import model_prefetch +from yeto.rl.engine import start_phases as sp + +REPO = Path(__file__).resolve().parents[1] + + +class _Collect: + def __init__(self): + self.events: list[tuple[str, dict]] = [] + + def __call__(self, event, **fields): + self.events.append((event, fields)) + + def named(self, name): + return [f for e, f in self.events if e == name] + + +# ------------------------------------------------------------ L0.1 phases + + +def _fake_start(tmp_path, *, with_shell=True): + """A fake launcher + setup + run: real shell marks, then the island replay.""" + + home = tmp_path / "home" + home.mkdir() + env = {**os.environ, "HOME": str(home)} + lines = [sp.shell_reset(), sp.shell_mark("container_ready"), "sleep 0.05", + sp.shell_mark("setup_done"), "sleep 0.05", sp.shell_mark("ray_ready")] + submitted = time.time() + if with_shell: + subprocess.run(["bash", "-c", "\n".join(lines)], env=env, check=True) + emit = _Collect() + phases = sp.StartPhases(emit) + phases.replay({sp.SUBMITTED_ENV: f"{submitted:.3f}"}, home / ".yeto-start-phases") + return phases, emit + + +def test_fake_launcher_start_has_seven_phases_in_order_and_the_total(tmp_path): + phases, emit = _fake_start(tmp_path) + phases.mark("actor_import_done") + phases.mark("actor_init_done") + phases.mark_once("first_step_done") + phases.mark_once("first_step_done") # later steps record nothing + records = emit.named(sp.START_PHASE_EVENT) + assert [r["phase"] for r in records] == list(sp.START_PHASES) + assert [r["index"] for r in records] == list(range(7)) + times = [r["t_unix"] for r in records] + assert times == sorted(times) + [summary] = emit.named(sp.START_SUMMARY_EVENT) + assert summary["missing"] == [] + assert summary["total_s"] == pytest.approx(times[-1] - times[0], abs=1e-3) + assert summary["phases"] == dict(zip(sp.START_PHASES, times)) + + +def test_phase_order_is_enforced(): + phases = sp.StartPhases(_Collect(), clock=lambda: 1.0) + phases.mark("submitted") + phases.mark("setup_done") # a skipped phase is allowed ... + with pytest.raises(ValueError, match="out of order"): + phases.mark("container_ready") # ... going back is not + with pytest.raises(ValueError, match="twice"): + phases.mark("setup_done") + with pytest.raises(ValueError, match="unknown"): + phases.mark("warm") + + +def test_missing_marks_are_named_in_the_summary(tmp_path): + phases, emit = _fake_start(tmp_path, with_shell=False) + phases.missing("actor_import_done", "no Ray record") + phases.mark("actor_init_done") + phases.mark("first_step_done") + [summary] = emit.named(sp.START_SUMMARY_EVENT) + assert summary["missing"] == ["container_ready", "setup_done", "ray_ready", "actor_import_done"] + assert summary["total_s"] is not None + + +def test_trainer_actor_import_time_from_ray_state(): + from yeto.rl.adapters.miles.entry import trainer_actor_import_done_unix + + tasks = [SimpleNamespace(name="MegatronTrainRayActor.__init__", end_time_ms=2000), + SimpleNamespace(name="MegatronTrainRayActor.__init__", end_time_ms=3500), + SimpleNamespace(name="RolloutManager.__init__", end_time_ms=9000)] + assert trainer_actor_import_done_unix(lambda **kw: tasks) == 3.5 + assert trainer_actor_import_done_unix(lambda **kw: []) is None + + def broken(**kw): + raise RuntimeError("no dashboard") + + assert trainer_actor_import_done_unix(broken) is None + + +def test_island_task_scripts_carry_the_marks(monkeypatch): + from tests.test_rl_engine_selection import _cli, _island_task + from yeto.launcher import _prepare_rl_args + + args = _cli(["--gpu", "nebius:2x1xl40s@eu-north1"]) + _prepare_rl_args(args) + task = _island_task(args, monkeypatch) + setup = task.setup + assert setup.index(sp.shell_reset()) < setup.index(sp.shell_mark("container_ready")) + assert setup.index(sp.shell_mark("container_ready")) < setup.index("SGLANG_IN_IMAGE=0") + assert setup.rstrip().endswith(sp.shell_mark("setup_done")) + run = task.run + assert run.index("ray start --head") < run.index(sp.shell_mark("ray_ready")) < run.index("python3 -m ") + assert abs(float(task.envs[sp.SUBMITTED_ENV]) - time.time()) < 600 + + +# ------------------------------------------------------- L0.2 run manifest + + +def test_summary_goes_into_the_run_manifest_and_t_s_reads_back(tmp_path, monkeypatch): + from yeto import runs + + run_dir = tmp_path / "run" + (run_dir / "events").mkdir(parents=True) + (run_dir / "run_manifest.json").write_text(json.dumps({"rl_engine": "ports"})) + monkeypatch.setattr(runs, "run_dir", lambda prefix: run_dir) + for island, total in ((0, 81.5), (1, 120.25)): + tape = run_dir / "events" / f"island-{island}.jsonl" + lines = [{"event": "rl_driver_start", "island_id": island}, + {"event": sp.START_SUMMARY_EVENT, "island_id": island, "total_s": total, + "missing": [], "phases": {"submitted": 1.0}}] + tape.write_text("".join(json.dumps(l) + "\n" for l in lines) + "{torn") + summaries = launch_manifest._record_trainer_start( + SimpleNamespace(cluster_prefix="r"), run_dir / "events") + assert set(summaries) == {"0", "1"} + manifest = json.loads((run_dir / "run_manifest.json").read_text()) + assert manifest["rl_engine"] == "ports" + assert manifest[sp.MANIFEST_KEY]["1"]["total_s"] == 120.25 + assert sp.t_s_from_manifest(manifest) == 120.25 + assert sp.t_s_from_manifest({"rl_engine": "ports"}) is None + + +def test_no_summary_leaves_the_manifest_alone(tmp_path, monkeypatch): + from yeto import runs + + (tmp_path / "events").mkdir() + monkeypatch.setattr(runs, "run_dir", lambda prefix: tmp_path) + assert launch_manifest._record_trainer_start(SimpleNamespace(cluster_prefix="r"), + tmp_path / "events") is None + assert not (tmp_path / "run_manifest.json").exists() + + + +def test_sky_islands_record_t_s_from_teardown_diagnostics(tmp_path, monkeypatch): + """Non-echo sky islands: T_s comes from the island tapes the pre-teardown + diagnostics fetched into // (other tapes ignored).""" + from yeto import runs + + monkeypatch.setattr(runs, "run_dir", lambda prefix: tmp_path) + diag = tmp_path / "diag" + for cluster, island, total in (("r-0", 0, 90.0), ("r-1", 1, 140.5)): + (diag / cluster).mkdir(parents=True) + (diag / cluster / f"rl-island-{island}.jsonl").write_text(json.dumps( + {"event": sp.START_SUMMARY_EVENT, "island_id": island, "total_s": total, + "missing": [], "phases": {}}) + "\n") + (diag / "r-1" / "yeto-tape.jsonl").write_text(json.dumps( + {"event": sp.START_SUMMARY_EVENT, "island_id": 7, "total_s": 1.0}) + "\n") + args = SimpleNamespace(cluster_prefix="r") + assert set(launch_manifest._record_sky_trainer_start(args, ["r-0", "r-1"], str(diag))) == {"0", "1"} + manifest = json.loads((tmp_path / "run_manifest.json").read_text()) + assert sp.t_s_from_manifest(manifest) == 140.5 + assert launch_manifest._record_sky_trainer_start(args, ["r-9"], str(diag)) is None + + +# --------------------------------------------- L2.3 manifest, not imports + + +def _run_check(tmp_path, manifest, package="peft", key="version", want=rl.MILES_PEFT_VERSION, + path=""): + """Run image_manifest_check against ``manifest`` (dict, or None = no file) + with a python3 shim that refuses to import peft/sglang.""" + + home = tmp_path / "home" + home.mkdir(exist_ok=True) + manifest_path = tmp_path / "image-manifest.json" + if manifest is not None: + manifest_path.write_text(json.dumps(manifest)) + poison = tmp_path / "poison" + for name in ("peft", "sglang"): + (poison / name).mkdir(parents=True, exist_ok=True) + (poison / name / "__init__.py").write_text( + f"import pathlib; pathlib.Path({str(tmp_path / 'imported')!r}).write_text({name!r})\n" + "__version__ = 'poison'\n__file__ = '/nowhere'\n") + script = images.image_manifest_check( + package, key, want, path=path, on_pass="echo PASS", + fallback="echo FALLBACK", + ).replace(rl.MILES_NEXT_IMAGE_MANIFEST, str(manifest_path)) + env = {**os.environ, "HOME": str(home), "PYTHONPATH": str(poison)} + proc = subprocess.run(["bash", "-c", "set -e\n" + script + "\necho AFTER"], env=env, + text=True, capture_output=True) + events = home / ".yeto-start-phases" / sp.SETUP_EVENTS_FILE + return proc, (sp.read_tape_records([events]) if events.exists() else []), \ + (tmp_path / "imported").exists() + + +def test_manifest_present_passes_without_importing(tmp_path): + proc, events, imported = _run_check( + tmp_path, {"peft": {"version": rl.MILES_PEFT_VERSION}}) + assert proc.returncode == 0, proc.stderr + assert "PASS" in proc.stdout and "FALLBACK" not in proc.stdout + assert events == [] and not imported + + +def test_sglang_manifest_checks_commit_and_path(tmp_path): + want = {"sglang": {"commit": rl.SGLANG_NEXT_COMMIT, "path": "/sgl-workspace/sglang"}} + proc, _, imported = _run_check(tmp_path, want, "sglang", "commit", rl.SGLANG_NEXT_COMMIT, + "/sgl-workspace/sglang") + assert proc.returncode == 0 and "PASS" in proc.stdout and not imported + want["sglang"]["path"] = "/elsewhere" + proc, _, _ = _run_check(tmp_path, want, "sglang", "commit", rl.SGLANG_NEXT_COMMIT, + "/sgl-workspace/sglang") + assert proc.returncode != 0 and "AFTER" not in proc.stdout + assert "sglang path /elsewhere != /sgl-workspace/sglang" in proc.stderr + + +@pytest.mark.parametrize("manifest", [None, {"miles": {"commit": "x"}}], ids=["no-file", "no-entry"]) +def test_manifest_missing_falls_back_with_an_event(tmp_path, manifest): + proc, events, _ = _run_check(tmp_path, manifest) + assert proc.returncode == 0, proc.stderr + assert "FALLBACK" in proc.stdout and "PASS" not in proc.stdout + assert [e["event"] for e in events] == ["image_manifest_fallback"] + assert events[0]["package"] == "peft" and events[0]["check"] == "import" + assert '"event": "image_manifest_fallback"' in proc.stderr + + +def test_manifest_mismatch_fails_naming_package_and_both_versions(tmp_path): + proc, _, imported = _run_check(tmp_path, {"peft": {"version": "0.13.2"}}) + assert proc.returncode == 1 + assert "AFTER" not in proc.stdout and "PASS" not in proc.stdout + assert f"peft version 0.13.2 != pinned {rl.MILES_PEFT_VERSION}" in proc.stderr + assert not imported + + +def test_setup_events_reach_the_tape_once(tmp_path): + proc, _, _ = _run_check(tmp_path, None) + emit = _Collect() + phases = sp.StartPhases(emit) + mark_dir = tmp_path / "home" / ".yeto-start-phases" + phases.replay({}, mark_dir) + assert phases.forward_setup_events(mark_dir) == 0 + [event] = emit.named("image_manifest_fallback") + assert event["source"] == "setup" and event["package"] == "peft" + + +# ------------------------------------------------- L2.4 peft in the image + + +def test_build_script_puts_peft_into_the_image_and_the_manifest(): + script = (REPO / "scripts" / "build_miles_ports_image.sh").read_text() + assert '"MILES_PEFT_VERSION"' in script + assert 'download -q --no-deps --only-binary=:all:' in script + assert '"peft": {"version": "$MILES_PEFT_VERSION"' in script + assert '--wheel "$PEFT_WHEEL"' in script + dockerfile = (REPO / "docker" / "miles-ports" / "Dockerfile").read_text() + assert f"ARG PEFT_VERSION={rl.MILES_PEFT_VERSION}" in dockerfile + + +def _wheel(path: Path, files: dict[str, bytes]) -> Path: + with zipfile.ZipFile(path, "w") as archive: + for name, data in files.items(): + archive.writestr(name, data) + return path + + +def test_layer_unpacks_a_pure_python_wheel(tmp_path): + sys.path.insert(0, str(REPO / "docker" / "miles-ports")) + try: + import make_layer + finally: + sys.path.pop(0) + wheel = _wheel(tmp_path / "peft-0.20.0-py3-none-any.whl", { + "peft/__init__.py": b"__version__ = '0.20.0'\n", + "peft-0.20.0.dist-info/METADATA": b"Name: peft\n", + "peft-0.20.0.dist-info/RECORD": b"", + }) + files = make_layer.wheel_files(wheel) + assert files["peft/__init__.py"].startswith(b"__version__") + assert files["peft-0.20.0.dist-info/INSTALLER"] == b"yeto-make-layer\n" + with pytest.raises(SystemExit, match="pure-Python"): + make_layer.wheel_files(_wheel(tmp_path / "x-1-cp312-cp312-linux_x86_64.whl", {})) + with pytest.raises(SystemExit, match=".data"): + make_layer.wheel_files(_wheel(tmp_path / "y-1-py3-none-any.whl", + {"y-1.data/scripts/y": b"", "y-1.dist-info/RECORD": b""})) + + +# --------------------------------------------------------- L2.5 prefetch + + +def _run_prefetch(tmp_path, exit_code: int): + home = tmp_path / "home" + home.mkdir() + shim = tmp_path / "bin" + shim.mkdir() + (shim / "python3").write_text(f"#!/bin/sh\necho \"$@\" > {tmp_path / 'argv'}\nexit {exit_code}\n") + (shim / "python3").chmod(0o755) + command = model_prefetch("org/model", "a" * 40) + env = {**os.environ, "HOME": str(home), "PATH": f"{shim}:{os.environ['PATH']}"} + proc = subprocess.run(["bash", "-c", f"set -e\n{command}\necho CONTINUES"], env=env, + text=True, capture_output=True, timeout=30) + events_file = home / ".yeto-start-phases" / sp.SETUP_EVENTS_FILE + deadline = time.time() + 10 + while time.time() < deadline and not (tmp_path / "argv").exists(): + time.sleep(0.05) + time.sleep(0.3) # the background shell writes the event after python3 exits + return proc, sp.read_tape_records([events_file]), (tmp_path / "argv").read_text() + + +def test_prefetch_uses_snapshot_download_not_the_removed_cli(): + command = model_prefetch("org/model", "a" * 40) + assert "huggingface-cli" not in command and "|| true" not in command.split("sh -c")[0] + assert "snapshot_download(repo_id=" in command and "org/model" in command and "a" * 40 in command + + +def test_prefetch_failure_is_recorded_and_the_setup_continues(tmp_path): + proc, events, argv = _run_prefetch(tmp_path, 5) + assert proc.returncode == 0 and "CONTINUES" in proc.stdout + assert "snapshot_download" in argv + assert events == [{"event": "model_prefetch_failed", "exit_code": 5, + "log": "/tmp/hf-prefetch.log", "model": "org/model"}] + + +def test_prefetch_success_records_nothing(tmp_path): + proc, events, _ = _run_prefetch(tmp_path, 0) + assert proc.returncode == 0 and events == [] diff --git a/yeto/cli.py b/yeto/cli.py index 89b736861..b0ab2ee83 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -603,6 +603,9 @@ def int_or_auto(value: str): default=30, help="fleet-controller health poll interval (seconds)", ) + from .launch import trainer_loop as _trainer_loop + + _trainer_loop.add_arguments(infra) # W&B telemetry. One group per fleet (this run's name), one run per # island plus one for the syncer's event tape; see docs/WANDB.md. diff --git a/yeto/codepack.py b/yeto/codepack.py new file mode 100644 index 000000000..abe86f691 --- /dev/null +++ b/yeto/codepack.py @@ -0,0 +1,493 @@ +"""Run code packs (rl-identity-tiering L9.4-L9.6, development mode). + +A code pack fixes the code of one run. The launcher makes it once, at the first +start of the run, from the local git work tree: + +* ``commit``: ``git rev-parse HEAD``; +* the dirty patch: ``git diff --binary --cached HEAD`` over a temporary index that + holds every tracked and every untracked, not ignored file (the real index is not + touched); ``patch_sha`` is its sha256, ``dirty`` is true when it is not empty; +* a deterministic tar of the resulting tree (sorted entries, zero times and owners, + modes reduced to 0644/0755), with :data:`PACK_INFO` as the first entry; +* ``digest``: the sha256 of the tar bytes. + +The tar is uploaded to the run's checkpoint store under ``code_packs/.tar`` +and recorded in the run manifest. Every later start of the run (resume, island +replacement, re-JOIN) takes the same pack from the manifest, unpacks it into a local +cache and uses that directory as the island workdir. The island checks the unpacked +tree against the digest before it starts (:func:`verify_tree`) and refuses to start +on a mismatch. The pack changes only through :func:`upgrade_record` (explicit +upgrade): it requires an equal semantic hash and appends "pack A -> pack B" to the +manifest record. + +Storage backends are pluggable (:func:`register_pack_store`); this module holds the +neutral core and a directory backend. Cloud backends live with their cloud code. +""" + +from __future__ import annotations + +import argparse +import fnmatch +import hashlib +import io +import json +import os +import shutil +import subprocess +import tarfile +import tempfile +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path, PurePosixPath +from typing import Any, Callable, Iterable, Mapping, Protocol + +CODE_PACK_MODE_GIT_PATCH = "git-patch" +PACK_INFO = "CODEPACK.json" +PACK_DIR = "code_packs" +PACK_FORMAT = 1 +# Never packed: the same paths the Modal workdir mount drops, plus Python caches. +PACK_EXCLUDE = (".git", ".venv", ".claude", "syncer/target", "RLinf-main", + "**/__pycache__", "**/*.pyc") + + +class CodePackError(RuntimeError): + """A pack cannot be made, found, or does not match its digest.""" + + +# --------------------------------------------------------------------------- +# pack record +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class CodePack: + """What the run manifest records for one pack.""" + + commit: str + dirty: bool + patch_sha: str + digest: str + pack_uri: str | None = None + mode: str = CODE_PACK_MODE_GIT_PATCH + upgrades: tuple[dict[str, Any], ...] = field(default_factory=tuple) + + def record(self) -> dict[str, Any]: + out = asdict(self) + out["upgrades"] = [dict(u) for u in self.upgrades] + return out + + @classmethod + def from_record(cls, record: Mapping[str, Any]) -> "CodePack": + if record.get("mode") != CODE_PACK_MODE_GIT_PATCH: + raise CodePackError(f"code pack mode {record.get('mode')!r} is not {CODE_PACK_MODE_GIT_PATCH!r}") + return cls(commit=str(record["commit"]), dirty=bool(record["dirty"]), + patch_sha=str(record["patch_sha"]), digest=str(record["digest"]), + pack_uri=record.get("pack_uri"), mode=str(record["mode"]), + upgrades=tuple(dict(u) for u in record.get("upgrades") or ())) + + +# --------------------------------------------------------------------------- +# building +# --------------------------------------------------------------------------- +def _git(repo: Path, *argv: str, env: Mapping[str, str] | None = None) -> bytes: + try: + return subprocess.run(["git", *argv], cwd=repo, check=True, capture_output=True, + env=None if env is None else {**os.environ, **env}).stdout + except (OSError, subprocess.CalledProcessError) as e: + detail = getattr(e, "stderr", b"") or b"" + raise CodePackError(f"git {' '.join(argv)} failed in {repo}: " + f"{detail.decode(errors='replace').strip() or e}") from e + + +def _excluded(path: str) -> bool: + parts = PurePosixPath(path).parts + for pattern in PACK_EXCLUDE: + if pattern.startswith("**/"): + leaf = pattern[3:] + if any(fnmatch.fnmatch(part, leaf) for part in parts): + return True + elif path == pattern or path.startswith(pattern + "/"): + return True + return False + + +def pack_file_list(repo: str | os.PathLike[str]) -> list[str]: + """Tracked and untracked, not ignored files that exist now, sorted, minus + :data:`PACK_EXCLUDE`.""" + + raw = _git(Path(repo), "ls-files", "-z", "--cached", "--others", "--exclude-standard") + names = {name for name in raw.decode("utf-8").split("\0") if name} + root = Path(repo) + return sorted(n for n in names + if not _excluded(n) and ((root / n).is_symlink() or (root / n).is_file())) + + +def dirty_patch(repo: str | os.PathLike[str]) -> bytes: + """``HEAD`` → work tree as one binary patch, untracked (not ignored) files + included. Uses a temporary index; the repository index is not changed.""" + + root = Path(repo) + git_dir = Path(_git(root, "rev-parse", "--absolute-git-dir").decode().strip()) + with tempfile.TemporaryDirectory() as tmp: + index = Path(tmp) / "index" + if (git_dir / "index").is_file(): + shutil.copyfile(git_dir / "index", index) + env = {"GIT_INDEX_FILE": str(index)} + _git(root, "add", "-A", "--", ".", env=env) + return _git(root, "diff", "--binary", "--cached", "--no-color", "--no-ext-diff", + "HEAD", "--", ".", env=env) + + +def _tar_bytes(root: Path, files: Iterable[str], info: Mapping[str, Any]) -> bytes: + """The deterministic tar of ``files`` under ``root``, :data:`PACK_INFO` first.""" + + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w", format=tarfile.USTAR_FORMAT) as tar: + def add(name: str, *, data: bytes | None = None, link: str | None = None, + mode: int = 0o644) -> None: + entry = tarfile.TarInfo(name) + entry.mtime = 0 + entry.uid = entry.gid = 0 + entry.uname = entry.gname = "" + entry.mode = mode + if link is not None: + entry.type = tarfile.SYMTYPE + entry.linkname = link + tar.addfile(entry) + else: + entry.size = len(data or b"") + tar.addfile(entry, io.BytesIO(data or b"")) + + add(PACK_INFO, data=(json.dumps(info, sort_keys=True, indent=1) + "\n").encode()) + for name in files: + path = root / name + if path.is_symlink(): + add(name, link=os.readlink(path), mode=0o777) + else: + executable = os.stat(path).st_mode & 0o111 + add(name, data=path.read_bytes(), mode=0o755 if executable else 0o644) + return buf.getvalue() + + +def build_pack(repo: str | os.PathLike[str]) -> tuple[CodePack, bytes]: + """Make the pack of the work tree at ``repo``: ``(record without URI, tar bytes)``.""" + + root = Path(repo) + commit = _git(root, "rev-parse", "HEAD").decode().strip() + patch = dirty_patch(root) + files = pack_file_list(root) + patch_sha = hashlib.sha256(patch).hexdigest() + info = {"format": PACK_FORMAT, "mode": CODE_PACK_MODE_GIT_PATCH, "commit": commit, + "dirty": bool(patch), "patch_sha": patch_sha, "files": files} + data = _tar_bytes(root, files, info) + pack = CodePack(commit=commit, dirty=bool(patch), patch_sha=patch_sha, + digest=hashlib.sha256(data).hexdigest()) + return pack, data + + +# --------------------------------------------------------------------------- +# checking and unpacking +# --------------------------------------------------------------------------- +def read_pack_info(tree: str | os.PathLike[str]) -> dict[str, Any]: + path = Path(tree) / PACK_INFO + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as e: + raise CodePackError(f"{path}: no readable code pack info ({e})") from e + + +def tree_digest(tree: str | os.PathLike[str]) -> str: + """Rebuild the tar from an unpacked tree (the files its :data:`PACK_INFO` lists) + and return its sha256. A missing or changed listed file changes the digest.""" + + root = Path(tree) + info = read_pack_info(root) + files = list(info.get("files") or []) + missing = [n for n in files if not ((root / n).is_symlink() or (root / n).is_file())] + if missing: + raise CodePackError(f"{root}: {len(missing)} packed file(s) missing, e.g. {missing[0]}") + return hashlib.sha256(_tar_bytes(root, files, info)).hexdigest() + + +def verify_tree(tree: str | os.PathLike[str], expected: str) -> str: + """Refuse a tree whose rebuilt tar does not have the digest ``expected``.""" + + actual = tree_digest(tree) + if actual != expected.lower(): + raise CodePackError(f"code pack digest mismatch: run manifest has {expected.lower()}, " + f"{tree} has {actual}") + return actual + + +def _safe_member(member: tarfile.TarInfo) -> None: + path = PurePosixPath(member.name) + if path.is_absolute() or ".." in path.parts or not (member.isfile() or member.issym()): + raise CodePackError(f"unsafe code pack entry {member.name!r}") + if member.issym(): + target = PurePosixPath(member.linkname) + if target.is_absolute() or ".." in (path.parent / target).parts: + raise CodePackError(f"code pack symlink {member.name!r} leaves the tree") + + +def unpack(data: bytes, digest: str, cache_root: str | os.PathLike[str]) -> Path: + """Check ``data`` against ``digest`` and unpack it to ``/`` + (reused when already there and intact). Returns the tree.""" + + if hashlib.sha256(data).hexdigest() != digest.lower(): + raise CodePackError(f"downloaded code pack does not have digest {digest}") + dest = Path(cache_root).expanduser() / digest.lower() + if dest.is_dir(): + try: + verify_tree(dest, digest) + return dest + except CodePackError: + shutil.rmtree(dest) + dest.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{digest[:12]}-", dir=dest.parent)) + try: + with tarfile.open(fileobj=io.BytesIO(data), mode="r:") as tar: + members = tar.getmembers() + for member in members: + _safe_member(member) + tar.extractall(staging, members=members, filter="tar") + verify_tree(staging, digest) + staging.rename(dest) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + return dest + + +# --------------------------------------------------------------------------- +# storage +# --------------------------------------------------------------------------- +class PackStore(Protocol): + def put(self, digest: str, data: bytes) -> str: + """Store the pack; return its URI.""" + + def get(self, uri: str) -> bytes: + """Return the bytes stored at ``uri``.""" + + +class DirPackStore: + """A checkpoint store the launcher sees as a directory (local path, a shared + file system mounted on the launcher too, or a test double).""" + + def __init__(self, root: str | os.PathLike[str], uri_base: str | None = None) -> None: + self.root = Path(root).expanduser() + self.uri_base = (uri_base or str(self.root)).rstrip("/") + + def _path(self, uri: str) -> Path: + if not uri.startswith(self.uri_base + "/"): + raise CodePackError(f"{uri!r} is not in store {self.uri_base!r}") + return self.root / uri[len(self.uri_base) + 1:] + + def put(self, digest: str, data: bytes) -> str: + uri = f"{self.uri_base}/{PACK_DIR}/{digest}.tar" + path = self._path(uri) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".tar.tmp") + tmp.write_bytes(data) + os.replace(tmp, path) + return uri + + def get(self, uri: str) -> bytes: + try: + return self._path(uri).read_bytes() + except OSError as e: + raise CodePackError(f"code pack {uri} not readable: {e}") from e + + +_PACK_STORES: dict[str, Callable[[str], PackStore]] = {} + + +def register_pack_store(scheme: str, factory: Callable[[str], PackStore]) -> None: + """Register the pack store for checkpoint store URIs ``://...``. The + factory receives the whole store URI.""" + + _PACK_STORES[scheme] = factory + + +def pack_store_for(store_uri: str) -> PackStore: + """The pack store for ``--rl-checkpoint-store``. A path must be visible on the + launcher (a share mounted there too); a ``scheme://`` URI needs a registered + backend.""" + + value = str(store_uri) + if "://" in value: + scheme = value.split("://", 1)[0] + if scheme not in _PACK_STORES: + _register_builtin_stores() + if scheme not in _PACK_STORES: + raise CodePackError(f"no code pack store for {scheme}:// (store {value!r})") + return _PACK_STORES[scheme](value) + root = Path(value).expanduser() + if not root.is_dir(): + raise CodePackError(f"checkpoint store path {value!r} is not visible on the launcher; " + "the code pack needs a bucket, a modal-volume:// store, or the " + "shared file system mounted on the launcher too") + return DirPackStore(root, uri_base=value.rstrip("/")) + + +class _S3PackStore: + """``s3://bucket[/prefix]`` (boto3; also S3-compatible endpoints via the usual + AWS environment).""" + + def __init__(self, store_uri: str) -> None: + self.base = store_uri.rstrip("/") + + @staticmethod + def _split(uri: str) -> tuple[str, str]: + bucket, _, key = uri.split("://", 1)[1].partition("/") + return bucket, key + + def put(self, digest: str, data: bytes) -> str: + import boto3 + + uri = f"{self.base}/{PACK_DIR}/{digest}.tar" + bucket, key = self._split(uri) + boto3.client("s3").put_object(Bucket=bucket, Key=key, Body=data) + return uri + + def get(self, uri: str) -> bytes: + import boto3 + + bucket, key = self._split(uri) + return boto3.client("s3").get_object(Bucket=bucket, Key=key)["Body"].read() + + +class _ModalVolumePackStore: + """``modal-volume://NAME[/PREFIX]`` (Modal SDK).""" + + def __init__(self, store_uri: str) -> None: + from .cloud.modal_ckpt_store import parse_modal_volume_uri + + self.base = store_uri.rstrip("/") + self.name, self.prefix = parse_modal_volume_uri(store_uri) + + def _volume(self, *, create: bool = False): + import modal + + return modal.Volume.from_name(self.name, create_if_missing=create) + + def put(self, digest: str, data: bytes) -> str: + # The first start stores the pack before the island creates the run's + # checkpoint volume (modal_runner, create_if_missing), so create it here too. + rel = "/".join(p for p in (self.prefix, PACK_DIR, f"{digest}.tar") if p) + with self._volume(create=True).batch_upload(force=True) as batch: + batch.put_file(io.BytesIO(data), "/" + rel) + return f"{self.base}/{PACK_DIR}/{digest}.tar" + + def get(self, uri: str) -> bytes: + rel = uri[len(self.base):].lstrip("/") + rel = "/".join(p for p in (self.prefix, rel) if p) + return b"".join(self._volume().read_file("/" + rel)) + + +def _register_builtin_stores() -> None: + _PACK_STORES.setdefault("s3", _S3PackStore) + _PACK_STORES.setdefault("modal-volume", _ModalVolumePackStore) + + +# --------------------------------------------------------------------------- +# run level +# --------------------------------------------------------------------------- +def make_and_store(repo: str | os.PathLike[str], store: PackStore) -> tuple[CodePack, bytes]: + pack, data = build_pack(repo) + uri = store.put(pack.digest, data) + return CodePack(commit=pack.commit, dirty=pack.dirty, patch_sha=pack.patch_sha, + digest=pack.digest, pack_uri=uri), data + + +def fetch_tree(pack: CodePack, store: PackStore, cache_root: str | os.PathLike[str]) -> Path: + """Download the manifest's pack and unpack it (checked) into the cache.""" + + if not pack.pack_uri: + raise CodePackError("code pack record has no pack_uri") + return unpack(store.get(pack.pack_uri), pack.digest, cache_root) + + +def upgrade_record(old: CodePack, new: CodePack, *, old_semantic: str | None, + new_semantic: str | None, reason: str, + now: float | None = None) -> CodePack: + """Explicit upgrade (L9.6): allowed only when both semantic hashes are known and + equal. Returns ``new`` carrying the old history plus "old → new".""" + + if not reason.strip(): + raise CodePackError("a code pack upgrade needs a reason") + if not old_semantic or not new_semantic: + raise CodePackError("code pack upgrade refused: the semantic hash of the old or the " + "new algorithm spec is unknown") + if old_semantic != new_semantic: + raise CodePackError(f"code pack upgrade refused: semantic hash changes " + f"({old_semantic} -> {new_semantic}); start a new run instead") + step = {"from": old.digest, "to": new.digest, "from_commit": old.commit, + "to_commit": new.commit, "semantic_sha256": new_semantic, "reason": reason, + "unix": time.time() if now is None else now} + return CodePack(commit=new.commit, dirty=new.dirty, patch_sha=new.patch_sha, + digest=new.digest, pack_uri=new.pack_uri, + upgrades=(*old.upgrades, step)) + + +@dataclass(frozen=True) +class RunPack: + pack: CodePack + tree: Path + reused: bool + + +def resolve_run_pack(*, previous: Mapping[str, Any] | None, repo: str | os.PathLike[str], + store: PackStore, cache_root: str | os.PathLike[str], + semantic_sha256: str | None, upgrade_reason: str | None = None, + log: Callable[[str], None] = print) -> RunPack: + """The pack this start of the run uses. + + * no pack in ``previous`` (the old run manifest): make, store and unpack one; + * a pack in ``previous``: take that same pack (resume / replacement), even when + the local tree changed since; + * with ``upgrade_reason``: make a pack from the local tree and switch to it + only under :func:`upgrade_record`'s rule.""" + + old_record = (previous or {}).get("code_pack") + if not old_record: + pack, data = make_and_store(repo, store) + log(f"[codepack] new run pack {pack.digest[:16]} (commit {pack.commit[:12]}" + f"{', dirty' if pack.dirty else ''}) -> {pack.pack_uri}") + return RunPack(pack, unpack(data, pack.digest, cache_root), reused=False) + old = CodePack.from_record(old_record) + if upgrade_reason is not None: + candidate, data = make_and_store(repo, store) + if candidate.digest == old.digest: + log("[codepack] upgrade requested but the local tree equals the run pack") + return RunPack(old, fetch_tree(old, store, cache_root), reused=True) + new = upgrade_record(old, candidate, old_semantic=(previous or {}).get("semantic_sha256"), + new_semantic=semantic_sha256, reason=upgrade_reason) + log(f"[codepack] upgrade {old.digest[:16]} -> {new.digest[:16]}: {upgrade_reason}") + return RunPack(new, unpack(data, new.digest, cache_root), reused=False) + log(f"[codepack] reusing run pack {old.digest[:16]} from the run manifest") + return RunPack(old, fetch_tree(old, store, cache_root), reused=True) + + +# --------------------------------------------------------------------------- +# CLI (island setup / operators) +# --------------------------------------------------------------------------- +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="python -m yeto.codepack") + sub = parser.add_subparsers(dest="cmd", required=True) + verify = sub.add_parser("verify", help="check an unpacked tree against a digest") + verify.add_argument("--dir", default=".") + verify.add_argument("--digest", required=True) + build = sub.add_parser("digest", help="print the pack digest of a git work tree") + build.add_argument("--repo", default=".") + args = parser.parse_args(argv) + try: + if args.cmd == "verify": + print(verify_tree(args.dir, args.digest)) + else: + print(build_pack(args.repo)[0].digest) + except CodePackError as e: + print(f"codepack: {e}") + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/yeto/launch/code_pack.py b/yeto/launch/code_pack.py new file mode 100644 index 000000000..7cb26d273 --- /dev/null +++ b/yeto/launch/code_pack.py @@ -0,0 +1,86 @@ +"""Launch side of the run code pack (rl-identity-tiering L9.4-L9.6). + +With ``--rl-code-pack git-patch`` the launcher resolves the run's pack before it +writes the run manifest: the first start makes and stores one; every later start +(resume, a new launcher process for the same run) takes the pack the old manifest +names. The unpacked pack becomes the workdir of every island (sky ``workdir`` and the +Modal ``add_local_dir``), so a local tree that changes after the first start does not +reach replacement or resumed islands. Islands check the tree against +``--code-pack-digest`` before they start. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from . import REPO_ROOT + +CODE_PACK_OFF = "off" +CODE_PACK_CACHE = Path("~/.cache/yeto/codepacks") + + +def code_workdir(args) -> str: + """The directory shipped as the island workdir: the run pack tree, else the repo.""" + + return str(getattr(args, "code_pack_tree", None) or REPO_ROOT) + + +def code_pack_island_flag(args) -> str: + pack = getattr(args, "code_pack", None) + return f" --code-pack-digest {pack.digest}" if pack is not None else "" + + +def _previous_manifest(cluster_prefix: str) -> dict | None: + from .. import runs + + path = runs.run_dir(cluster_prefix) / "run_manifest.json" + try: + return json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return None + + +def _semantic_sha256(args) -> str | None: + raw = getattr(args, "rl_algorithm_spec_json", None) + if not raw: + return None + from ..rl.engine.algorithm import AlgorithmSpec + + return AlgorithmSpec.from_dict(json.loads(raw)).semantic_sha256() + + +def prepare_run_code_pack(args, *, store=None, repo: str | Path | None = None, + cache_root: str | Path | None = None) -> None: + """Resolve the run pack and point the launch at it (sets ``args.code_pack``, + ``args.code_pack_tree`` and re-attests ``args.source_sha256`` on the pack tree). + No-op unless ``--rl-code-pack git-patch`` on an RL launch.""" + + mode = getattr(args, "rl_code_pack", None) or CODE_PACK_OFF + upgrade = getattr(args, "rl_code_pack_upgrade", None) + if mode == CODE_PACK_OFF or getattr(args, "training_mode", "sft") != "rl": + if upgrade is not None: + raise ValueError("--rl-code-pack-upgrade requires --rl-code-pack git-patch") + return + from .. import codepack + from ..provenance import source_tree_sha256 + + if store is None: + store_uri = getattr(args, "rl_checkpoint_store", None) + if not store_uri: + raise ValueError("--rl-code-pack git-patch requires --rl-checkpoint-store " + "(the pack is stored next to the run's checkpoints)") + store = codepack.pack_store_for(store_uri) + resolved = codepack.resolve_run_pack( + previous=_previous_manifest(args.cluster_prefix), repo=repo or REPO_ROOT, + store=store, cache_root=Path(cache_root or CODE_PACK_CACHE).expanduser(), + semantic_sha256=_semantic_sha256(args), upgrade_reason=upgrade) + # sky syncs a workdir minus its .gitignore matches; the pack holds no ignored + # untracked file, and .skyignore (when present) replaces .gitignore, so a + # force-added ignored file still ships. The extra file is not part of the digest. + skyignore = resolved.tree / ".skyignore" + if not skyignore.exists(): + skyignore.write_text("/.git\n", encoding="utf-8") + args.code_pack = resolved.pack + args.code_pack_tree = str(resolved.tree) + args.source_sha256 = source_tree_sha256(resolved.tree / "yeto") diff --git a/yeto/launch/fleet.py b/yeto/launch/fleet.py index ec6cdfcf6..e6cf979e4 100644 --- a/yeto/launch/fleet.py +++ b/yeto/launch/fleet.py @@ -76,8 +76,10 @@ class FixedRosterIslandAbandoned(RuntimeError): # Modal moved an island to a different container mid-run (preempted/rescheduled and the # run script re-ran): the run is failed and the Modal app stopped. Registered exit codes: # 1 failed, 2 artifact not fetchable, 3 incomplete tape, 4 island failed, 5 teardown -# unverified, 6 stalled, 7 Modal container id changed. +# unverified, 6 stalled, 7 Modal container id changed, 8 trainer scheduler gave up. CONTAINER_CHANGED_EXIT = 7 +# The trainer scheduler stopped after too many failed trainer incarnations in a row. +TRAINER_GAVE_UP_EXIT = 8 class FleetController: @@ -171,6 +173,7 @@ def __init__( # writes it): once it exists no island is relaunched any more. self.stop_flag = stop_flag self.poll_interval = poll_interval + self._abort_error: BaseException | None = None # set by abort() self.recover_timeout = recover_timeout self.on_relaunch = on_relaunch self.thread_cls = thread_cls @@ -244,15 +247,23 @@ def _make_record(name, task, job_id): "exit": None, } + def abort(self, error: BaseException) -> None: + """Ask :meth:`run` (any thread) to end with ``error`` at its next poll, e.g. + the trainer scheduler gave up; the launcher then tears the run down.""" + self._abort_error = error + def run(self) -> dict: """Poll until every learner is done or abandoned. Returns {learner name: final status string}; raises RuntimeError - (after downing the syncer) if every learner was abandoned. + (after downing the syncer) if every learner was abandoned, and the + :meth:`abort` error when one was set. """ for name in self.learners: self._fleet("island_ready", name, initial=True) while True: + if self._abort_error is not None: + raise self._abort_error self._fleet("cost_tick") if self.syncer is not None: self._poll(self.syncer, is_syncer=True) diff --git a/yeto/launch/images.py b/yeto/launch/images.py index dc0e9bf21..9d9ebdff8 100644 --- a/yeto/launch/images.py +++ b/yeto/launch/images.py @@ -109,6 +109,59 @@ def parse_image_spec(value: str | None): return images +# Exit codes of the image-manifest check (:func:`image_manifest_check`). +MANIFEST_MATCH, MANIFEST_ABSENT, MANIFEST_MISMATCH = 0, 3, 4 +# Stdlib-only: reading the manifest must not import the packages it describes +# (``import peft`` ~27 s, ``import sglang`` ~18 s; coldstart-gpu-review 6.1). +_MANIFEST_CHECK_PY = """\ +import json, os, sys +path, package, key, want, want_path = sys.argv[1:6] +try: + entry = json.load(open(path)).get(package) or {} +except FileNotFoundError: + sys.exit(3) +have = entry.get(key) +if have is None: + sys.exit(3) +if have != want: + print(f"[yeto-setup] image manifest {path}: {package} {key} {have} " + f"!= pinned {want}", file=sys.stderr) + sys.exit(4) +if want_path and os.path.realpath(entry.get("path", "")) != os.path.realpath(want_path): + print(f"[yeto-setup] image manifest {path}: {package} path {entry.get('path')} " + f"!= {want_path}", file=sys.stderr) + sys.exit(4) +""" + + +def image_manifest_check(package: str, key: str, want: str, *, fallback: str, + path: str = "", on_pass: str = ":") -> str: + """Shell lines: check ``package`` against the image manifest, without importing it. + + Manifest entry equal to ``want`` (and at ``path`` when given): run + ``on_pass``. Manifest or entry absent: record ``image_manifest_fallback`` + and run ``fallback`` (the old import check). Different value: fail the + setup (before Ray starts), naming the package and both values. + """ + + from ..rl import MILES_NEXT_IMAGE_MANIFEST + from ..rl.engine.start_phases import shell_event + + argv = " ".join(shlex.quote(a) for a in (MILES_NEXT_IMAGE_MANIFEST, package, key, want, path)) + event = shell_event("image_manifest_fallback", + {"package": package, "manifest": MILES_NEXT_IMAGE_MANIFEST, + "check": "import"}) + return ( + f"rc=0; python3 -c {shlex.quote(_MANIFEST_CHECK_PY)} {argv} || rc=$?\n" + 'case "$rc" in\n' + f"{MANIFEST_MATCH}) echo '[yeto-setup] image manifest: {package} {key} {want}'; {on_pass} ;;\n" + f"{MANIFEST_ABSENT})\n{event}\n{fallback} ;;\n" + f"*) echo \"[yeto-setup] image manifest check of {package} failed (exit $rc); " + "refusing to start\" >&2; exit 1 ;;\n" + "esac" + ) + + def _miles_source_setup(overlay: str | None = None) -> tuple[str, str]: """Return the (miles_setup, sglang_setup) remote steps of the pinned Miles/SGLang forks. @@ -159,14 +212,17 @@ def checkout( # PYTHONPATH) becomes a link to it. Any other image, or an existing # ~/sglang directory, gets the clone + install at ~/sglang as before. image_root = shlex.quote(MILES_NEXT_IMAGE_SGLANG_ROOT) - sglang_in_image = ( + sglang_checkout_in_image = ( f'{{ [ ! -e ~/sglang ] || [ "$(readlink ~/sglang)" = {image_root} ]; }} && ' f'[ "$(git -C {image_root} rev-parse HEAD 2>/dev/null)" = ' f"{SGLANG_NEXT_COMMIT} ] && " f'[ "$(git -C {image_root} config --get remote.origin.url)" = ' f"{shlex.quote(SGLANG_NEXT_REPOSITORY)} ] && " f'[ -z "$(git -C {image_root} status --porcelain ' - '--untracked-files=all)" ] && ' + '--untracked-files=all)" ]' + ) + # Fallback for an image without a manifest entry: the old import check. + sglang_import_ok = ( "python3 -c 'import os, sys, sglang; sys.exit(0 if os.path.realpath(" f'sglang.__file__).startswith("{MILES_NEXT_IMAGE_SGLANG_ROOT}/python/' "\") else 1)'" @@ -193,6 +249,15 @@ def checkout( f"sys.exit(0 if peft.__version__ == \"{MILES_PEFT_VERSION}\" else 1)' " "2>/dev/null" ) + peft_check = image_manifest_check( + "peft", "version", MILES_PEFT_VERSION, + fallback=f"{peft_ok} || python3 -m pip install -q 'peft=={MILES_PEFT_VERSION}'", + ) + sglang_check = image_manifest_check( + "sglang", "commit", SGLANG_NEXT_COMMIT, path=MILES_NEXT_IMAGE_SGLANG_ROOT, + on_pass="SGLANG_IN_IMAGE=1", + fallback=f"if {sglang_import_ok}; then SGLANG_IN_IMAGE=1; fi", + ) return ( "set -e\n" "MILES_REFRESHED=0\n" @@ -202,13 +267,17 @@ def checkout( + f"if {miles_in_image}; then\n" f"echo \"[yeto-setup] image provides miles {MILES_NEXT_COMMIT} " f'({manifest}); editable install kept"\n' - f"{peft_ok} || python3 -m pip install -q 'peft=={MILES_PEFT_VERSION}'\n" + f"{peft_check}\n" "else\n" "python3 -m pip install -q --no-deps -e ~/miles " f"'peft=={MILES_PEFT_VERSION}'\n" "fi" + (f"\n{_miles_overlay().overlay_setup(overlay)}" if overlay else ""), - f"if {sglang_in_image}; then\n" + "SGLANG_IN_IMAGE=0\n" + f"if {sglang_checkout_in_image}; then\n" + f"{sglang_check}\n" + "fi\n" + 'if [ "$SGLANG_IN_IMAGE" = 1 ]; then\n' f"ln -sfn {image_root} ~/sglang\n" f"echo '[yeto-setup] image provides sglang {SGLANG_NEXT_COMMIT}'\n" "else\n" diff --git a/yeto/launch/learner_task.py b/yeto/launch/learner_task.py index 2d9dd0b38..64fdcc711 100644 --- a/yeto/launch/learner_task.py +++ b/yeto/launch/learner_task.py @@ -6,10 +6,10 @@ import shlex from ..gpu_spec import ClusterSpec -from . import REPO_ROOT +from .code_pack import code_workdir from .remote_setup import ( HF_TOKEN_ENV, HF_TOKEN_PATH, HUNYUAN3D_ROOT, HUNYUAN3D_SETUP, NVME_ENV, NVME_SETUP, - PROTENIX_SETUP, TORCH_SETUP, WAN_TUNING, causal_kernel_setup_steps, + PROTENIX_SETUP, TORCH_SETUP, WAN_TUNING, causal_kernel_setup_steps, model_prefetch, ) from .topology import multinode_network_tier from .loss import pickled_loss_path @@ -374,15 +374,7 @@ def make_learner_task(args, spec: ClusterSpec, learner_id: int, num_learners: in elif is_local_reference(repo): prefetch = ": # local model; no Hub prefetch" else: - revision_flag = ( - f" --revision {shlex.quote(args.model_revision)}" - if getattr(args, "model_revision", None) - else "" - ) - prefetch = ( - f"(nohup huggingface-cli download {shlex.quote(repo)}{revision_flag} " - ">/tmp/hf-prefetch.log 2>&1 &) || true" - ) + prefetch = model_prefetch(repo, getattr(args, "model_revision", None)) run = ( f"{NVME_ENV}\n" f"{HF_TOKEN_ENV}\n" @@ -400,7 +392,7 @@ def make_learner_task(args, spec: ClusterSpec, learner_id: int, num_learners: in envs=envs, **({"secrets": env_secrets} if env_secrets else {}), num_nodes=spec.num_nodes, - workdir=str(REPO_ROOT), + workdir=code_workdir(args), file_mounts=file_mounts, ) infra = f"{spec.cloud}/{spec.region}" if spec.region else spec.cloud diff --git a/yeto/launch/manifest.py b/yeto/launch/manifest.py index c70972e6f..54ba6b85d 100644 --- a/yeto/launch/manifest.py +++ b/yeto/launch/manifest.py @@ -29,6 +29,27 @@ def _island_tasks_manifest(args) -> dict: return {"island_tasks": {str(island): data.record() for island, data in sorted(plan.items())}} +def _code_identity_manifest(args) -> dict: + """rl-identity-tiering L9.1 dual write: ``semantic_sha256`` of the prepared + algorithm spec and ``code_pack_digest`` (transition: the source tree hash). + ``semantic_sha256`` is absent when no spec was prepared.""" + + from ..rl.engine.algorithm import AlgorithmSpec + from ..provenance import CODE_PACK_MODE_SOURCE_TREE + + pack = getattr(args, "code_pack", None) + if pack is not None: # L9.4: the run code pack replaces the transition value + out = {"code_pack_digest": pack.digest, "code_pack_mode": pack.mode, + "code_pack": pack.record()} + else: + out = {"code_pack_digest": getattr(args, "source_sha256", None), + "code_pack_mode": CODE_PACK_MODE_SOURCE_TREE} + raw = getattr(args, "rl_algorithm_spec_json", None) + if raw: + out["semantic_sha256"] = AlgorithmSpec.from_dict(json.loads(raw)).semantic_sha256() + return out + + def _write_run_manifest(args) -> dict | None: """Record the engine pins actually used (launch.log + /run_manifest.json).""" @@ -43,12 +64,16 @@ def _write_run_manifest(args) -> dict | None: "miles_commit": _rl.MILES_NEXT_COMMIT, "sglang_commit": _rl.SGLANG_NEXT_COMMIT, "source_sha256": getattr(args, "source_sha256", None), + **_code_identity_manifest(args), "cluster_prefix": args.cluster_prefix, **_miles_overlay_manifest(args), **_island_tasks_manifest(args), # launch-preflight-guards: preflight readings, negative_test, island_overrides **(getattr(args, "_launch_preflight_manifest", None) or {}), "island_placements": rl_island_placement_records(args), + # rl-offpolicy-correction-passthrough D8: correction group per island + **({"offpolicy_correction": args._offpolicy_correction_manifest} + if getattr(args, "_offpolicy_correction_manifest", None) else {}), "written_unix": time.time(), } print(f"[launcher] RL image {manifest['rl_image']}, " @@ -85,3 +110,37 @@ def _no_sync_events_dir(args) -> Path: from .. import runs return runs.run_dir(args.cluster_prefix) / "events" + + +def _record_trainer_start(args, events_dir) -> dict | None: + """rl-trainer-cold-start-cost L0.2: copy each island's ``trainer_start_summary`` + from the settled tapes into ``run_manifest.json`` (``trainer_start``), where the + scheduler reads ``T_s`` of the last run. Never fails the run.""" + + return _record_trainer_start_from(args, sorted(Path(events_dir).glob("*.jsonl"))) + + +def _record_sky_trainer_start(args, clusters, diag_root: str | None = None) -> dict | None: + """``trainer_start`` for sky islands whose tape is not echoed to the launcher: + read the island tapes (``rl-island-*.jsonl``) that the pre-teardown diagnostics + fetched into ``//``. Never fails the run.""" + + from .teardown import TEARDOWN_DIAG_DIR + + root = Path(diag_root or TEARDOWN_DIAG_DIR).expanduser() + tapes = [p for c in clusters for p in sorted((root / c).glob("rl-island-*.jsonl"))] + return _record_trainer_start_from(args, tapes) if tapes else None + + +def _record_trainer_start_from(args, tapes) -> dict | None: + from .. import runs + from ..rl.engine import start_phases + + try: + summaries = start_phases.summaries_from_events(start_phases.read_tape_records(tapes)) + path = runs.run_dir(args.cluster_prefix) / "run_manifest.json" + if start_phases.write_manifest_summary(path, summaries): + return summaries + except (OSError, ValueError) as e: + print(f"[launcher] could not record trainer start phases: {e}", file=sys.stderr) + return None diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index 143fb8f19..a7f12af6f 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -12,6 +12,7 @@ from ..gpu_spec import ClusterSpec from . import REPO_ROOT +from .code_pack import code_workdir from .remote_setup import HF_TOKEN_PATH from .syncer import critic_syncer_address, island_hmac_secret from .credentials import SECRET_ENV_NAMES, _secret_value, registry_login_for @@ -124,7 +125,7 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy memory_gib_override=getattr(args, "modal_memory_gib", None), tape_subdir=(modal_tape_subdir(args.cluster_prefix, learner_id) if getattr(args, "modal_tape_volume", None) else None), - workdir=str(REPO_ROOT), + workdir=code_workdir(args), # Opt-in overrides (default: Modal runner defaults, 10 retries / 24 h). # Acceptance runs pass --modal-retries 0 so a learner exit is final # (no re-run billing) and --modal-timeout-s as the Modal-side hard stop. diff --git a/yeto/launch/remote_setup.py b/yeto/launch/remote_setup.py index 6522dc0ae..5d2f75a19 100644 --- a/yeto/launch/remote_setup.py +++ b/yeto/launch/remote_setup.py @@ -2,6 +2,7 @@ from __future__ import annotations +import shlex from pathlib import Path @@ -285,3 +286,31 @@ def causal_kernel_setup_steps(args) -> list[str]: ] ) return steps + + +MODEL_PREFETCH_LOG = "/tmp/hf-prefetch.log" + + +def model_prefetch(repo: str, revision: str | None) -> str: + """Background Hub prefetch of ``repo`` into the HF cache (setup step). + + ``huggingface_hub.snapshot_download`` through the image's python3: the + ``huggingface-cli`` the old step called is deprecated and exits 1 in the + ports image (coldstart-gpu-review 6.2 #6). A failure does not stop the + run (the island downloads again in the foreground); it is recorded as + ``model_prefetch_failed`` with the exit code (start_phases.shell_event). + """ + + from ..rl.engine.start_phases import shell_event + + fetch_py = ( + "from huggingface_hub import snapshot_download; " + f"snapshot_download(repo_id={repo!r}, revision={revision!r})" + ) + failed = shell_event("model_prefetch_failed", {"model": repo, "log": MODEL_PREFETCH_LOG}, + {"exit_code": "$rc"}) + body = ( + f"python3 -c {shlex.quote(fetch_py)} >{MODEL_PREFETCH_LOG} 2>&1; rc=$?\n" + f'if [ "$rc" != 0 ]; then\n{failed}\nfi' + ) + return f"(nohup sh -c {shlex.quote(body)} >/dev/null 2>&1 &)" diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 67360ad70..8145352f4 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -249,8 +249,12 @@ def _prepare_rl_args( ) if resolve_model_kind(args.model, args.model_kind) != "causal-lm": raise ValueError("RL v0 supports only causal language models") - if args.tuning != "lora": - raise ValueError("RL v0 requires --tuning lora") + from ..rl.engine import sample_roles + + if args.tuning != "lora" and not sample_roles.buffer_full_parameter(args): + raise ValueError("RL v0 requires --tuning lora (--tuning full: only sampler / trainer " + "islands, --rl-sample-role or --rl-island-sample-role, with " + "--rl-single-island-no-sync)") if args.lora_r <= 0: raise ValueError("RL v0 requires a positive LoRA rank") if args.total_steps <= 0: diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 94298c5ec..a902b9ec3 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -6,13 +6,15 @@ import os import shlex import sys +import time from ..gpu_spec import ClusterSpec from ..hw.catalog import COMPAT_GROUP_ENV, compat_group -from . import REPO_ROOT +from .code_pack import code_pack_island_flag, code_workdir from .rl_backend import _miles_overlay, _rl_backend_module from .remote_setup import ( HF_TOKEN_ENV, HF_TOKEN_PATH, PORTS_MEGATRON_PATH, WAN_TUNING, island_pre_run_block, + model_prefetch, ) from .topology import ( multinode_env_prelude, multinode_network_tier, rl_cross_node_flags, rl_island_bundle_map_flag, @@ -32,6 +34,7 @@ from .rl_ports import _ports_algorithm_flags from .rl_args import sao_contract_mounts from .images import _miles_source_setup +from ..rl.engine.start_phases import SUBMITTED_ENV, shell_mark, shell_reset from .in_vm_docker import ( DOCKER_LOGIN_UNSET, IN_VM_DOCKER_CLOUDS, in_vm_docker_run, in_vm_docker_setup, ) @@ -42,6 +45,12 @@ ) +def submitted_unix() -> float: + """Phase ``submitted`` of the island being built (a seam for byte-equal task tests).""" + + return time.time() + + def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: """LoRA rank/targets; for the full-parameter preset the sweep identity (the probed layout hash, the frozen roster size and the syncer profile hash); @@ -57,6 +66,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: f" --sao-context-sha256 {island.context_sha256}") lora = f" --lora-r {args.lora_r} --lora-targets {args.lora_targets}" if not rl_dense_sweep(args): + if getattr(args, "tuning", "lora") == "full": # L4.6 buffer island (checked at launch) + return " --parameter-mode full" return lora sweep = (f" --parameter-layout-sha256 {args.rl_parameter_layout_sha256}" f" --num-learners {num_learners}" @@ -109,7 +120,8 @@ def make_miles_island_task( f" --reward-function {shlex.quote(args.reward_function)}" f" --reward-sha256 {shlex.quote(args.reward_sha256)}" f" --source-sha256 {shlex.quote(args.source_sha256)}" - f" --global-rounds {args.total_steps}" + + code_pack_island_flag(args) + + f" --global-rounds {args.total_steps}" f" --sync-preset {args.rl_sync_preset}" f" --fragments {args.fragments}" f" --pipeline {args.pipeline}" @@ -207,6 +219,8 @@ def make_miles_island_task( flags += f" --apply-chat-template-kwargs {shlex.quote(chat_kwargs)}" if not getattr(args, "sglang_deterministic_inference", True): flags += " --no-sglang-deterministic-inference" + if getattr(args, "rl_sampler_snapshot", False): + flags += " --rl-sampler-snapshot" if args.custom_generate_function_path: flags += ( " --custom-generate-function-path " @@ -269,15 +283,7 @@ def make_miles_island_task( if is_local_reference(model): prefetch = ": # local model; no Hub prefetch" else: - revision = ( - f" --revision {shlex.quote(args.model_revision)}" - if args.model_revision - else "" - ) - prefetch = ( - f"(nohup huggingface-cli download {shlex.quote(model)}{revision} " - ">/tmp/hf-prefetch.log 2>&1 &) || true" - ) + prefetch = model_prefetch(model, args.model_revision) worker_model_fetch = "" if spec.num_nodes > 1 and not is_local_reference(model): # Same call and cache as the learner's own snapshot_download on the @@ -312,6 +318,8 @@ def make_miles_island_task( **({"CRITIC_SYNCER_ADDR": critic_syncer_address(syncer_addr)} if rl_needs_critic(args) and syncer_addr != "none" else {}), "LEARNER_ID": str(learner_id), + # rl-trainer-cold-start-cost L0: phase ``submitted`` (launcher clock) + SUBMITTED_ENV: f"{submitted_unix():.3f}", # decoupling 7.7a: "-" from the card catalog; an unknown # card raises here, before any cloud spend (no default). COMPAT_GROUP_ENV: compat_group(spec.gpu), @@ -369,7 +377,10 @@ def make_miles_island_task( envs["YETO_RUN_GROUP"] = args.cluster_prefix if os.environ.get("WANDB_API_KEY"): envs["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"] - setup_steps = [WAN_TUNING, HF_TOKEN_ENV, miles_setup, sglang_setup] + # rl-trainer-cold-start-cost L0: the setup shell records its start/end and + # the run shell records ray_ready (start_phases.MARK_DIR); the island replays them. + setup_steps = [shell_reset(), shell_mark("container_ready"), + WAN_TUNING, HF_TOKEN_ENV, miles_setup, sglang_setup] if getattr(args, "wandb", False): setup_steps.append( "pip install -q wandb || echo '[yeto-setup] wandb install failed; " @@ -419,6 +430,7 @@ def make_miles_island_task( # dataclasses.asdict on it: "asdict() should be called on dataclass # instances"; B1 nsmoke). SkyPilot exports secrets into setup/run, so both # scripts unset them first. Default (public image): nothing injected. + setup_steps.append(shell_mark("setup_done")) registry_login = registry_login_for(args) login_unset = DOCKER_LOGIN_UNSET if registry_login else "" setup_script = login_unset + "\n".join(setup_steps) @@ -462,6 +474,7 @@ def make_miles_island_task( '--port=6379 --include-dashboard=true --temp-dir="$MILES_RAY_DIR"' + rl_island_ray_gpus(args, spec) + (" --dashboard-host=0.0.0.0" if spec.num_nodes > 1 else "") + "\n" + f" {shell_mark('ray_ready')}\n" " trap stop_miles_ray EXIT\n" # Miles calls ray.init(address="auto"), which reads RAY_ADDRESS # first and otherwise sky's /tmp/ray/ray_current_cluster file. @@ -515,7 +528,7 @@ def make_miles_island_task( run=run_script, envs=envs, num_nodes=spec.num_nodes, - workdir=str(REPO_ROOT), + workdir=code_workdir(args), file_mounts=file_mounts or None, ) resources = { diff --git a/yeto/launch/rl_ports.py b/yeto/launch/rl_ports.py index 811dea937..ede091d59 100644 --- a/yeto/launch/rl_ports.py +++ b/yeto/launch/rl_ports.py @@ -77,11 +77,17 @@ def _prepare_ports_algorithm(args) -> None: len(parse_gpu_spec(args.gpu)) if getattr(args, "gpu", None) else 1 ) + max(0, getattr(args, "external_learners", 0) or 0) no_sync = bool(getattr(args, "rl_single_island_no_sync", False)) - if no_sync and (islands != 1 or getattr(args, "rl_sync_preset", "strict-avg") + from ..rl.engine import sample_roles + + # L5.4: a store-coupled sampler/trainer pair is two no-sync islands + # (sample_roles.check_pair holds its own rules). + no_sync_islands = 2 if sample_roles.is_pair(args) else 1 + if no_sync and (islands != no_sync_islands or getattr(args, "rl_sync_preset", "strict-avg") != "strict-avg" or getattr(args, "rl_initial_adapter", None)): raise AlgorithmSpecError( "--rl-single-island-no-sync needs exactly one island (one --gpu entry, no " - "external learners), the default sync preset and no initial adapter" + "external learners; two with --rl-island-sample-role), the default sync " + "preset and no initial adapter" ) args.rl_allow_unverified_mechanism = list( check_unverified_allowance( @@ -115,6 +121,25 @@ def _prepare_ports_algorithm(args) -> None: _check_critic_warmup_options(args, spec) args.rl_algorithm_spec_json = spec.canonical_json() args.rl_expected_algorithm_sha256 = spec.sha256() + args._offpolicy_correction_manifest = _offpolicy_correction_manifest(args, spec) + + +def _offpolicy_correction_manifest(args, spec) -> dict: + """rl-offpolicy-correction-passthrough D8/L6.5: every island's correction + record and the warnings (unmappable fields, island differences). Phase 1 + prints the warnings and goes on; the run manifest keeps both.""" + + from ..rl.engine.offpolicy_correction import run_correction_manifest + from .topology import rl_island_args + + backends = {str(m): getattr(rl_island_args(args, m), "rl_backend", None) or "miles" + for m in range(rl_island_count(args))} + section = run_correction_manifest(spec, backends) + for warning in section["warnings"]: + print("[launcher] correction_config_warning: " + + json.dumps({k: v for k, v in warning.items() if k != "event"}, sort_keys=True), + flush=True) + return section def _check_critic_warmup_options(args, spec) -> None: @@ -347,6 +372,12 @@ def _check_ports_infra_switches(args) -> None: raise ValueError("--rl-megatron-ref-load needs an absolute path without whitespace") if getattr(args, "rl_boot_only", False) and not getattr(args, "rl_single_island_no_sync", False): raise ValueError("--rl-boot-only needs --rl-single-island-no-sync") + from ..rl.engine import sample_roles + + if sample_roles.is_pair(args): # L5.4 --rl-island-sample-role: per-island roles + sample_roles.check_pair(args, rl_island_count(args)) + else: + sample_roles.check_launch(args) # L4.6 --rl-sample-role (both checks nothing) pre_run = getattr(args, "rl_island_pre_run", None) if pre_run is not None: if not getattr(args, "rl_single_island_no_sync", False): @@ -494,6 +525,9 @@ def _ports_infra_flags(args) -> tuple[str, str]: flags += " --rl-boot-only" if getattr(args, "rl_max_policy_age", 0): # 0 (default) adds nothing flags += f" --rl-max-policy-age {int(args.rl_max_policy_age)}" + from ..rl.engine.sample_roles import island_flags as sample_role_flags + + flags += sample_role_flags(args) # L4.6: role both (default) adds nothing if (getattr(args, "rl_lr_schedule", None) or "auto") != "auto": # auto adds nothing flags += f" --rl-lr-schedule {args.rl_lr_schedule}" if island_scheduling_mode(args) == "elastic": # legacy adds nothing diff --git a/yeto/launch/trainer_island.py b/yeto/launch/trainer_island.py new file mode 100644 index 000000000..8d35111e8 --- /dev/null +++ b/yeto/launch/trainer_island.py @@ -0,0 +1,311 @@ +"""Cloud binding of the trainer scheduler: the trainer island on Modal +(rl-trainer-start-stop L5.3c). + +With ``--rl-trainer-scheduler s1`` the launcher holds back the island named by +``--rl-trainer-island``: its Modal function is defined on the run's app but the +fleet controller does not start it. :func:`modal_binding` gives the scheduler loop: + +* :class:`ModalTrainerIslands` (``trainer_remote.TrainerIslands``): each + incarnation spawns the held-back function with the incarnation environment + (``trainer_remote.island_env``), tails its log stream for the island events, and + cancels the call at release. The call ids are kept in the head's run directory, + so a restarted head can release an orphan; +* :class:`ModalVolumeFiles`: the head's reads and writes in the run's + ``modal-volume://`` checkpoint store (commands, cut pointers, the buffer index); +* the head's buffer view: the sample buffer index mirrored from the Volume, + backlog counted from the consumed watermark the trainer committed: the running + incarnation's reported one, else the newest cut's (``CommittedBacklog``). + +Sky trainer islands are not bound yet (``check`` refuses them). +""" + +from __future__ import annotations + +import io +import json +import threading +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from yeto.rl import event_echo +from yeto.rl.engine.resume import CheckpointStore +from yeto.rl.engine.sample_buffer import BUFFER_DIR, INDEX, CommittedBacklog, FileSampleBuffer +from yeto.rl.engine.start_phases import SUBMITTED_ENV +from yeto.rl.engine.trainer_lifecycle import CutLog, CutRef, DurableSampleBuffer, TrainerLifecycle +from yeto.rl.engine.trainer_remote import ( + EVENT, + RemoteTrainerLifecycle, + StoreCutLog, + StoreFiles, + consumed_watermark, + island_env, +) + +JOBS_FILE = "trainer-islands.json" # in the trainer loop's journal dir: incarnation -> call id +MIRROR_DIR = "buffer-mirror" +RELEASE_CHECKS = 12 +RELEASE_CHECK_DELAY_S = 5.0 + + +@dataclass +class TrainerIsland: + """The held-back trainer island the launcher hands to the binding.""" + + name: str + cloud: str + modal_cfg: Any = None # modal_runner.ModalIslandConfig + modal_ops: Any = None # modal_runner.ModalOps (app deployed) + + +@dataclass +class CloudTrainerBinding: + """What the loop needs from the cloud layer for one run. + + ``lifecycle(first_incarnation)`` builds the trainer lifecycle (e.g. + ``DriverLifecycle`` whose ``build_driver`` acquires a trainer island and whose + ``release`` gives it back). ``release_orphan(incarnation)`` gives back the island + a dead head started. ``data_ended`` tells the samplers finished. + ``after_restore()`` runs after the journal is read back and an orphan released + (e.g. reset the head's buffer to the watermark of the newest cut).""" + + lifecycle: Callable[[int], TrainerLifecycle] + buffer: DurableSampleBuffer + cuts: CutLog + release_orphan: Callable[[int], None] | None = None + data_ended: Callable[[], bool] = lambda: False + start_events: Callable[[], Iterable[Mapping[str, Any]]] | None = None + spec: Mapping[str, Any] | None = None + after_restore: Callable[[], None] | None = None + + +# --------------------------------------------------------------------------- +# Store access from the head +# --------------------------------------------------------------------------- +class ModalVolumeFiles: + """:class:`~yeto.rl.engine.trainer_remote.StoreFiles` over a Modal Volume via + the SDK (``read_file`` / ``batch_upload``). ``prefix``: the store's path in the + Volume (``modal-volume://NAME/PREFIX``).""" + + def __init__(self, volume_name: str, prefix: str = "", *, + volume_factory: Callable[[str], Any] | None = None) -> None: + self.volume_name = volume_name + self.prefix = prefix.strip("/") + self._factory = volume_factory + self._volume: Any = None + + def _vol(self) -> Any: + if self._volume is None: + if self._factory is not None: + self._volume = self._factory(self.volume_name) + else: + import modal + + self._volume = modal.Volume.from_name(self.volume_name) + return self._volume + + def _path(self, rel: str) -> str: + return f"{self.prefix}/{rel}" if self.prefix else rel + + def read(self, rel: str) -> bytes | None: + try: + return b"".join(self._vol().read_file(self._path(rel))) + except Exception as error: # noqa: BLE001 - the SDK names a missing file differently by version + if isinstance(error, FileNotFoundError) or type(error).__name__ == "NotFoundError": + return None + raise + + def write(self, rel: str, data: bytes) -> None: + with self._vol().batch_upload(force=True) as batch: + batch.put_file(io.BytesIO(data), self._path(rel)) + + +class MirrorStore(CheckpointStore): + """A read-only local copy of some store files: ``reload`` copies ``paths`` from + ``files`` (whole files). Lets :class:`FileSampleBuffer` read an index that + lives in a store this host does not mount.""" + + kind = "mirror" + + def __init__(self, root: str | Path, files: StoreFiles, paths: Iterable[str]) -> None: + super().__init__(root) + self.files = files + self.paths = tuple(paths) + + def reload(self) -> None: + for rel in self.paths: + data = self.files.read(rel) + if data is None: + continue + path = Path(self.root) / rel + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + tmp.write_bytes(data) + tmp.replace(path) + + def commit(self) -> float: + raise RuntimeError("the head's mirror of the store is read-only") + + +# --------------------------------------------------------------------------- +# Islands +# --------------------------------------------------------------------------- +class IslandEvents: + """Feeds on an island's log stream (``collector.feed`` like ``TapeCollector``); + keeps the ``rl_trainer_lifecycle`` and start-phase events.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._new: list[dict[str, Any]] = [] + self.all: list[dict[str, Any]] = [] + + def feed(self, chunk: Any) -> None: + for line in str(chunk).splitlines(): + raw = event_echo.parse_line(line) + if raw is None or raw is event_echo.INVALID: + continue + record = json.loads(raw) + with self._lock: + self.all.append(record) + if record.get("event") == EVENT: + self._new.append(record) + + def take(self) -> list[dict[str, Any]]: + with self._lock: + out, self._new = self._new, [] + return out + + +@dataclass +class _Job: + incarnation: int + call_id: str + events: IslandEvents + + +class ModalTrainerIslands: + """:class:`~yeto.rl.engine.trainer_remote.TrainerIslands` on Modal (module doc). + + ``tail(call_id, label, collector)`` streams the call's logs into ``collector`` + (default: the launcher's ``_tail_modal`` on a daemon thread).""" + + def __init__(self, modal_ops: Any, cfg: Any, *, name: str, jobs_path: Path, + poll_s: float | None = None, + tail: Callable[[str, str, IslandEvents], None] | None = None, + sleep: Callable[[float], None] = time.sleep, + clock: Callable[[], float] = time.time) -> None: + self.modal_ops = modal_ops + self.cfg = cfg + self.name = name + self.jobs_path = Path(jobs_path) + self.poll_s = poll_s + self._tail = tail or self._tail_thread + self.sleep = sleep + self.clock = clock + self.jobs: list[_Job] = [] + + # -- TrainerIslands -------------------------------------------------------- + def launch(self, incarnation: int, resume_from: CutRef | None) -> _Job: + # Phase ``submitted`` of this incarnation is the spawn, not the run's launch: + # the held-back task was built when the fleet started (T_s, start_phases). + envs = {**self.cfg.envs, **island_env(incarnation, resume_from, poll_s=self.poll_s), + SUBMITTED_ENV: f"{self.clock():.3f}"} + call_id = self.modal_ops.spawn(replace(self.cfg, envs=envs)) + self._save_call(incarnation, call_id) + job = _Job(int(incarnation), str(call_id), IslandEvents()) + self.jobs.append(job) + print(f"[launcher] trainer incarnation {incarnation} on Modal as {self.name} " + f"(call {call_id}, resume {None if resume_from is None else resume_from.cut_id})", + flush=True) + self._tail(job.call_id, f"{self.name}-i{incarnation}", job.events) + return job + + def events(self, job: _Job) -> list[Mapping[str, Any]]: + return job.events.take() + + def alive(self, job: _Job) -> bool: + return self.modal_ops.status(job.call_id) == "RUNNING" + + def release(self, job: _Job) -> bool: + return self._release_call(job.call_id) + + def release_incarnation(self, incarnation: int) -> bool: + call_id = self._calls().get(str(int(incarnation))) + if call_id is None: + print(f"[launcher] trainer incarnation {incarnation}: no Modal call recorded; " + "nothing to release", flush=True) + return False + return self._release_call(call_id) + + def start_events(self) -> list[Mapping[str, Any]]: + """Every event of every incarnation (``trainer_start_summary`` gives T_s).""" + return [e for job in self.jobs for e in job.events.all] + + # -- internals ------------------------------------------------------------- + def _release_call(self, call_id: str) -> bool: + """Cancel the call (containers terminated) and check it is no longer + running: the cards are back when the call ended.""" + try: + self.modal_ops.cancel(call_id) + except Exception as error: # noqa: BLE001 - an ended call may refuse; the status decides + print(f"[launcher] cancel of trainer call {call_id} failed: {error}", flush=True) + for attempt in range(RELEASE_CHECKS): + if self.modal_ops.status(call_id) != "RUNNING": + return True + if attempt + 1 < RELEASE_CHECKS: + self.sleep(RELEASE_CHECK_DELAY_S) + return False + + def _calls(self) -> dict[str, str]: + if not self.jobs_path.is_file(): + return {} + return dict(json.loads(self.jobs_path.read_text(encoding="utf-8"))) + + def _save_call(self, incarnation: int, call_id: str) -> None: + calls = {**self._calls(), str(int(incarnation)): str(call_id)} + self.jobs_path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.jobs_path.with_name(self.jobs_path.name + ".tmp") + tmp.write_text(json.dumps(calls, sort_keys=True), encoding="utf-8") + tmp.replace(self.jobs_path) + + def _tail_thread(self, call_id: str, label: str, collector: IslandEvents) -> None: + from .modal_islands import _tail_modal + + threading.Thread(target=_tail_modal, args=(self.modal_ops, call_id, label, collector), + name=f"tail-{label}", daemon=True).start() + + +# --------------------------------------------------------------------------- +# Binding +# --------------------------------------------------------------------------- +def modal_binding(args: Any, island: TrainerIsland, *, journal_dir: Path, + files: StoreFiles | None = None, islands: Any = None) -> Any: + """The :class:`CloudTrainerBinding` of a Modal trainer + island (``files`` / ``islands`` replace the SDK in tests).""" + from .storage import MODAL_CHECKPOINT_STORE_MOUNT, modal_checkpoint_store + + if files is None: + volume, mount = modal_checkpoint_store(args) + files = ModalVolumeFiles(volume, mount[len(MODAL_CHECKPOINT_STORE_MOUNT):]) + if islands is None: + islands = ModalTrainerIslands(island.modal_ops, island.modal_cfg, name=island.name, + jobs_path=Path(journal_dir) / JOBS_FILE, + poll_s=getattr(args, "rl_trainer_command_poll_seconds", None)) + cuts = StoreCutLog(files) + mirror = MirrorStore(Path(journal_dir) / MIRROR_DIR, files, [f"{BUFFER_DIR}/{INDEX}"]) + buffer = FileSampleBuffer(mirror, batch_size=int(args.rollout_batch_size)) + poll_s = float(getattr(args, "rl_trainer_command_poll_seconds", None) or 5.0) + lifecycles: list[RemoteTrainerLifecycle] = [] + + def lifecycle(first: int) -> RemoteTrainerLifecycle: + lifecycles.append(RemoteTrainerLifecycle(islands, files, first_incarnation=first, poll_s=poll_s)) + return lifecycles[-1] + + def watermark() -> int: + return consumed_watermark(cuts, lifecycles[-1] if lifecycles else None) + + return CloudTrainerBinding( + lifecycle=lifecycle, buffer=CommittedBacklog(buffer, watermark), cuts=cuts, + release_orphan=islands.release_incarnation, start_events=islands.start_events) diff --git a/yeto/launch/trainer_loop.py b/yeto/launch/trainer_loop.py new file mode 100644 index 000000000..dfaffef15 --- /dev/null +++ b/yeto/launch/trainer_loop.py @@ -0,0 +1,255 @@ +"""Trainer scheduler loop on the head (rl-trainer-start-stop L5.3b). + +``--rl-trainer-scheduler s1`` runs :class:`~yeto.rl.engine.trainer_loop.TrainerLoop` +on its own thread next to the fleet controller's poll loop. Off (the default), the +launcher behaves as before: :func:`run_with_trainer_loop` only calls ``run``. + +Inputs and outputs of the loop on the head: + +* journal: ``/trainer-loop/journal.jsonl`` (``yeto.rl.engine.journal``). + A head restart with the same ``--cluster-prefix`` reopens it and continues; +* ``T_s``: the configured ``--rl-trainer-start-seconds`` until a start-phase + summary event is available (``start_events`` of the cloud binding); +* ``A_hard``: ``--rl-max-policy-age``; +* the trainer island itself: a :class:`CloudTrainerBinding` from + :data:`binding_factory`. It acquires a trainer island per incarnation and + releases it, and gives the sample buffer and the cut log of the run. + +The launcher holds back the island ``--rl-trainer-island`` (index into ``--gpu``): +the fleet controller does not start it; the binding starts one per incarnation. +Only a Modal trainer island is bound (``yeto.launch.trainer_island``); :func:`check` +refuses anything else before any launch. +""" + +from __future__ import annotations + +import dataclasses +import sys +import threading +from collections.abc import Callable +from pathlib import Path +from typing import Any, TypeVar + +from yeto.rl.engine.convertible import NodeSpec, parse_node_spec, validate_nodes +from yeto.rl.engine.journal import Journal +from yeto.rl.engine.trainer_loop import ( + DEFAULT_MAX_CONSECUTIVE_FAILURES, + TrainerLoop, + TrainerLoopGaveUp, + TrainerLoopThread, + next_incarnation, +) +from yeto.rl.engine.trainer_scheduler import S1Config, S1Scheduler + +from .trainer_island import CloudTrainerBinding, TrainerIsland + +OFF = "off" +S1 = "s1" +MODES = (OFF, S1) +DEFAULT_INTERVAL_S = 30.0 +DEFAULT_START_SECONDS = 600.0 +DEFAULT_DRAIN_SECONDS = 3600.0 +TRAINER_PRECISION = "bf16" # the trainer's precision; a convertible node must match it +JOURNAL_DIR = "trainer-loop" + +T = TypeVar("T") + + +def default_binding(args, island: TrainerIsland | None) -> CloudTrainerBinding: + if island is None or island.cloud != "modal": + raise SystemExit("--rl-trainer-scheduler: the trainer island must be a Modal island " + f"(got {None if island is None else island.cloud}); sky is not bound yet") + from .trainer_island import modal_binding + + return modal_binding(args, island, journal_dir=journal_dir(args)) + + +# (args, held-back trainer island) -> CloudTrainerBinding +binding_factory: Callable[[Any, TrainerIsland | None], CloudTrainerBinding] | None = default_binding + + +def add_arguments(group) -> None: + group.add_argument( + "--rl-trainer-scheduler", choices=MODES, default=OFF, + help="start and stop the trainer island by scheduler S1 from the sample-buffer " + "backlog (rl-trainer-start-stop); off (default) keeps the trainer up for the run") + group.add_argument( + "--rl-trainer-scheduler-interval", type=float, default=DEFAULT_INTERVAL_S, + help="seconds between two scheduler decisions on the head") + group.add_argument( + "--rl-trainer-island", type=int, default=None, metavar="M", + help="with --rl-trainer-scheduler: the island (index into --gpu) the scheduler starts " + "and stops; the fleet does not start it. Must be a Modal island with a " + "modal-volume:// --rl-checkpoint-store") + group.add_argument( + "--rl-trainer-command-poll-seconds", type=float, default=None, + help="seconds between two reads of the command file on the trainer island and of " + "the island events on the head (default 5)") + group.add_argument( + "--rl-trainer-drain-seconds", type=float, default=DEFAULT_DRAIN_SECONDS, + help="after the fleet (the samplers) ended: max seconds the scheduler keeps the " + "trainer to consume the rest of the buffer before it stops it (default 3600)") + group.add_argument( + "--rl-trainer-max-failures", type=int, default=DEFAULT_MAX_CONSECUTIVE_FAILURES, + help="trainer incarnations that may fail in a row (no trained round between them) " + "before the scheduler stops, releases the trainer and fails the run (default " + f"{DEFAULT_MAX_CONSECUTIVE_FAILURES}); a trained round resets the count") + group.add_argument( + "--rl-trainer-start-seconds", type=float, default=DEFAULT_START_SECONDS, + help="configured trainer start time T_s (seconds); a measured start replaces it") + group.add_argument( + "--rl-trainer-start-seconds-fixed", action="store_true", + help="keep --rl-trainer-start-seconds as T_s after a start was measured (an override " + "for check runs; the measured start is still journaled)") + group.add_argument( + "--rl-node-class", action="append", default=[], metavar="NAME=CLASS[:spot][:PRECISION]", + help="node class of one node (rollout, trainer, convertible; rl-trainer-start-stop " + "D7). A node without ':spot' follows --spot/--on-demand. A convertible node must be " + f"on-demand and at the trainer precision ({TRAINER_PRECISION}). Repeat per node") + + +def node_specs(args) -> list[NodeSpec]: + """``--rl-node-class`` entries; a node is spot if it says so or the run is spot.""" + run_spot = bool(getattr(args, "spot", False)) + specs = [] + for text in getattr(args, "rl_node_class", None) or (): + spec = parse_node_spec(text) + if run_spot and not spec.spot: + spec = NodeSpec(spec.name, spec.node_class, spot=True, + precision=spec.precision, capacity=spec.capacity) + specs.append(spec) + return specs + + +def check_node_classes(args) -> list[NodeSpec]: + """Refuse a launch whose node classes break the class rule (L7.1).""" + try: + specs = node_specs(args) + validate_nodes(specs, trainer_precision=TRAINER_PRECISION) + except ValueError as exc: # NodeClassError, or a bad capacity number + raise SystemExit(f"--rl-node-class: {exc}") from None + return specs + + +def enabled(args) -> bool: + return (getattr(args, "rl_trainer_scheduler", OFF) or OFF) != OFF + + +def journal_dir(args) -> Path: + from yeto import runs + + return runs.run_dir(args.cluster_prefix) / JOURNAL_DIR + + +def held_island(args) -> int | None: + """Index (into ``--gpu``) of the island the scheduler owns; None with the switch off.""" + return args.rl_trainer_island if enabled(args) else None + + +def check(args) -> None: + """Refuse bad node classes, and the switch without a trainer island binding + (before any launch).""" + check_node_classes(args) + if not enabled(args): + return + if binding_factory is None: + raise SystemExit(f"--rl-trainer-scheduler {args.rl_trainer_scheduler}: no trainer " + "island binding to the cloud layer") + if args.rl_trainer_scheduler_interval <= 0: + raise SystemExit("--rl-trainer-scheduler-interval must be positive") + if getattr(args, "rl_trainer_drain_seconds", DEFAULT_DRAIN_SECONDS) < 0: + raise SystemExit("--rl-trainer-drain-seconds must be >= 0") + if getattr(args, "rl_trainer_max_failures", DEFAULT_MAX_CONSECUTIVE_FAILURES) < 1: + raise SystemExit("--rl-trainer-max-failures must be >= 1") + if binding_factory is default_binding: + _check_modal_island(args) + + +def _check_modal_island(args) -> None: + from yeto.gpu_spec import parse_gpu_spec + + m = getattr(args, "rl_trainer_island", None) + if m is None: + raise SystemExit("--rl-trainer-scheduler needs --rl-trainer-island M (the island it " + "starts and stops)") + specs = parse_gpu_spec(args.gpu) + if not 0 <= m < len(specs): + raise SystemExit(f"--rl-trainer-island {m}: --gpu has {len(specs)} island(s)") + if specs[m].cloud != "modal": + raise SystemExit(f"--rl-trainer-island {m} is a {specs[m].cloud} island; only a Modal " + "trainer island is bound to the scheduler (sky: not implemented)") + if not str(getattr(args, "rl_checkpoint_store", None) or "").startswith("modal-volume://"): + raise SystemExit("--rl-trainer-scheduler needs a modal-volume:// --rl-checkpoint-store: " + "the head reads the cuts and the sample buffer there and writes the " + "trainer commands there") + + +def build(args, binding: CloudTrainerBinding, journal: Journal) -> TrainerLoop: + """Scheduler S1 + lifecycle + loop over ``journal``; restored from its records.""" + records = journal.records + scheduler = S1Scheduler( + S1Config(a_hard=int(getattr(args, "rl_max_policy_age", 0) or 0), + start_seconds=float(args.rl_trainer_start_seconds), + fixed_start_seconds=bool(getattr(args, "rl_trainer_start_seconds_fixed", False))), + start_events=binding.start_events, record=journal.append) + loop = TrainerLoop(scheduler=scheduler, lifecycle=binding.lifecycle(next_incarnation(records)), + buffer=binding.buffer, cuts=binding.cuts, record=journal.append, + data_ended=binding.data_ended, release_orphan=binding.release_orphan, + spec=binding.spec, + max_consecutive_failures=int(getattr( + args, "rl_trainer_max_failures", DEFAULT_MAX_CONSECUTIVE_FAILURES))) + if loop.restore(records): + print(f"[launcher] trainer scheduler continues from {journal.dir}", flush=True) + if binding.after_restore is not None: + binding.after_restore() + return loop + + +def run_with_trainer_loop(args, run: Callable[[], T], *, island: TrainerIsland | None = None, + factory: Callable[[Any, TrainerIsland | None], CloudTrainerBinding] + | None = None, + abort: Callable[[BaseException], None] | None = None) -> T: + """``run()`` (the fleet controller); with the switch on, the trainer loop ticks + on its own thread meanwhile. ``island``: the held-back trainer island + (:func:`held_island`). + + When ``run`` returns, the fleet (the samplers) has ended: the data ended. The + loop keeps ticking up to ``--rl-trainer-drain-seconds`` so the trainer consumes + the rest of the buffer and the loop finishes (backlog under one batch); then + the trainer is stopped. When ``run`` raises, the trainer is stopped at once. + + ``abort(error)``: called once from the loop thread when the loop gives up + (:class:`TrainerLoopGaveUp`); the fleet controller then ends ``run`` with that + error so the samplers stop too. Without it the error is raised after ``run``.""" + if not enabled(args): + return run() + factory = binding_factory if factory is None else factory + if factory is None: + check(args) # raises + binding = factory(args, island) + fleet_ended = threading.Event() + cloud_data_ended = binding.data_ended + binding = dataclasses.replace( + binding, data_ended=lambda: fleet_ended.is_set() or bool(cloud_data_ended())) + with Journal(journal_dir(args)) as journal: + def on_error(error: BaseException) -> None: + if abort is not None and isinstance(error, TrainerLoopGaveUp): + abort(error) + + runner = TrainerLoopThread(build(args, binding, journal), + args.rl_trainer_scheduler_interval, on_error=on_error).start() + try: + result = run() + if isinstance(runner.error, TrainerLoopGaveUp): + raise runner.error + fleet_ended.set() + drain_s = float(getattr(args, "rl_trainer_drain_seconds", DEFAULT_DRAIN_SECONDS)) + if not runner.join(drain_s): + print(f"[launcher] trainer scheduler: buffer not drained after {drain_s:.0f} s " + "since the fleet ended; stopping the trainer", file=sys.stderr, flush=True) + return result + finally: + runner.stop() + if runner.error is not None: + print(f"[launcher] trainer scheduler stopped on error: {runner.error!r}", + file=sys.stderr, flush=True) diff --git a/yeto/launch_preflight.py b/yeto/launch_preflight.py index 7cc025fca..419432c2e 100644 --- a/yeto/launch_preflight.py +++ b/yeto/launch_preflight.py @@ -362,6 +362,13 @@ def add_cli_args(p) -> None: "the latest weights and the cut) or droppable (spot, comes back from the syncer " "base, never from a checkpoint store); repeatable; an island without a role is " "an anchor; replaces --spot") + g.add_argument("--rl-island-sample-role", action="append", default=None, metavar="ISLAND:ROLE", + help="rl-trainer-start-stop L5.4: one launch of a sampler island and a trainer " + "island over one --rl-checkpoint-store (ROLE sampler or trainer, one entry per " + "--gpu island; exactly two islands). Needs --rl-single-island-no-sync: no " + "syncer, each island runs as a one-learner no-sync island with that " + "--rl-sample-role. Combine with --rl-trainer-scheduler s1 " + "--rl-trainer-island M to start and stop the trainer island") g.add_argument("--rl-negative-test-run", action="store_true", help="mark this run as a negative-test run (required by --rl-island-override); " "a normal run refuses to resume its checkpoint store and `yeto merge` refuses " diff --git a/yeto/launcher.py b/yeto/launcher.py index fb1120f57..17f021faf 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -39,7 +39,7 @@ from .models import MODEL_WEIGHT_GB from .rl.event_echo import TapeCollector as EventCollector from .launch.rl_backend import _rl_backend_module -from .launch.topology import rl_island_args, rl_island_placement_records, rl_island_spec +from .launch.topology import rl_island_args, rl_island_count, rl_island_placement_records, rl_island_spec from .launch.syncer import ( SYNCER_PORT, island_scheduling_mode, island_syncer_addrs, make_syncer_task, probe_syncer_ports, resolve_island_syncer_epoch, rl_role_syncers, role_syncer_commands, syncer_event_tape, @@ -66,15 +66,20 @@ ) from .launch.fleet import ( CONTAINER_CHANGED_EXIT, DEFAULT_RL_STALL_TIMEOUT_S, FAILED_RUN_DRAIN_S, ISLAND_FAILED_EXIT, - RUN_STALLED_EXIT, TEARDOWN_UNVERIFIED_EXIT, FixedRosterIslandAbandoned, FleetController, + RUN_STALLED_EXIT, TEARDOWN_UNVERIFIED_EXIT, TRAINER_GAVE_UP_EXIT, FixedRosterIslandAbandoned, + FleetController, RunStalled, controller_fixed_roster, ) from .launch.teardown import collect_diagnostics_parallel, teardown_island from .launch.verda import _rename_hook, launch_verda_island, prepare_verda_islands -from .launch.manifest import _dashboard_fleet_log, _no_sync_events_dir, _write_run_manifest +from .launch.manifest import ( + _dashboard_fleet_log, _no_sync_events_dir, _record_sky_trainer_start, _record_trainer_start, + _write_run_manifest, +) # Re-exported: callers and tests reach these as ``yeto.launcher.``. from .launch import REPO_ROOT # noqa: F401 +from .launch.code_pack import prepare_run_code_pack from .launch.rl_backend import _rl_backend_image_ok, resolve_default_rl_image # noqa: F401 from .launch.remote_setup import ( # noqa: F401 HF_TOKEN_ENV, HF_TOKEN_PATH, ISLAND_PRE_RUN_DONE, ISLAND_PRE_RUN_FAILED, ISLAND_PRE_RUN_START, @@ -174,9 +179,13 @@ def _island_task(task_factory, args, spec, m, num_learners, syncer_addr, overrid island's billing and drops the checkpoint store of a droppable island.""" from . import island_overrides as io from .cloud import droppable + from .rl.engine import sample_roles iargs = droppable.role_args(io.island_args(rl_island_args(args, m), m, overrides), role) - task = task_factory(iargs, spec, m, num_learners, syncer_addr) + # L5.4 pair: the island's sample role; it runs as learner 0 of 1 (no syncer). + iargs = sample_roles.island_args(iargs, m, rl_island_count(args)) + learner_id, island_learners = sample_roles.island_learner(args, m, num_learners) + task = task_factory(iargs, spec, learner_id, island_learners, syncer_addr) io.apply_to_task(task, {**io.island_env(args, m, overrides), **(droppable.role_env(role, spec.cloud, spec.region) if role else {})}) return task @@ -224,13 +233,17 @@ def run(args, on_clusters=None, local_syncer=None, on_instance_ids=None) -> int: # Checked first: the flag must never turn an SFT launch syncer-less. raise ValueError("--rl-single-island-no-sync requires --training-mode rl") from . import launch_preflight, sky_patches + from .launch import trainer_loop # launch-preflight-guards 1.3: thread count before anything touches a cloud. args._launch_preflight_manifest = {"preflight": {"threads": launch_preflight.thread_preflight(args)}} sky_patches.install() prepare_launch_args(args) + # rl-identity-tiering L9.4: the run code pack (first start makes it; later starts reuse it). + prepare_run_code_pack(args) # 2.3 / 3.1 / 3.6: memory estimate, island overrides, negative-test resume guard. launch_preflight.pre_cloud_checks(args, threads=False) + trainer_loop.check(args) # rl-trainer-start-stop L5.3b: refuse before any launch _write_run_manifest(args) head_mode = local_syncer is not None # --rl-single-island-no-sync: one ports island, no syncer at all @@ -242,11 +255,16 @@ def run(args, on_clusters=None, local_syncer=None, on_instance_ids=None) -> int: nodes_by_cluster = dict(zip(learner_cluster_names(args.cluster_prefix, specs), (s.num_nodes for s in specs))) # rl-multinode-island D10 no_sync_incomplete: list[str] = [] + # rl-trainer-start-stop L5.3c: the trainer scheduler starts and stops this island + # (one Modal call per incarnation); the fleet controller never sees it. + held = trainer_loop.held_island(args) + held_name = None if held is None else learner_cluster_names(args.cluster_prefix, specs)[held] + held_cfg = None # Islands whose tape travels over the log stream (no fetchable # ~/yeto-output): every Modal RL ports island, and a no-sync island. echo_names = { name for name, spec in zip(learner_cluster_names(args.cluster_prefix, specs), specs) - if _echoes_events(args, spec) + if _echoes_events(args, spec) and name != held_name } events_dir = _no_sync_events_dir(args) if echo_names else None if echo_names: @@ -401,6 +419,11 @@ def run(args, on_clusters=None, local_syncer=None, on_instance_ids=None) -> int: cfg = build_modal_island_config(role_args(args, role), spec, m, task, modal_addr) print(f"[launcher] {name} dials {island_syncer_addrs(cfg.envs)}", flush=True) + if name == held_name: + held_cfg = cfg # defined on the app below; the scheduler spawns it + print(f"[launcher] {name} is the trainer island: the trainer scheduler " + "starts and stops it", flush=True) + continue tasks[name] = cfg modal_cfgs[name] = cfg continue @@ -458,13 +481,14 @@ def resolve(name: str, m: int, rid) -> None: ] for t in threads: t.start() - if modal_cfgs: + if modal_cfgs or held_cfg is not None: from .modal_runner import ModalIslandOps, ModalOps, modal_app_name modal_ops = ModalOps(modal_app_name(prefix)) - for cfg in modal_cfgs.values(): + for cfg in [*modal_cfgs.values(), *([] if held_cfg is None else [held_cfg])]: modal_ops.define(cfg) - print(f"[launcher] deploying Modal app {modal_ops.app_name} ({len(modal_cfgs)} island(s))") + print(f"[launcher] deploying Modal app {modal_ops.app_name} " + f"({len(modal_cfgs) + (held_cfg is not None)} island(s))") modal_ops.deploy() modal_island_ops = ModalIslandOps(modal_ops) for name, cfg in modal_cfgs.items(): @@ -635,9 +659,21 @@ def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: no_sync_incomplete.extend(settle_echo_tapes( event_collectors, echo_names, modal_cfgs, events_dir, modal_tape_dir=modal_tape_dir)) + _record_trainer_start(args, events_dir) try: - exit_codes = controller.run() + exit_codes = trainer_loop.run_with_trainer_loop( + args, controller.run, + island=None if held_name is None else trainer_loop.TrainerIsland( + held_name, specs[held].cloud, held_cfg, modal_ops), + abort=controller.abort) + except trainer_loop.TrainerLoopGaveUp as error: + # the trainer scheduler released the trainer and stopped; the samplers + # cannot finish the run usefully: tear everything down + print(f"[launcher] ERROR: {error}; stopping the run (exit {TRAINER_GAVE_UP_EXIT})", + file=sys.stderr) + drain_tapes(min(NO_SYNC_EVENT_DRAIN_S, FAILED_RUN_DRAIN_S)) + return CONTAINER_CHANGED_EXIT if container_changed() else TRAINER_GAVE_UP_EXIT except RunStalled as error: print(f"[launcher] ERROR: {error}; stopping the run (exit {RUN_STALLED_EXIT})", file=sys.stderr) @@ -802,6 +838,11 @@ def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: if not teardown_island(sky, name, head=head_mode, collect=None, num_nodes=nodes_by_cluster.get(name, 1)): unverified.append(name) + # rl-trainer-cold-start-cost L0.2: sky islands without an echoed tape + # (their tapes are in the pre-teardown diagnostics, also for islands + # the controller tore down earlier) still give the manifest T_s. + _record_sky_trainer_start( + args, [n for n in nodes_by_cluster if n not in modal_cfgs and n not in echo_names]) if modal_ops is not None: # Belt and braces: stop the whole per-run Modal app so no # function call of this run outlives the launcher. @@ -891,6 +932,10 @@ def dry_run_plan(args) -> dict: head = getattr(args, "controller", "local") == "head" if no_sync and head: raise ValueError("--rl-single-island-no-sync has no syncer; use --controller local") + from .launch import trainer_loop + from .rl.engine import sample_roles + + trainer_loop.check(args) # the launch refuses the same scheduler combinations specs = parse_gpu_spec(args.gpu) external = max(0, getattr(args, "external_learners", 0) or 0) islands = [] @@ -926,6 +971,11 @@ def dry_run_plan(args) -> dict: entry["island_role"] = {"role": role, "billing": billing(role), "rejoin": rejoin_source(role)} + if sample_roles.is_pair(args): # L5.4: per-island role, learner 0 of 1 + entry["sample_role"] = sample_roles.island_args( + args, learner_id, len(specs)).rl_sample_role + if learner_id == trainer_loop.held_island(args): + entry["held_by_trainer_scheduler"] = True # the fleet never starts it if task_plan is not None: entry["island_tasks"] = task_plan[learner_id].record() if island_overrides.get(learner_id): @@ -951,6 +1001,10 @@ def dry_run_plan(args) -> dict: "role_syncers": [] if not rl or no_sync else _dry_run_role_syncers( args, len(specs) + external), "outer_sync": not no_sync, + # rl-trainer-start-stop L5.3b/L5.4: the head's scheduler and the island it owns + "trainer_scheduler": (None if not trainer_loop.enabled(args) else { + "mode": args.rl_trainer_scheduler, "island": trainer_loop.held_island(args), + "interval_s": args.rl_trainer_scheduler_interval}), "algorithm_spec_sha256": getattr(args, "rl_expected_algorithm_sha256", None), "algorithm_spec": getattr(args, "rl_algorithm_spec_json", None), "unverified_mechanisms": list(getattr(args, "rl_allow_unverified_mechanism", None) or ()), diff --git a/yeto/provenance.py b/yeto/provenance.py index 1beed779b..c7950cc2b 100644 --- a/yeto/provenance.py +++ b/yeto/provenance.py @@ -117,6 +117,11 @@ def python_spec_sha256(spec: str, *, base_dir: str | Path | None = None) -> str: return file_sha256(python_spec_path(spec, base_dir=base_dir)) +# rl-identity-tiering L9.1 transition: until run code packs exist, an RL run's +# ``code_pack_digest`` is the attested :func:`source_tree_sha256`. +CODE_PACK_MODE_SOURCE_TREE = "source-tree" + + def source_tree_sha256(package_root: str | Path | None = None) -> str: """Stable digest of the installed Yeto Python source tree.""" diff --git a/yeto/rl/adapters/miles/algorithm_flags.py b/yeto/rl/adapters/miles/algorithm_flags.py index d4f7ab1c1..5e534a04f 100644 --- a/yeto/rl/adapters/miles/algorithm_flags.py +++ b/yeto/rl/adapters/miles/algorithm_flags.py @@ -140,17 +140,9 @@ def use_kl_loss(spec): def _correction_rows() -> list[FlagMapping]: def use_tis(spec): - c = spec.correction - if c.method not in ("tis", "custom"): - return [] - argv = ["--use-tis"] - if c.tis_clip is not None: - argv += ["--tis-clip", _num(c.tis_clip)] - if c.tis_clip_low is not None: - argv += ["--tis-clip-low", _num(c.tis_clip_low)] - if c.function is not None: - argv += ["--custom-tis-function-path", c.function.path] - return argv + from .offpolicy_correction import tis_argv # neutral fields -> Miles (D3) + + return tis_argv(spec) def use_opsm(spec): c = spec.correction @@ -161,10 +153,10 @@ def use_opsm(spec): return [ FlagMapping("--use-tis", "correction.method", True, lambda raw: True, lambda _: [("correction.method", "tis")], use_tis), - FlagMapping("--tis-clip", "correction.tis_clip", False, _float, - lambda v: [("correction.tis_clip", v)], lambda spec: []), - FlagMapping("--tis-clip-low", "correction.tis_clip_low", False, _float, - lambda v: [("correction.tis_clip_low", v)], lambda spec: []), + FlagMapping("--tis-clip", "correction.upper", False, _float, + lambda v: [("correction.upper", v)], lambda spec: []), + FlagMapping("--tis-clip-low", "correction.lower", False, _float, + lambda v: [("correction.lower", v)], lambda spec: []), FlagMapping("--custom-tis-function-path", "correction.function", False, _plugin, lambda v: [("correction.method", "custom"), ("correction.function", v)], lambda spec: []), diff --git a/yeto/rl/adapters/miles/buffer_rollout.py b/yeto/rl/adapters/miles/buffer_rollout.py new file mode 100644 index 000000000..46ba2fe98 --- /dev/null +++ b/yeto/rl/adapters/miles/buffer_rollout.py @@ -0,0 +1,403 @@ +"""Miles rollout through the durable sample buffer (rl-trainer-start-stop L4.5). + +Two Miles rollout functions (``--rollout-function-path``, classes in +:mod:`.buffer_rollout_miles`) connect a sampler island and a trainer island +through a :class:`~yeto.rl.engine.sample_buffer.FileSampleBuffer`: + +* sampler (``BufferWriterRolloutFn``): Miles generates the round as usual; every + group is written to the buffer as one payload plus one index entry + (:class:`GroupWriter`). The payload keeps each sample whole, so the per-token + behavior log-probs (``rollout_log_probs``) and the per-token engine versions + (``weight_versions``) reach the trainer for the D2 version gate and decoupled + PPO. The entry's ``behavior_version`` is the OLDEST engine version in the group + (the start version; P0-S schema 4). +* trainer (``BufferReaderRolloutFn``): the round is selected from the buffer, not + generated (:class:`BatchReader`). Groups older than ``max_policy_age`` versions + are passed over (consumed, not trained) and counted. + +The consumed watermark: the reader remembers the buffer position after each +rollout id. The round cut before rollout ``N`` (all rollouts below ``N`` are +trained) commits the position after rollout ``N - 1`` and writes it into the cut +record (``RoundCutCheckpoint.buffer_watermark``); a restore resets the reader to +the watermark of the restored cut (``RoundCutCheckpoint.buffer_resume``). The +driver reaches the reader in the rollout executor's process with +:func:`executor_watermark_before` / :func:`executor_reset` (``__ray_call__``, +``MilesRolloutPool._in_executor``); :func:`wire_sample_buffer` connects both. + +This module does not import Miles (it is tested without it). + +Limits (first implementation): + +* one sampler island writes a buffer (``FileSampleBuffer`` serializes appenders + with ``flock``, which holds only on one host) and one trainer reads it; +* the buffer lives in the checkpoint store (:func:`~yeto.rl.engine.resume.store_for` + at ``BUFFER_ROOT_ENV``): across hosts both islands must mount the same store + (bucket mount or Modal Volume); visibility is the store's ``commit``/``reload``; +* the trainer does no evaluation through the buffer (eval rounds are refused); +* ``--fully-async`` is not supported on either side (it selects its own rollout + function). +""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from yeto.rl.engine.policy_versions import engine_version_number +from yeto.rl.engine.trainer_lifecycle import SampleIndexEntry + +SAMPLER, TRAINER = "sampler", "trainer" +ROLES = (SAMPLER, TRAINER) + +BUFFER_ROOT_ENV = "YETO_SAMPLE_BUFFER_ROOT" # checkpoint store path holding sample-buffer/ +SAMPLER_ID_ENV = "YETO_SAMPLE_BUFFER_SAMPLER_ID" # prefix of the sample ids a sampler writes +WAIT_ENV = "YETO_SAMPLE_BUFFER_WAIT_S" # reader: max seconds to wait for a full round +POLL_ENV = "YETO_SAMPLE_BUFFER_POLL_S" # reader: seconds between index reads while waiting +MAX_POLICY_AGE_ENV = "YETO_SAMPLE_BUFFER_MAX_POLICY_AGE" +# the rollout functions run in Ray workers, which inherit the raylet's environment: +# connect_island_ray forwards these to every worker +ISLAND_ENVS = (BUFFER_ROOT_ENV, SAMPLER_ID_ENV, WAIT_ENV, POLL_ENV, MAX_POLICY_AGE_ENV) + +DEFAULT_WAIT_S = 3600.0 +DEFAULT_POLL_S = 2.0 +TRAINER_CONSUMER = "trainer" + +_MODULE = "yeto.rl.adapters.miles.buffer_rollout_miles" +ROLLOUT_FUNCTION_PATHS = { + SAMPLER: f"{_MODULE}.BufferWriterRolloutFn", + TRAINER: f"{_MODULE}.BufferReaderRolloutFn", +} +PAYLOAD_SCHEMA = "yeto.rl.sample-buffer-payload/v1" +_NDARRAY = "__ndarray__" + + +class SampleBufferRolloutError(RuntimeError): + pass + + +# ---------------------------------------------------------------------- argv +def rollout_argv(role: str | None, *, max_policy_age: int, fully_async: bool = False, + rollout_function_path: str | None = None) -> tuple[str, ...]: + """The Miles flags of an island that writes (``sampler``) or reads (``trainer``) + the sample buffer; ``()`` when ``role`` is None. Raises ``ValueError`` on a + combination the buffer rollout does not support.""" + if role is None: + return () + if role not in ROLES: + raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") + if fully_async: + raise ValueError("the sample buffer rollout does not support --fully-async " + "(it selects its own rollout function)") + if rollout_function_path is not None: + raise ValueError(f"the sample buffer rollout replaces the rollout function; " + f"{rollout_function_path!r} was also given") + if role == TRAINER and int(max_policy_age) < 1: + raise ValueError("a trainer that reads the sample buffer trains on samples of earlier " + f"versions: it needs max_policy_age >= 1 (got {max_policy_age})") + return ("--rollout-function-path", ROLLOUT_FUNCTION_PATHS[role]) + + +def island_env(role: str, *, store_root: str, sampler_id: str | None = None, + max_policy_age: int | None = None, wait_s: float | None = None) -> dict[str, str]: + """The environment the launch layer exports on the island (read in the rollout process).""" + if role not in ROLES: + raise ValueError(f"sample buffer role {role!r}: expected one of {ROLES}") + env = {BUFFER_ROOT_ENV: str(store_root)} + if role == SAMPLER: + if not sampler_id: + raise ValueError("a sampler needs a sampler id (sample id prefix)") + env[SAMPLER_ID_ENV] = str(sampler_id) + else: + if max_policy_age is None: + raise ValueError("a trainer needs max_policy_age") + env[MAX_POLICY_AGE_ENV] = str(int(max_policy_age)) + if wait_s is not None: + env[WAIT_ENV] = f"{float(wait_s):g}" + return env + + +def open_buffer(environ: Mapping[str, str] | None = None, *, batch_size: int = 1) -> Any: + """The island's buffer, from :data:`BUFFER_ROOT_ENV` (store chosen by ``store_for``).""" + from yeto.rl.engine.resume import store_for + from yeto.rl.engine.sample_buffer import FileSampleBuffer + + env = os.environ if environ is None else environ + root = env.get(BUFFER_ROOT_ENV) + if not root: + raise SampleBufferRolloutError(f"{BUFFER_ROOT_ENV} is not set") + return FileSampleBuffer(store_for(root, environ=env), batch_size=batch_size) + + +# ---------------------------------------------------------------------- payload +def _encode_value(value: Any) -> Any: + tolist = getattr(value, "tolist", None) + if callable(tolist) and hasattr(value, "dtype") and hasattr(value, "shape"): + import numpy + + array = numpy.ascontiguousarray(value) + return {_NDARRAY: base64.b64encode(array.tobytes()).decode("ascii"), + "dtype": str(array.dtype), "shape": list(array.shape)} + if hasattr(value, "value") and type(value).__module__ != "builtins": # enum + return value.value + raise TypeError(f"sample field of type {type(value).__name__} is not serializable") + + +def _decode_hook(obj: dict[str, Any]) -> Any: + if _NDARRAY in obj: + import numpy + + data = base64.b64decode(obj[_NDARRAY]) + return numpy.frombuffer(data, dtype=obj["dtype"]).reshape(obj["shape"]).copy() + return obj + + +def encode_group(samples: Sequence[Mapping[str, Any]]) -> bytes: + """One payload: the group's samples as ``Sample.to_dict()`` mappings.""" + record = {"schema": PAYLOAD_SCHEMA, "samples": [dict(s) for s in samples]} + return json.dumps(record, default=_encode_value, sort_keys=True, allow_nan=True).encode() + + +def decode_group(data: bytes) -> list[dict[str, Any]]: + record = json.loads(data, object_hook=_decode_hook) + if record.get("schema") != PAYLOAD_SCHEMA: + raise SampleBufferRolloutError(f"payload schema {record.get('schema')!r} != {PAYLOAD_SCHEMA}") + return list(record["samples"]) + + +def _sample_versions(sample: Mapping[str, Any]) -> list[int]: + versions = [] + for call in sample.get("weight_versions") or (): + for span in call or (): + versions.append(span["version"]) + if not versions and (sample.get("metadata") or {}).get("weight_version") is not None: + versions.append(sample["metadata"]["weight_version"]) + numbers = [engine_version_number(v) for v in versions] + if not numbers or None in numbers: + raise SampleBufferRolloutError( + f"sample {sample.get('index')}: no numeric engine weight version ({versions!r}); " + "the buffer needs the behavior version of every sample") + return [int(n) for n in numbers] + + +def group_entry_fields(samples: Sequence[Mapping[str, Any]]) -> tuple[int, int]: + """``(behavior_version, tokens)`` of a group: the oldest engine version of any + token, and the response tokens. Every sample must carry per-token behavior + log-probs (decoupled PPO needs them).""" + if not samples: + raise SampleBufferRolloutError("empty group") + oldest, tokens = None, 0 + for sample in samples: + length = int(sample.get("response_length") or 0) + logprobs = sample.get("rollout_log_probs") + if logprobs is None or len(logprobs) != length: + raise SampleBufferRolloutError( + f"sample {sample.get('index')}: rollout_log_probs " + f"{'missing' if logprobs is None else len(logprobs)} != response_length {length}") + version = min(_sample_versions(sample)) + oldest = version if oldest is None else min(oldest, version) + tokens += length + return int(oldest), tokens + + +def renumber_group(samples: Sequence[dict[str, Any]], *, position: int, + samples_per_group: int) -> list[dict[str, Any]]: + """Give a group read from the buffer an identity unique in the buffer. + + Miles numbers groups and samples per sampler process from 0, so groups from two + sampler rounds (or a restarted sampler) repeat them; the trainer's batch checks + and trajectory ids need them unique. The buffer position is unique for the + buffer's life and the same after a restore re-reads it, so ``group_index`` is + the position and ``index`` is ``position * samples_per_group + j``. The + sampler's numbers stay in ``metadata["yeto_sampler_identity"]``.""" + if len(samples) > int(samples_per_group): + raise SampleBufferRolloutError( + f"group at buffer position {position} has {len(samples)} samples > " + f"samples per group {samples_per_group}") + out = [] + for j, sample in enumerate(samples): + sample = dict(sample) + metadata = dict(sample.get("metadata") or {}) + metadata["yeto_sampler_identity"] = {"group_index": sample.get("group_index"), + "index": sample.get("index")} + sample.update(metadata=metadata, group_index=int(position), + index=int(position) * int(samples_per_group) + j) + out.append(sample) + return out + + +# ---------------------------------------------------------------------- sampler side +class GroupWriter: + """Writes finished groups to the buffer: payload first, then the index entry.""" + + def __init__(self, buffer: Any, sampler_id: str) -> None: + if not sampler_id: + raise ValueError("sampler_id must be non-empty") + self.buffer = buffer + self.sampler_id = str(sampler_id) + + def write(self, samples: Sequence[Mapping[str, Any]], *, group_key: str) -> SampleIndexEntry: + behavior_version, tokens = group_entry_fields(samples) + sample_id = f"{self.sampler_id}-{group_key}" + ref = self.buffer.write_payload(sample_id, encode_group(samples)) + entry = SampleIndexEntry(sample_id=sample_id, behavior_version=behavior_version, + payload_ref=ref, tokens=tokens) + self.buffer.append(entry) + return entry + + +# ---------------------------------------------------------------------- trainer side +@dataclass +class TakeReport: + rollout_id: int + groups: int + stale_groups: int # passed over: older than max_policy_age (consumed, not trained) + start: int # buffer position of the first selected entry + end: int # buffer position after the last selected entry + behavior_versions: dict[int, int] = field(default_factory=dict) + wait_s: float = 0.0 + + def metrics(self) -> dict[str, float]: + return {"sample_buffer/groups": float(self.groups), + "sample_buffer/stale_groups": float(self.stale_groups), + "sample_buffer/start": float(self.start), "sample_buffer/end": float(self.end), + "sample_buffer/wait_s": float(self.wait_s)} + + +class BatchReader: + """Selects the trainer's rounds from the buffer (single consumer).""" + + def __init__(self, buffer: Any, *, max_policy_age: int, wait_s: float = DEFAULT_WAIT_S, + poll_s: float = DEFAULT_POLL_S, consumer: str = TRAINER_CONSUMER, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep) -> None: + if int(max_policy_age) < 1: + raise ValueError("max_policy_age must be >= 1") + self.buffer = buffer + self.max_policy_age = int(max_policy_age) + self.wait_s = float(wait_s) + self.poll_s = float(poll_s) + self.consumer = consumer + self.clock = clock + self.sleep = sleep + self._ends: dict[int, int] = {} # rollout id -> buffer position after its round + + def take(self, rollout_id: int, groups: int, *, current_version: int | None + ) -> tuple[list[tuple[int, list[dict[str, Any]]]], TakeReport]: + """``groups`` ``(buffer position, decoded group)`` pairs for ``rollout_id``. + Waits up to ``wait_s`` for them; after the data end returns what is left + (raises when nothing is).""" + if int(groups) < 1: + raise ValueError("groups must be positive") + rollout_id = int(rollout_id) + started = self.clock() + start = self.buffer.read_cursor + out: list[tuple[int, list[dict[str, Any]]]] = [] + versions: dict[int, int] = {} + stale = 0 + floor = None if current_version is None else int(current_version) - self.max_policy_age + while len(out) < groups: + selected = self.buffer.select(self.consumer, groups - len(out)) + for item in selected: + if floor is not None and item.entry.behavior_version < floor: + stale += 1 + continue + out.append((item.position, + decode_group(self.buffer.read_payload(item.entry.payload_ref)))) + versions[item.entry.behavior_version] = versions.get(item.entry.behavior_version, 0) + 1 + if len(out) >= groups or selected: + continue + if self.buffer.data_ended: + break + if self.clock() - started > self.wait_s: + raise SampleBufferRolloutError( + f"rollout {rollout_id}: {len(out)}/{groups} groups after {self.wait_s:g} s " + f"(buffer position {self.buffer.read_cursor})") + self.sleep(self.poll_s) + end = self.buffer.read_cursor + if not out: + raise SampleBufferRolloutError(f"rollout {rollout_id}: the sample buffer ended " + f"with no trainable group (position {end})") + self._ends[rollout_id] = end + return out, TakeReport(rollout_id=rollout_id, groups=len(out), stale_groups=stale, + start=start, end=end, behavior_versions=dict(sorted(versions.items())), + wait_s=round(self.clock() - started, 3)) + + def watermark_before(self, rollout_id: int) -> int: + """Commit and return the position after the last round below ``rollout_id`` + (every round below it is trained when the cut before ``rollout_id`` runs).""" + done = [rid for rid in self._ends if rid < int(rollout_id)] + if done: + self.buffer.commit_consumed(self._ends[max(done)]) + for rid in done: + del self._ends[rid] + return int(self.buffer.watermark) + + def reset(self, watermark: int) -> int: + """Continue at the restored cut's watermark (rounds after it are selected again).""" + self.buffer.reset_to(int(watermark)) + self._ends.clear() + return int(self.buffer.watermark) + + +# ---------------------------------------------------------------------- executor link +_ACTIVE: dict[str, BatchReader] = {} + + +def register_reader(reader: BatchReader) -> None: + """Called by the reader rollout function in the rollout executor's process.""" + _ACTIVE[TRAINER] = reader + + +def _active_reader() -> BatchReader: + reader = _ACTIVE.get(TRAINER) + if reader is None: + raise SampleBufferRolloutError("no sample buffer reader in this rollout process") + return reader + + +def executor_watermark_before(_executor: Any, rollout_id: int) -> int: + """Runs INSIDE the rollout executor actor (``__ray_call__``).""" + return _active_reader().watermark_before(int(rollout_id)) + + +def executor_read_cursor(_executor: Any) -> int | None: + """Runs INSIDE the rollout executor actor: the buffer position after the last + selected group (None before the reader exists).""" + reader = _ACTIVE.get(TRAINER) + return None if reader is None else int(reader.buffer.read_cursor) + + +def executor_reset(_executor: Any, watermark: int) -> int: + """Runs INSIDE the rollout executor actor (``__ray_call__``).""" + return _active_reader().reset(int(watermark)) + + +class BufferCutLink: + """Driver side: the round cut's watermark and restore hooks over ``in_executor`` + (``MilesRolloutPool._in_executor``: runs ``fn(executor, *args)`` in the actor).""" + + def __init__(self, in_executor: Callable[..., Any]) -> None: + self.in_executor = in_executor + + def watermark(self, rollout_id: int) -> int: + return int(self.in_executor(executor_watermark_before, int(rollout_id))) + + def resume(self, watermark: int) -> int: + return int(self.in_executor(executor_reset, int(watermark))) + + +def wire_sample_buffer(checkpoint: Any, rollout: Any) -> BufferCutLink: + """Give a ``RoundCutCheckpoint`` the trainer's buffer watermark and restore hook.""" + in_executor = getattr(rollout, "_in_executor", None) + if not callable(in_executor): + raise SampleBufferRolloutError(f"rollout pool {type(rollout).__name__} cannot run " + "code in the rollout executor") + link = BufferCutLink(in_executor) + checkpoint.buffer_watermark = link.watermark + checkpoint.buffer_resume = link.resume + return link diff --git a/yeto/rl/adapters/miles/buffer_rollout_miles.py b/yeto/rl/adapters/miles/buffer_rollout_miles.py new file mode 100644 index 000000000..f796be748 --- /dev/null +++ b/yeto/rl/adapters/miles/buffer_rollout_miles.py @@ -0,0 +1,110 @@ +"""Miles rollout functions for the durable sample buffer (L4.5). + +Loaded by Miles in the rollout executor process (``--rollout-function-path``, set +by :func:`.buffer_rollout.rollout_argv`); the contract and the logic live in +:mod:`.buffer_rollout`. +""" + +from __future__ import annotations + +import asyncio +import os +import uuid +from typing import Any + +from miles.rollout.base_types import ( + BaseRolloutFn, + RolloutFnConstructorInput, + RolloutFnInput, + RolloutFnTrainOutput, +) +from miles.rollout.inference_rollout.inference_rollout_common import InferenceRolloutFn +from miles.utils.function_registry import load_function +from miles.utils.types import Sample + +from .buffer_rollout import ( + DEFAULT_POLL_S, + DEFAULT_WAIT_S, + MAX_POLICY_AGE_ENV, + POLL_ENV, + SAMPLER_ID_ENV, + WAIT_ENV, + BatchReader, + GroupWriter, + SampleBufferRolloutError, + open_buffer, + register_reader, + renumber_group, +) + + +def _flat(group: Any) -> list[Any]: + out: list[Any] = [] + for item in group if isinstance(group, list) else [group]: + out.extend(_flat(item) if isinstance(item, list) else [item]) + return out + + +class BufferWriterRolloutFn(InferenceRolloutFn): + """Sampler: Miles' rollout; every finished group is also written to the buffer.""" + + def __init__(self, input: RolloutFnConstructorInput) -> None: + super().__init__(input) + sampler_id = os.environ.get(SAMPLER_ID_ENV) + if not sampler_id: + raise SampleBufferRolloutError(f"{SAMPLER_ID_ENV} is not set") + # a per-process suffix: a restarted sampler repeats Miles group indexes, so the + # sample ids stay unique; the trainer renumbers groups by buffer position + self._writer = GroupWriter(open_buffer(), f"{sampler_id}-{uuid.uuid4().hex[:8]}") + + async def _call_train(self, input: Any) -> Any: + output = await super()._call_train(input) + await asyncio.to_thread(self._write, input.rollout_id, output.samples) + return output + + def _write(self, rollout_id: int, groups: list[Any]) -> None: + for group in groups: + samples = _flat(group) + if any(s.status == Sample.Status.ABORTED for s in samples): + continue # an aborted group is regenerated, never trained + key = f"r{int(rollout_id)}-g{samples[0].group_index}" + self._writer.write([s.to_dict() for s in samples], group_key=key) + + +class BufferReaderRolloutFn(BaseRolloutFn): + """Trainer: the round is selected from the buffer, nothing is generated.""" + + def __init__(self, input: RolloutFnConstructorInput) -> None: + super().__init__(input) + args = self.args = input.args + age = os.environ.get(MAX_POLICY_AGE_ENV) + if age is None: + raise SampleBufferRolloutError(f"{MAX_POLICY_AGE_ENV} is not set") + self.groups = int(args.rollout_batch_size) + self.samples_per_group = int(args.n_samples_per_prompt) + self.reader = BatchReader( + open_buffer(batch_size=self.groups), max_policy_age=int(age), + wait_s=float(os.environ.get(WAIT_ENV, DEFAULT_WAIT_S)), + poll_s=float(os.environ.get(POLL_ENV, DEFAULT_POLL_S))) + register_reader(self.reader) + + def __call__(self, input: RolloutFnInput) -> RolloutFnTrainOutput: + if input.evaluation: + raise SampleBufferRolloutError("the sample buffer reader does not evaluate; " + "set --eval-function-path or disable eval") + groups, report = self.reader.take(input.rollout_id, self.groups, + current_version=input.weight_version) + samples = [[Sample.from_dict(d) for d in renumber_group( + group, position=position, samples_per_group=self.samples_per_group)] + for position, group in groups] + self._run_round_hooks(samples) + return RolloutFnTrainOutput(samples=samples, metrics=report.metrics()) + + def _run_round_hooks(self, samples: list[list[Any]]) -> None: + """Call the end-of-round hooks Miles' generating rollout calls, so the yeto + rollout metadata reaches the driver. Every selected group is trained, so the + kept groups and all groups are the same list; there is no data source.""" + if f := load_function(self.args.rollout_sample_filter_path): + f(self.args, samples) + if f := load_function(self.args.rollout_all_samples_process_path): + f(self.args, samples, None) diff --git a/yeto/rl/adapters/miles/carry_over.py b/yeto/rl/adapters/miles/carry_over.py index 215829f32..4e4eca4a0 100644 --- a/yeto/rl/adapters/miles/carry_over.py +++ b/yeto/rl/adapters/miles/carry_over.py @@ -39,6 +39,12 @@ from collections.abc import Callable, Sequence from typing import Any +from yeto.rl.engine.offpolicy_correction import ( + GATE_GROUP_DROPPED, + GATE_DROPPED_LENGTHS, + GATE_UNKNOWN_VERSION_DROPPED, + TRUNCATED_BY_TURN, +) from yeto.rl.engine.policy_versions import engine_version_number from yeto.rl.engine.version_segments import ( ProvenanceError, @@ -48,6 +54,7 @@ ) CARRY_STATS_ATTR = "_yeto_carry_over_stats" +DROPPED_LENGTHS = "dropped_lengths" CARRY_FIELDS = ( "carried_in_groups", "carried_in_trajectories", @@ -129,14 +136,19 @@ def _trajectories(group: Sequence[Any]) -> int: return sum(1 for s in _flat(group) if int(getattr(s, "response_length", 0) or 0) > 0) -def _stats_for(args: Any, current_version: int) -> dict[str, int]: +def _stats_for(args: Any, current_version: int) -> dict[str, Any]: stats = getattr(args, CARRY_STATS_ATTR, None) if not isinstance(stats, dict) or stats.get("version") != current_version: - stats = {"version": current_version, **{k: 0 for k in CARRY_FIELDS}} + stats = {"version": current_version, **{k: 0 for k in CARRY_FIELDS}, + DROPPED_LENGTHS: {GATE_GROUP_DROPPED: [], GATE_UNKNOWN_VERSION_DROPPED: []}} setattr(args, CARRY_STATS_ATTR, stats) return stats +def _trajectory_lengths(group: Sequence[Any]) -> list[int]: + return [n for n in (int(getattr(s, "response_length", 0) or 0) for s in _flat(group)) if n] + + def carry_over_buffer_filter(args: Any, buffer: list, num_samples: int, *, current_version: int, limit: int) -> list: """Select up to ``num_samples`` buffered groups to continue at ``current_version``. @@ -150,8 +162,10 @@ def carry_over_buffer_filter(args: Any, buffer: list, num_samples: int, *, known, unknown = group_versions(group) if unknown: stats["unknown_version_discarded_groups"] += 1 + stats[DROPPED_LENGTHS][GATE_UNKNOWN_VERSION_DROPPED] += _trajectory_lengths(group) continue if known and current_version - min(known) > limit: + stats[DROPPED_LENGTHS][GATE_GROUP_DROPPED] += _trajectory_lengths(group) stats["over_age_discarded_groups"] += 1 stats["over_age_discarded_trajectories"] += _trajectories(group) stats["over_age_discarded_tokens"] += _response_tokens(group) @@ -167,14 +181,30 @@ def carry_over_buffer_filter(args: Any, buffer: list, num_samples: int, *, return selected -def take_carry_stats(args: Any, current_version: int | None) -> dict[str, int]: - """This round's carry-over counts (zeros when nothing was buffered); consumed.""" +def take_carry_round(args: Any, current_version: int | None + ) -> tuple[dict[str, int], dict[str, list[int]]]: + """This round's carry-over counts and the response lengths of the + trajectories discarded per ``offpolicy_gate`` reason; consumed.""" stats = getattr(args, CARRY_STATS_ATTR, None) if hasattr(args, CARRY_STATS_ATTR): setattr(args, CARRY_STATS_ATTR, None) if not isinstance(stats, dict) or stats.get("version") != current_version: - return {k: 0 for k in CARRY_FIELDS} - return {k: int(stats[k]) for k in CARRY_FIELDS} + return ({k: 0 for k in CARRY_FIELDS}, + {GATE_GROUP_DROPPED: [], GATE_UNKNOWN_VERSION_DROPPED: []}) + lengths = stats.get(DROPPED_LENGTHS) or {} + return ({k: int(stats[k]) for k in CARRY_FIELDS}, + {reason: list(lengths.get(reason, ())) + for reason in (GATE_GROUP_DROPPED, GATE_UNKNOWN_VERSION_DROPPED)}) + + +def take_carry_stats(args: Any, current_version: int | None) -> dict[str, Any]: + """This round's carry-over counts (zeros when nothing was buffered); consumed. + + Also carries the dropped response lengths under :data:`GATE_DROPPED_LENGTHS` + (rl-offpolicy-correction-passthrough D6); the driver turns them into + ``offpolicy_gate`` and removes the key from ``rl_rollout_carry_over``.""" + counts, lengths = take_carry_round(args, current_version) + return {**counts, GATE_DROPPED_LENGTHS: lengths} def group_versions_record(group: Sequence[Any]) -> dict[str, Any]: @@ -293,6 +323,7 @@ def estimate_cross_version_truncation(args: Any, samples: Sequence[Any], current unknown_tokens = 0 near_zero = 0 near_zero_examples: list[dict[str, Any]] = [] + by_turn: dict[int, list[int]] = {} # turn index -> [truncated, cross-version tokens] def fail(why: str) -> None: nonlocal failed @@ -325,11 +356,16 @@ def fail(why: str) -> None: tuple(min(float(p), 0.0) for p in generated)) corrections.append(cross_version_is(provenance, current, current_version, clip_low=clip_low, clip_high=clip_high)) + turns = turn_indices(sample, len(versions)) for i, (v, g, c) in enumerate(zip(provenance.versions, provenance.logprobs, current, strict=True)): if v == current_version: continue r = math.exp(float(c) - g) ratios.append(r) + if turns[i] is not None: + counts = by_turn.setdefault(turns[i], [0, 0]) + counts[0] += not clip_low <= r <= clip_high + counts[1] += 1 if r < NEAR_ZERO_RATIO: # 5.5 follow-up: where do zero ratios come from near_zero += 1 if len(near_zero_examples) < NEAR_ZERO_EXAMPLES: @@ -342,6 +378,8 @@ def fail(why: str) -> None: extra["cross_version_unscored_reasons"] = dict(sorted(reasons.items())[:8]) if unknown_tokens: extra["cross_version_unknown_version_tokens"] = unknown_tokens + if by_turn: # rl-offpolicy-correction D6 (absent when none: old key set kept) + extra[TRUNCATED_BY_TURN] = {turn: tuple(c) for turn, c in sorted(by_turn.items())} if near_zero: # absent when none (old key set kept) extra["cross_version_ratio_near_zero"] = near_zero extra["cross_version_ratio_near_zero_examples"] = near_zero_examples @@ -355,6 +393,22 @@ def fail(why: str) -> None: } +def turn_indices(sample: Any, n: int) -> list[int | None]: + """Per response token: its model turn (0, 1, ...; runs of ``loss_mask`` 1), + None for a token outside the loss (tool output). No mask: one turn.""" + mask = list(getattr(sample, "loss_mask", None) or ()) + if len(mask) != n: + return [0] * n + out: list[int | None] = [] + turn, inside = -1, False + for m in mask: + if m and not inside: + turn += 1 + inside = bool(m) + out.append(turn if inside else None) + return out + + NEAR_ZERO_RATIO = 1e-6 NEAR_ZERO_EXAMPLES = 8 diff --git a/yeto/rl/adapters/miles/config.py b/yeto/rl/adapters/miles/config.py index 35a749d5b..b0f4da2e0 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -45,6 +45,7 @@ objective_flags, ) from .placement import PlacementRequest, check_placement_not_rewritten +from . import sampler_snapshot from .binding import filter_argv ROLLOUT_META_HOOK_PATH = ( @@ -423,6 +424,7 @@ def _check_sglang_tp(value, config): "serving.page_size": _ok, "serving.max_running_requests": _ok, "serving.chunked_prefill_size": _ok, + "serving.sampler_snapshot": _ok, # gated in translate_run_config (sampler_snapshot) "ports.rollout_engine_base_port": _must_be_none( "upstream Miles has no --rollout-engine-base-port (fork-only port isolation, task 1.4b)" ), @@ -462,6 +464,11 @@ def _check_sglang_tp(value, config): # d2-wire: read by the yeto-side ElasticHook only; no Miles option. "edge_costs_path": _ok, "elastic_window_s": _ok, + # L4.6: the buffer rollout functions (translate_run_config); the island + # environment and the sampler loop read the rest. + "sample_buffer": _ok, + **{f"sample_buffer.{name}": _ok + for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s")}, } @@ -594,6 +601,18 @@ def reject_fault_tolerance_flags(tokens: Sequence[str]) -> None: ) +def _refuse_decoupled_ppo_breakers(tokens: Sequence[str]) -> None: + """rl-offpolicy-correction-passthrough D2: flags that remove pi_prox or mask + tokens by version are refused, with the flag named.""" + from .offpolicy_correction import refuse_forbidden_argv + from yeto.rl.engine.offpolicy_correction import CorrectionConfigError + + try: + refuse_forbidden_argv(tokens) + except CorrectionConfigError as exc: + raise MilesConfigError(str(exc)) from exc + + def check_extra_argv(extra_argv: Sequence[str], algorithm: AlgorithmSpec | None = None): """Refuse FT/adapter-owned flags; absorb mapped algorithm flags. @@ -604,6 +623,7 @@ def check_extra_argv(extra_argv: Sequence[str], algorithm: AlgorithmSpec | None """ reject_fault_tolerance_flags(extra_argv) + _refuse_decoupled_ppo_breakers(extra_argv) mapped = mapped_flags() unmapped_objective = objective_flags() - mapped for flag in _flags(extra_argv): @@ -761,7 +781,13 @@ def translate_run_config( agent = config.agent columns = config.data.columns request = placement_request(config) - rollout_overlap = max_policy_age > 0 and request.kind != "colocated" + sample_buffer = getattr(config, "sample_buffer", None) + if sample_buffer is not None and trainable.parameter_mode != "full": + # the sampler loads each release as a full HF checkpoint (update_weights_from_disk) + raise MilesConfigError(f"--rl-sample-role {sample_buffer.role} supports full-parameter " + "mode only (a LoRA release cannot be loaded by the sampler yet)") + # a buffer island selects its own rollout function, so never Miles' fully_async + rollout_overlap = max_policy_age > 0 and request.kind != "colocated" and sample_buffer is None if request.kind == "colocated": placement_values = ["--colocate"] @@ -913,6 +939,8 @@ def translate_run_config( max_policy_age, agentic=bool(getattr(config.agent, "custom_agent_function_path", None)), overlap=rollout_overlap)) values.extend(critic_load_argv(config, algorithm)) + if sample_buffer is not None: + values.extend(sample_buffer_argv(sample_buffer.role, max_policy_age=max_policy_age)) evaluation = config.eval if evaluation is not None: @@ -941,6 +969,11 @@ def translate_run_config( values.extend((flag, str(value))) if serving.deterministic_inference: values.append("--sglang-enable-deterministic-inference") + if getattr(serving, "sampler_snapshot", False): + try: + values.extend(sampler_snapshot.snapshot_argv(request, extra_argv=extra_argv)) + except sampler_snapshot.SamplerSnapshotRefused as exc: + raise MilesConfigError(str(exc)) from exc if parallel.uneven_pipeline_layers is not None: first_layers, last_layers = parallel.uneven_pipeline_layers values.extend( @@ -1078,6 +1111,7 @@ def translate_run_config( values.extend((flag, engine)) values.extend(extra_argv) reject_fault_tolerance_flags(values[1:]) # generated argv never carries FT flags + _refuse_decoupled_ppo_breakers(values) # also when the spec emitted the flag # IR-1: the final argv never pairs the session server with partial rollout # (also re-checked on the parsed namespace by validate_parsed_args). check_session_server_partial_rollout(agent.use_session_server, "--partial-rollout" in values) @@ -1100,6 +1134,27 @@ def translate_run_config( ) +# Miles builds the HF snapshot exporter (actor.export_hf) only when one of its HF +# outputs is set. yeto never calls save_model, so this path is never written; the +# trainer's policy releases call export_hf with their own directory. +RELEASE_EXPORTER_SAVE_HF = "/tmp/yeto-rl-unused-save-hf/rollout_{rollout_id}" + + +SAMPLER_ONLY_FLAG = "--debug-rollout-only" + + +def sample_buffer_argv(role: str, *, max_policy_age: int) -> tuple[str, ...]: + """L4.6: a sampler writes the buffer and builds no trainer model + (``--debug-rollout-only``: rollout GPUs only, engines serve the base weights); + a trainer reads the buffer and can export HF snapshots for its releases.""" + from .buffer_rollout import SAMPLER, rollout_argv + + argv = rollout_argv(role, max_policy_age=max_policy_age) + if role == SAMPLER: + return (*argv, SAMPLER_ONLY_FLAG) + return (*argv, "--save-hf", RELEASE_EXPORTER_SAVE_HF) + + def _full_parameter_optimizer_attrs(trainable: Any) -> dict[str, Any]: """Keep every FP32 master whole in full-parameter mode. @@ -1169,6 +1224,11 @@ def validate_parsed_args( cells = counter(args) if cells != 1: raise SingleCellError(f"actor TrainGroup would have {cells} cells; ports require exactly 1") + if SAMPLER_ONLY_FLAG in launch.argv: + # A sampler island (L4.6) asked for --debug-rollout-only: Miles places the + # rollout engines only and builds no trainer, so the trainer partition of the + # request is unused by design, not rewritten. + return check_placement_not_rewritten(launch.placement, args) diff --git a/yeto/rl/adapters/miles/convertible.py b/yeto/rl/adapters/miles/convertible.py new file mode 100644 index 000000000..96e22a637 --- /dev/null +++ b/yeto/rl/adapters/miles/convertible.py @@ -0,0 +1,138 @@ +"""Miles/SGLang :class:`~yeto.rl.engine.convertible.ConversionPort` (L7.2). + +The calls are the ones of the upstream colocated branch of ``train.py`` that the +adapter already uses (``rollout.py`` ``_offload_after_rollout``, ``publish.py`` +``_publish``, ``trainer.py`` ``onload``/``offload``): + +* stop requests: ``controller.abort_all()`` (Miles puts the aborted groups back + in its data buffer, so they are regenerated); the finished segments come from + ``collect_inflight``, which the sampler side gives (it knows its partial + trajectories); +* release rollout memory: ``controller.offload(tags=[weights, kv_cache, + cuda_graph])`` (SGLang ``release_memory_occupation``, all tags); +* trainer: ``trainer.onload()`` / ``trainer.offload()`` (Megatron wake/sleep with + ``--offload-train``); +* resume rollout: ``controller.onload(tags=[weights])`` (not ``onload_weights``, + which does nothing unless ``--offload-rollout-level`` has ``weight``, while + the release above always frees the weights), then ``load_weights(version)`` + (the existing weight update; in this fork ``begin_weight_update`` -> + ``update_weights_from_tensor`` -> ``end_weight_update``, checked on one GPU in + SESSION24 L8.1), then ``controller.onload_kv()``, then ``read_checksum()``: + the engine weight checksum of the loaded version. The read comes after the KV + onload: in LoRA mode it serves one generate to load the adapter into the LoRA + pool, which crashed the engine scheduler with the KV cache released (G-L7.4). + +:func:`miles_conversion_port` wires the three callbacks from the island's +adapter objects: + +* ``load_weights`` / ``read_checksum``: :meth:`.publish.MilesPublisher.load_version` + and :meth:`.publish.MilesPublisher.served_checksum` (digest of the engines' + common ``check_weights(checksum)`` body; compare with + :attr:`.publish.MilesPublisher.published_checksum` of the same version); +* ``collect_inflight``: :func:`engine_requests_inflight`, the requests running + or queued on the engines (Miles puts their groups back into its data buffer + on ``abort_all`` and generates them again, so every one is an unfinished + trajectory; a Miles group has no finished segment of its own to return); +* ``set_admission``: :meth:`.rollout.MilesRolloutPool.set_admission` (harness + admission; no-op without a harness board). + +Not checked on GPU yet: the whole port (L7.4). +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from yeto.rl.engine.convertible import InflightTrajectory + + +def _weight_tags() -> list[str]: + from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS + + return [GPU_MEMORY_TYPE_WEIGHTS] + + +def _all_memory_tags() -> list[str]: + from sglang.srt.constants import ( + GPU_MEMORY_TYPE_CUDA_GRAPH, + GPU_MEMORY_TYPE_KV_CACHE, + GPU_MEMORY_TYPE_WEIGHTS, + ) + + return [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH] + + +class MilesConversionPort: + """``controller``: the Miles rollout manager handle of this node's engines; + ``trainer``: the adapter trainer (``MilesTrainer``) on the same GPUs; + ``run``: runs a coroutine to completion (the adapter's ``_run``).""" + + def __init__(self, *, controller: Any, trainer: Any, + run: Callable[[Awaitable[Any]], Any], + load_weights: Callable[[int], None], + read_checksum: Callable[[], str | None] = lambda: None, + collect_inflight: Callable[[], list[InflightTrajectory]] = list, + set_admission: Callable[[bool], None] | None = None, + memory_tags: Callable[[], list[str]] = _all_memory_tags, + weight_tags: Callable[[], list[str]] = _weight_tags) -> None: + self._controller = controller + self._trainer = trainer + self._run = run + self._load_weights = load_weights + self._read_checksum = read_checksum + self._collect_inflight = collect_inflight + self._set_admission = set_admission + self._memory_tags = memory_tags + self._weight_tags = weight_tags + + def stop_requests(self) -> list[InflightTrajectory]: + if self._set_admission is not None: + self._set_admission(False) + inflight = self._collect_inflight() # read before abort clears the partial state + self._run(self._controller.abort_all()) + return inflight + + def release_rollout(self) -> None: + self._run(self._controller.offload(tags=self._memory_tags())) + + def onload_trainer(self) -> None: + self._trainer.onload() + + def offload_trainer(self) -> None: + self._trainer.offload() + + def resume_rollout(self, version: int) -> str | None: + self._run(self._controller.onload(tags=self._weight_tags())) + self._load_weights(version) + self._run(self._controller.onload_kv()) + return self._read_checksum() + + def accept_requests(self) -> None: + if self._set_admission is not None: + self._set_admission(True) + + +def engine_requests_inflight(rollout: Any) -> list[InflightTrajectory]: + """The requests running or queued on the engines of ``rollout`` (a + :class:`.rollout.MilesRolloutPool`), one unfinished trajectory each, read from + its load sample. Empty when the load is unknown (no router address, engine + endpoint failed): the count is an observation only, since Miles puts the + aborted groups back into its data buffer whatever this returns.""" + sample = rollout.load_sample() + if sample is None or sample.get("running_requests") is None: + return [] + count = int(sample["running_requests"]) + int(sample.get("queued_requests") or 0) + return [InflightTrajectory(f"engine-request-{i}") for i in range(count)] + + +def miles_conversion_port(*, controller: Any, rollout: Any, trainer: Any, publisher: Any, + run: Callable[[Awaitable[Any]], Any]) -> MilesConversionPort: + """The port of a colocated Miles island: ``controller`` is its inference + controller, ``rollout`` its :class:`.rollout.MilesRolloutPool`, ``trainer`` its + trainer (``onload``/``offload``), ``publisher`` its :class:`.publish.MilesPublisher`.""" + return MilesConversionPort( + controller=controller, trainer=trainer, run=run, + load_weights=publisher.load_version, read_checksum=publisher.served_checksum, + collect_inflight=lambda: engine_requests_inflight(rollout), + set_admission=rollout.set_admission) diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 5da36bd2e..dee59779a 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -8,7 +8,9 @@ One shard per rank holds (4.1 audit, ``cut-audit.md``): -* LoRA adapter model copies (bf16 or fp32), keyed by parameter name; +* the model copies (bf16 or fp32) of the trainable parameters, keyed by + parameter name: the LoRA adapters of a LoRA trainer, every trainable + parameter of a full-parameter trainer (key ``adapter`` in both cases); * the optimizer state keyed by parameter name, including the FP32 main copy (``param``), the moments and the step/hyper-parameters -- produced by the fork-M5 helpers ``miles.backends.megatron_utils.lora.dp_invariant_state`` @@ -25,8 +27,10 @@ ranges of the rebuilt optimizer -- no peer shard is read (TP/PP>1 and EP>1 included). Refused like fork-M5: fp16 (loss-scaler state not saved), precision-aware optimizer, partial DistOpt instances, CP>1, and a trainer -without LoRA adapters; the DP-resharding restore also refuses EP>1 and -TP/PP>1 with DistOpt. torch / miles are imported lazily. +whose trainable parameters mix LoRA adapters with others; the DP-resharding +restore also refuses EP>1 and TP/PP>1 with DistOpt. A full-parameter shard +holds the whole model copy on every DP rank plus this rank's optimizer range +(Qwen3-0.6B at DP=1, estimated: 1.2 GB bf16 model + 7.2 GB FP32 master and moments). torch / miles are imported lazily. """ from __future__ import annotations @@ -262,13 +266,25 @@ def _require(actor: Any, backend: Any, *, reshard: bool = False) -> dict[str, in def _adapters(actor: Any, backend: Any) -> list[tuple[str, Any]]: - named = [(n, p) for n, p in backend.named_parameters(actor.model) if backend.is_adapter(n)] - if not named: - raise CutPluginError("no LoRA adapter parameters: only LoRA trainers are covered by the cut (4.1)") - trainable = {id(p) for _, p in backend.named_parameters(actor.model) if p.requires_grad} - if {id(p) for _, p in named} != trainable: + """The parameters a cut carries: the LoRA adapters of a LoRA trainer, else every + trainable parameter of a full-parameter trainer (the shard's ``adapter`` entry + then holds the full model copy). A mix -- LoRA adapters plus other trainable + parameters -- is refused: neither set alone is the trainer's state.""" + every = backend.named_parameters(actor.model) + trainable = [(n, p) for n, p in every if p.requires_grad] + adapters = [(n, p) for n, p in every if backend.is_adapter(n)] + if not adapters: + if not trainable: + raise CutPluginError("the trainer has no trainable parameters: nothing to cut") + return trainable + if {id(p) for _, p in adapters} != {id(p) for _, p in trainable}: raise CutPluginError("trainable parameters other than LoRA adapters exist; the cut would miss them") - return named + return adapters + + +def _tuning(backend: Any, named: list) -> str: + """``lora`` or ``full`` (what :func:`_adapters` selected).""" + return "lora" if any(backend.is_adapter(n) for n, _ in named) else "full" # -------------------------------------------------------------------------- @@ -569,6 +585,7 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "weight_version": (snap.get("miles_counters") or {}).get("weight_version"), "rng_digest": state_digest(rng), "components": component_digests({"state": snap, "rng": rng}), + "tuning": _tuning(backend, named), "adapter_tensors": len(named), "adapter_names": sorted(n for n, _ in named), "optimizer_names": sorted((snap["optimizer_named"] or {}).get("entries", {})), diff --git a/yeto/rl/adapters/miles/entry.py b/yeto/rl/adapters/miles/entry.py index 76455ef1a..a1e52f142 100644 --- a/yeto/rl/adapters/miles/entry.py +++ b/yeto/rl/adapters/miles/entry.py @@ -672,14 +672,20 @@ def _is_full_parameter_sync(sync: Any) -> bool: return isinstance(sync, FullParameterSync) -def _policy_state_for(sync: Any, **lora: Any) -> Any: - """The full-parameter actor cut for a FullParameterSync island, else LoRA state.""" +def _policy_state_for(sync: Any, *, full_parameter: bool = False, **lora: Any) -> Any: + """The full-parameter actor cut for a FullParameterSync island; a trainer-side + full-parameter digest for a no-sync full-parameter island (L4.6); else LoRA state.""" from .full_parameter.session import MilesFullParameterPolicyState - from .state import MilesPolicyState + from .state import MilesFullParameterLocalPolicyState, MilesPolicyState if _is_full_parameter_sync(sync): return MilesFullParameterPolicyState(sync) + if full_parameter: + return MilesFullParameterLocalPolicyState( + actor_model=lora["actor_model"], base_model_revision=lora["base_model_revision"], + config_hash=lora["config_hash"], expected_layout_hash=lora["expected_layout_hash"], + runner=lora.get("runner")) return MilesPolicyState(**lora) @@ -745,8 +751,11 @@ def expected_policy() -> tuple[int, str]: raise DriverError("rollout requested before any complete publication") return parse_policy_snapshot_token(driver.expected_token) + full_parameter = _is_full_parameter_sync(sync) or getattr( + miles_args, "yeto_rl_parameter_mode", None) == "full" policy_state = _policy_state_for( sync, + full_parameter=full_parameter, actor_model=actor_model, base_model_revision=base_model_revision, config_hash=lora_config_hash, @@ -785,7 +794,7 @@ def expected_policy() -> tuple[int, str]: critic_model=critic_model, release_refs=release_refs, runner=runner, - install_actor_recorders=_is_full_parameter_sync(sync), + install_actor_recorders=full_parameter, ), policy_state=policy_state, publisher=MilesPublisher( @@ -876,6 +885,10 @@ def expected_policy() -> tuple[int, str]: _wire_round_cuts(driver, elastic=resume, miles_args=miles_args, algorithm=algorithm, base_model_revision=base_model_revision) _wire_droppable_reclaim(sync, driver) + from .sampler_island import wire_policy_releases + + # L4.6: a trainer island (--rl-sample-role trainer) releases every publish to the samplers + wire_policy_releases(driver, miles_args=miles_args, actor_model=actor_model, runner=runner) # rl-spot-cost-saving 4.1 (opt-in, YETO_SPOT_INFLIGHT_SAVE_DIR): timed in-flight save from .inflight_save import install as _install_inflight_save @@ -1022,13 +1035,42 @@ def _wire_round_cuts(driver, *, elastic, miles_args, algorithm, base_model_revis else {"ref_load": str(ref_load), "base_model_revision": base_model_revision}), ) config = getattr(miles_args, "yeto_rl_resume", None) or {} - wire_round_cuts(driver, controller=controller, source=source, - every=int(getattr(elastic, "every", None) or config.get("every") - or getattr(miles_args, "yeto_rl_resume_every", None) or 1), - fingerprint=run_fingerprint_of(miles_args, algorithm), - allow_config_change=bool(getattr(elastic, "allow_config_change", False) - or config.get("allow_config_change", False)), - next_lr=next_lr_for(miles_args)) + checkpoint = wire_round_cuts( + driver, controller=controller, source=source, + every=int(getattr(elastic, "every", None) or config.get("every") + or getattr(miles_args, "yeto_rl_resume_every", None) or 1), + fingerprint=run_fingerprint_of(miles_args, algorithm), + allow_config_change=bool(getattr(elastic, "allow_config_change", False) + or config.get("allow_config_change", False)), + next_lr=next_lr_for(miles_args)) + _wire_sample_buffer(checkpoint, driver=driver, miles_args=miles_args) + + +def _wire_sample_buffer(checkpoint: Any, *, driver: Any, miles_args: Any) -> None: + """L4.5: a trainer whose rollout reads the durable sample buffer writes the + consumed watermark into each round cut and continues there after a restore.""" + from .buffer_rollout import ROLLOUT_FUNCTION_PATHS, TRAINER, wire_sample_buffer + + if getattr(miles_args, "rollout_function_path", None) != ROLLOUT_FUNCTION_PATHS[TRAINER]: + return + if checkpoint is None: + raise RuntimeError("a trainer that reads the sample buffer needs round cuts " + "(a checkpoint store and no outer syncer) to keep its watermark") + wire_sample_buffer(checkpoint, driver.rollout) + + +def run_commanded_trainer(driver: Any) -> Any: + """One incarnation of a trainer the head's scheduler starts and stops + (``trainer_remote``): commands come through the round-cut store.""" + from yeto.rl.engine.trainer_remote import run_commanded_island + + controller = getattr(getattr(driver.sync, "round_cuts", None), "controller", None) + store = getattr(controller, "store", None) + if store is None: + raise RuntimeError("a trainer started by the head's scheduler needs round cuts " + "(a checkpoint store and no outer syncer)") + run_commanded_island(driver, store=store) + return driver.published_state def _wire_trainer_rebuild(driver, *, elastic, miles_args, algorithm, actor_model, @@ -1296,6 +1338,13 @@ def connect_island_ray(*, environ=None, ray_module=None, miles_args=None) -> str for key in (OVERRIDE_ENV, NEGATIVE_RUN_ENV): # launch-preflight-guards: negative-test island if environ.get(key): # (identity_test_salt is read where the sync client is built) env_vars[key] = environ[key] + from yeto.rl.engine.resume import STORE_IMPL_ENV + + from .buffer_rollout import ISLAND_ENVS as SAMPLE_BUFFER_ENVS + + for key in (*SAMPLE_BUFFER_ENVS, STORE_IMPL_ENV): # L4.6: the buffer rollout functions + if environ.get(key): # open the store inside the rollout executor actor + env_vars[key] = environ[key] if environ.get(ELASTIC_METADATA_ENV) == "1": # --rl-elastic: the rollout metadata hook runs inside Ray workers, which # inherit the raylet's environment, not the driver's. @@ -1828,6 +1877,34 @@ def manifest_pool_gpus(resources: Any) -> tuple[str, ...]: return tuple(gpus) +TRAINER_ACTOR_INIT_SUFFIX = "TrainRayActor.__init__" + + +def trainer_actor_import_done_unix(list_tasks: Callable[..., Any] | None = None) -> float | None: + """When the last trainer actor finished its creation task (Unix seconds). + + Ray runs the actor class import and ``__init__`` in the actor creation + task; Miles' trainer ``__init__`` is light (the model is built in + ``init``), so its end is the import end. Read from the Ray state API + (dashboard on, see the island run script); ``None`` when Ray has no such + record. Unverified on GPU (L2.6 checks it). + """ + + try: + if list_tasks is None: + from ray.util.state import list_tasks + tasks = list_tasks(filters=[("type", "=", "ACTOR_CREATION_TASK")], detail=True, + limit=10000, raise_on_missing_output=False) + except Exception: # noqa: BLE001 - a missing record is reported as a missing phase + return None + ends = [ + getattr(task, "end_time_ms", None) for task in tasks + if str(getattr(task, "name", "")).endswith(TRAINER_ACTOR_INIT_SUFFIX) + ] + ends = [end for end in ends if end] + return max(ends) / 1000.0 if ends else None + + def run_ports_island( miles_args: Any, launch: Any, @@ -1884,6 +1961,13 @@ def run_ports_island( heartbeat_interval_s=getattr(miles_args, "yeto_rl_heartbeat_interval_s", None), resource_interval_s=getattr(miles_args, "yeto_rl_resource_sample_interval_s", None), ).__enter__() # stopped before driver.run() and in the finally below + from yeto.rl.engine.start_phases import StartPhases + + # rl-trainer-cold-start-cost L0: submitted / container_ready / setup_done / + # ray_ready come from the launcher env and the setup shell's mark files. + start_phases = StartPhases( + lambda event, **fields: append_tape_event(miles_args, {"event": event, **fields})) + start_phases.replay() connect_island_ray(miles_args=miles_args) topology = getattr(launch.placement, "topology", None) if topology is not None and topology.nodes > 1: @@ -1926,6 +2010,13 @@ async def init(): startup.step("engine_ready") actor, critic = await create_training_models(miles_args, executor) startup.step("weights_loaded") + imported = trainer_actor_import_done_unix() + if imported is None: + start_phases.missing("actor_import_done", "Ray state API has no trainer actor creation time") + else: + start_phases.mark("actor_import_done", imported) + start_phases.mark("actor_init_done") + start_phases.forward_setup_events() # e.g. a late model_prefetch_failed # rl-algo-critic-family 3.1: Miles creates a critic iff use_critic # (estimator ppo); it must agree with the AlgorithmSpec. if (critic is not None) != bool(algorithm.execution.needs_critic): @@ -1998,6 +2089,9 @@ async def evaluate(rollout_id: int) -> dict[str, float]: miles_args, "yeto_rl_resource_sample_interval_s", None) eval_wiring.attach(driver, miles_args, eval_guard) # rl-eval-difficulty-buckets 2.2/5.2 driver.island_tasks = getattr(miles_args, "yeto_rl_island_tasks", None) # P3 + driver.code_pack_digest = getattr(miles_args, "yeto_rl_code_pack_digest", None) + driver.code_pack_mode = getattr(miles_args, "yeto_rl_code_pack_mode", None) + driver.start_phases = start_phases # L0: first_step_done if e2_plan is not None: # TEST ONLY: E2 GPU harness instead of the training loop from .e2_harness import HarnessContext, run_harness @@ -2008,6 +2102,10 @@ async def evaluate(rollout_id: int) -> dict[str, float]: )) return driver.published_state startup.__exit__(None, None, None) # the driver's own heartbeat starts in run() + from yeto.rl.engine import trainer_remote + + if trainer_remote.commanded(): # rl-trainer-start-stop L5.3c: the head starts/stops us + return run_commanded_trainer(driver) return driver.run() except BaseException as exc: error = exc diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 2532d2e1f..f74bd3339 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -37,6 +37,16 @@ def build_parser() -> argparse.ArgumentParser: return parser +def verify_island_code_pack(digest: str | None, tree: str | Path | None = None) -> None: + """rl-identity-tiering L9.5: refuse to start when the workdir is not the run pack.""" + + if not digest: + return + from yeto.codepack import verify_tree + + verify_tree(tree if tree is not None else Path(__file__).resolve().parents[4], digest) + + def parse_args(argv=None): parser = build_parser() args = parser.parse_args(argv) @@ -50,10 +60,12 @@ def parse_args(argv=None): SAO_STREAMING_SYNC_PRESET, ) - if args.parameter_mode == "full" and args.sync_preset not in FULL_PARAMETER_SYNC_PRESETS: + if (args.parameter_mode == "full" and args.sync_preset not in FULL_PARAMETER_SYNC_PRESETS + and not _buffer_island(args)): parser.error( "--parameter-mode full requires a full-parameter --sync-preset " - f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))})" + f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))}) or a no-sync " + "--rl-sample-role sampler / trainer island" ) if args.sync_preset == FULL_PARAMETER_SYNC_PRESET and not _trains_base_parameters(args): parser.error( @@ -218,6 +230,23 @@ def _elastic_cells(value: str | None) -> tuple[str, ...]: return tuple(c.strip() for c in (value or "").split(",") if c.strip()) +def apply_sample_role(args, miles_args, environ=None) -> None: + """L4.6 (--rl-sample-role sampler|trainer): the buffer config on ``miles_args`` + and the environment of the buffer rollout functions (forwarded to the Ray + workers by ``connect_island_ray``). Role ``both`` sets nothing.""" + from yeto.rl.engine import sample_roles + + from .buffer_rollout import island_env + + if sample_roles.sample_role(args) == sample_roles.BOTH: + return + config = sample_roles.resolve(args, groups_per_round=args.groups_per_round) + miles_args.yeto_rl_sample_buffer = config + (os.environ if environ is None else environ).update(island_env( + config.role, store_root=config.store, sampler_id=f"island{int(args.learner_id)}", + max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), wait_s=config.wait_s)) + + def apply_ports_infra_switches(args, miles_args, environ=None) -> None: """Carry the opt-in 2.3/3.x switches onto ``miles_args`` before any rollout process starts. Without them nothing is set (default path unchanged).""" @@ -238,6 +267,7 @@ def apply_ports_infra_switches(args, miles_args, environ=None) -> None: custom_agent=getattr(miles_args, "custom_agent_function_path", None), recompute_prefill=bool(getattr(miles_args, "recompute_logprobs_via_prefill", False))) miles_args.yeto_rl_max_policy_age = int(args.rl_max_policy_age) + apply_sample_role(args, miles_args, environ) observe = bool(getattr(args, "rl_observe_timeline", False)) if observe: miles_args.yeto_rl_observe_timeline = True @@ -397,14 +427,20 @@ class AlgorithmMismatchError(RuntimeError): def verify_ports_algorithm(args, miles_args, launch) -> None: """D9/D11: before the island joins outer sync (no bridge exists yet). - Plugins are re-hashed and imported; the island's algorithm hash is + Plugins are re-hashed and imported (a stale source hash: warning event, + rl-identity-tiering L9.2); the island's algorithm hash is compared with ``--rl-expected-algorithm-sha256`` (mismatch: event + refusal; missing: warning event, for manually started islands); absorbed flags and unverified allowances are recorded. """ algorithm = launch.algorithm - algorithm.verify_plugins() + for warning in algorithm.verify_plugins(): + _append_ports_event(args, miles_args, { + "event": "rl_plugin_source_mismatch", + "level": "warning", + "message": warning, + }) from yeto.rl.engine.algorithm import island_problems problems = island_problems(algorithm, { @@ -444,6 +480,7 @@ def verify_ports_algorithm(args, miles_args, launch) -> None: f"island {args.learner_id} algorithm hash {actual} differs from the launcher's " f"expected {expected.lower()}; refusing to join outer sync" ) + _append_correction_record(args, miles_args, algorithm) miles_args.yeto_rl_algorithm_absorbed_flags = dict(launch.absorbed_flags) miles_args.yeto_rl_outer_sync = not getattr(args, "rl_single_island_no_sync", False) # INFRA compares it with the runtime AlgorithmSpec before connect_island_ray @@ -456,6 +493,25 @@ def verify_ports_algorithm(args, miles_args, launch) -> None: ) +def _append_correction_record(args, miles_args, algorithm) -> None: + """rl-offpolicy-correction-passthrough D8: the island tape records the full + correction group, the Miles values and ``prox_anchor``; each field Miles + cannot map gets a ``correction_config_warning`` (phase 1: no refusal).""" + from yeto.rl.engine.offpolicy_correction import CORRECTION_RECORD_EVENT, correction_record + + from .offpolicy_correction import BACKEND, correction_mapping + + if algorithm.correction.is_default(): + return # default tapes unchanged; the run manifest still records the group + mapping = correction_mapping(algorithm) + _append_ports_event(args, miles_args, { + "event": CORRECTION_RECORD_EVENT, + **correction_record(algorithm.correction, mapping, backend=BACKEND), + }) + for warning in mapping.warnings: + _append_ports_event(args, miles_args, warning.event()) + + def _require_ports_supported(args, extra_argv: Sequence[str] = ()) -> None: """Reject combinations outside the R0 ports matrix before any startup.""" @@ -598,6 +654,23 @@ def _is_sha256(value: Any) -> bool: return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None +def _full_parameter_identity(model_config_sha256: str) -> tuple[str, str]: + """``(layout hash, config hash)`` of a full-parameter island: no LoRA layout, + the model config is the identity.""" + + return model_config_sha256, _canonical_sha256( + {"schema": "yeto-full-parameter-mode-v1", "model_config_sha256": model_config_sha256} + ) + + +def _buffer_island(args) -> bool: + """A no-sync sampler / trainer island (L4.6): it has no outer sync, so full + parameter mode runs without the full-parameter sync preset.""" + from yeto.rl.engine.sample_roles import BOTH, sample_role + + return bool(getattr(args, "rl_single_island_no_sync", False)) and sample_role(args) != BOTH + + def _full_parameter_mode(args, sync_preset: str) -> bool: """True for a full-parameter island. @@ -610,8 +683,9 @@ def _full_parameter_mode(args, sync_preset: str) -> bool: full_mode = getattr(args, "parameter_mode", "lora") == "full" full_preset = sync_preset in FULL_PARAMETER_SYNC_PRESETS - if full_mode and not full_preset: - raise ValueError("full parameter mode needs a full-parameter sync preset") + if full_mode and not full_preset and not _buffer_island(args): + raise ValueError("full parameter mode needs a full-parameter sync preset " + "(or a no-sync --rl-sample-role sampler / trainer island)") if full_preset and not _trains_base_parameters(args): raise ValueError(f"the {sync_preset} sync needs full parameter mode or expert-full tuning") if full_preset and not full_mode and sync_preset != FULL_PARAMETER_SYNC_PRESET: @@ -1623,7 +1697,7 @@ def run_miles( # full-parameter strict sweep also carries DeepSeek V4 expert-full (P0 1.7b.2). full_parameter_sync = sync_preset in FULL_PARAMETER_SYNC_PRESETS dense_full = sync_preset == FULL_PARAMETER_SYNC_PRESET - if _trains_base_parameters(args) and not full_parameter_sync: + if _trains_base_parameters(args) and not full_parameter_sync and not _buffer_island(args): raise ValueError(f"expert-full tuning needs the {FULL_PARAMETER_SYNC_PRESET} sync") sao_runtime = None if dense_full: @@ -1785,13 +1859,7 @@ def run_miles( args, miles_argv, model_config_sha256=model_config_hash ) if full_parameter: - layout_hash = model_config_hash - lora_config_hash = _canonical_sha256( - { - "schema": "yeto-full-parameter-mode-v1", - "model_config_sha256": model_config_hash, - } - ) + layout_hash, lora_config_hash = _full_parameter_identity(model_config_hash) else: from yeto.rl.core import canonical_layout_hash @@ -1973,6 +2041,26 @@ def run_miles( _configure_grad_audit(args, miles_args) require_ports_router_mode(miles_args) miles_args.yeto_rl_island_tasks = dict(island_tasks) if island_tasks else None + if full_parameter and not yeto_policy_sync: + # A no-sync buffer island (L4.6) trains every parameter without the sync + # block above: bind the same full-parameter identity for _run_ports. + from yeto.provenance import file_sha256 + + model_config_path = Path(model_path) / "config.json" + if not model_config_path.is_file(): + raise ValueError("full-parameter model has no config.json") + miles_args.yeto_rl_parameter_mode = parameter_mode + miles_args.yeto_rl_layout_hash, miles_args.yeto_rl_lora_config_hash = ( + _full_parameter_identity(file_sha256(model_config_path))) + # rl-identity-tiering: the run code pack digest (L9.4, checked at startup), else + # (L9.1 transition) the attested source tree hash. + if getattr(args, "code_pack_digest", None): + from yeto.codepack import CODE_PACK_MODE_GIT_PATCH + + miles_args.yeto_rl_code_pack_digest = args.code_pack_digest + miles_args.yeto_rl_code_pack_mode = CODE_PACK_MODE_GIT_PATCH + else: + miles_args.yeto_rl_code_pack_digest = args.source_sha256 _run_ports( args, miles_args, @@ -2167,6 +2255,13 @@ def _run_ports( apply_ports_infra_switches(args, miles_args) if _run_test_debug_rollout_replay(args, miles_args): return + from yeto.rl.engine.sample_roles import SAMPLER, sample_role + + if sample_role(args) == SAMPLER: # L4.6: generate into the sample buffer, never train + from yeto.rl.adapters.miles.sampler_island import run_sampler_island + + run_sampler_island(miles_args) + return run_ports_island( miles_args, launch, @@ -2190,6 +2285,7 @@ def main(argv=None) -> None: ) verify_source_tree_sha256(args.source_sha256) + verify_island_code_pack(getattr(args, "code_pack_digest", None)) if not is_immutable_commit(args.model_revision): raise ValueError("RL model revision must be an immutable commit") if not is_local_reference(args.data) and not is_immutable_commit( diff --git a/yeto/rl/adapters/miles/offpolicy_correction.py b/yeto/rl/adapters/miles/offpolicy_correction.py new file mode 100644 index 000000000..93e7b440b --- /dev/null +++ b/yeto/rl/adapters/miles/offpolicy_correction.py @@ -0,0 +1,173 @@ +"""Miles mapping of the neutral correction group +(change ``rl-offpolicy-correction-passthrough`` L6.3, L6.7, L6.9). + +:func:`correction_mapping` turns ``AlgorithmSpec.correction`` into Miles +argv (design D2/D3); :func:`algorithm_flags.algorithm_argv` emits the same +argv through its ``--use-tis`` row. Mapping table (design D3): + +==================== ========================================================== +neutral Miles +==================== ========================================================== +tis, token ``--use-tis --tis-clip upper --tis-clip-low lower`` +ratio_mask=icepop the same plus ``--custom-tis-function-path`` IcePop +ratio_mask=mis the same with the MIS function; mask mode, ``level``, + bounds, ``batch_normalize`` and ``veto`` as namespace + attributes (``mismatch_correction.neutral_mis_config``) +opsm ``--use-opsm --opsm-delta`` +==================== ========================================================== + +Rows Miles cannot express give a :class:`CorrectionWarning` (phase 1 runs +on): ``level`` other than token without the MIS mask, and ``veto`` / +``batch_normalize`` without the MIS mask. + +Decoupled PPO (design D2) needs pi_prox, which Miles recomputes on the +weights before each update (``prox_anchor=per_step``). :data:`FORBIDDEN` +lists the flags that remove it or mask by version; :func:`refuse_forbidden_argv` +refuses them with the flag named. + +Admission (design D5, approximate): :func:`admission_argv` maps the +consumption limit ``A_hard`` and an optional admission limit ``k`` to +``--max-weight-staleness`` and the fully_async buffer capacity. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from yeto.rl.algos.mismatch_correction import ICEPOP_PATH, MIS_PATH, mis_config +from yeto.rl.engine.algorithm import AlgorithmSpec +from yeto.rl.engine.offpolicy_correction import ( + PROX_PER_STEP, + CorrectionMapping, + CorrectionWarning, + ForbiddenSetting, + refuse_forbidden, +) +from yeto.rl.engine.policy_age import validate_limit + +BACKEND = "miles" +PROX_ANCHOR = PROX_PER_STEP + +FORBIDDEN = ( + ForbiddenSetting(BACKEND, "--use-rollout-logprobs", + "it makes the rollout log-probabilities the PPO clip anchor, so pi_prox " + "is lost"), + ForbiddenSetting(BACKEND, "--skip-actor-forward-only", + "it skips the actor forward pass that computes pi_prox"), + ForbiddenSetting(BACKEND, "--mask-offpolicy-in-partial-rollout", + "it masks the older-version tokens of a continued sample; every response " + "token must enter the loss"), +) +FORBIDDEN_FLAGS = tuple(f.setting for f in FORBIDDEN) + +# Miles fully_async data buffer: capacity = factor * rollout_batch_size +# (miles/rollout/fully_async_data_buffer.py:136, Miles fork read 2026-10-14). +BUFFER_CAPACITY_FLAG = "--async-data-buffer-capacity-factor" + + +def _num(value: float) -> str: + return repr(float(value)) + + +def _function_path(spec: AlgorithmSpec) -> str | None: + c = spec.correction + if c.function is not None: + return c.function.path + return {"icepop": ICEPOP_PATH, "mis": MIS_PATH}.get(c.ratio_mask) + + +def tis_argv(spec: AlgorithmSpec) -> list[str]: + """``--use-tis`` and its bounds / function for a tis or custom correction.""" + + c = spec.correction + if c.method not in ("tis", "custom"): + return [] + argv = ["--use-tis"] + if c.upper is not None: + argv += ["--tis-clip", _num(c.upper)] + if c.lower is not None: + argv += ["--tis-clip-low", _num(c.lower)] + function = _function_path(spec) + if function is not None: + argv += ["--custom-tis-function-path", function] + return argv + + +def _warnings(spec: AlgorithmSpec) -> list[CorrectionWarning]: + c = spec.correction + if c.method != "tis" or c.ratio_mask == "mis": + return [] + out = [] + if c.level != "token": + out.append(CorrectionWarning( + "level", BACKEND, f"Miles TIS{' / IcePop' if c.ratio_mask == 'icepop' else ''} " + "weights tokens only; level other than token needs ratio_mask='mis'", "token")) + if c.veto is not None: + out.append(CorrectionWarning( + "veto", BACKEND, "Miles has a veto only in the MIS function (ratio_mask='mis')", None)) + if c.batch_normalize: + out.append(CorrectionWarning( + "batch_normalize", BACKEND, + "Miles normalizes weights only in the MIS function (ratio_mask='mis')", False)) + return out + + +def _values(argv: Sequence[str]) -> dict[str, Any]: + values: dict[str, Any] = {} + i = 0 + while i < len(argv): + flag = argv[i] + if i + 1 < len(argv) and not argv[i + 1].startswith("--"): + values[flag] = argv[i + 1] + i += 2 + else: + values[flag] = True + i += 1 + return values + + +def correction_mapping(spec: AlgorithmSpec) -> CorrectionMapping: + """The Miles settings of ``spec.correction`` with warnings and ``prox_anchor``.""" + + c = spec.correction + argv = tis_argv(spec) + if c.opsm_delta is not None: + argv += ["--use-opsm", "--opsm-delta", _num(c.opsm_delta)] + if c.use_rollout_logprobs: + argv.append("--use-rollout-logprobs") # refused by refuse_forbidden_argv + if c.mismatch_metrics: + argv.append("--get-mismatch-metrics") + attrs = mis_config(spec) + values = {**_values(argv), **{f"namespace.{k}": v for k, v in sorted(attrs.items())}} + return CorrectionMapping(backend=BACKEND, prox_anchor=PROX_ANCHOR, values=values, + warnings=tuple(_warnings(spec)), argv=tuple(argv), + runtime_attrs=attrs) + + +def refuse_forbidden_argv(argv: Sequence[str]) -> None: + """Refuse a Miles argv that holds a :data:`FORBIDDEN` flag (names the flag).""" + + refuse_forbidden((token.split("=", 1)[0] for token in argv), FORBIDDEN) + + +def admission_argv(max_policy_age: int, *, admission_limit: int | None = None) -> tuple[str, ...]: + """Approximate admission and consumption limits on Miles fully_async (D5). + + Consumption: ``--max-weight-staleness A_hard`` drops a group whose oldest + token is more than ``A_hard`` versions behind (whole group). Admission: + Miles has no limit by start version; a buffer of ``k`` rounds of groups + (``--async-data-buffer-capacity-factor k``) bounds how far generation can + run ahead. ``k`` None keeps Miles' buffer default. ``k <= A_hard``. + """ + + limit = validate_limit(max_policy_age) + if limit == 0: + return () + argv = ["--max-weight-staleness", str(limit)] + if admission_limit is not None: + k = validate_limit(admission_limit) + if not 1 <= k <= limit: + raise ValueError(f"admission limit k={k} must be in [1, A_hard={limit}]") + argv += [BUFFER_CAPACITY_FLAG, str(k)] + return tuple(argv) diff --git a/yeto/rl/adapters/miles/pins.py b/yeto/rl/adapters/miles/pins.py index 6a2d26dc8..2bc5533b6 100644 --- a/yeto/rl/adapters/miles/pins.py +++ b/yeto/rl/adapters/miles/pins.py @@ -49,7 +49,11 @@ class MilesRevisionPins(NamedTuple): # transport (95533a0cd, 5f7ac4e32, 587bd258a; P0-1.7b.2) + SAO streaming # primitives ported from agentenv/miles d5eb5e82a: content identity, per-role # chunks, critic group RPC (P0-1.7b.3)). -MILES_NEXT_COMMIT = "386ae3b07138ccd45ce4a07fc0383587f7968a3a" +# s24-miles-import2 (44d41efd0 = 17aeb5cb5 + the megatron_utils package +# import stays light, so RayWorkerManager no longer loads megatron.bridge +# (rl-trainer-cold-start-cost L2.2)). 17aeb5cb5 = 386ae3b07 + light Ray actors +# stay off megatron_utils (L2.1) + the master port is held until use (S23 defect 9). +MILES_NEXT_COMMIT = "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9" # sgl-project/sglang ``sglang-miles`` head when radixark/miles@9e4260d was # committed (upstream's Dockerfile follows that branch unpinned). SGLANG_NEXT_REPOSITORY = "https://github.com/michaellchung/sglang" @@ -100,6 +104,10 @@ class MilesRevisionPins(NamedTuple): "587bd258a350cc4757e3e774436eb48085d808cc", # 386ae3b07 = + SAO streaming primitives (P0-1.7b.3); no loss path. "386ae3b07138ccd45ce4a07fc0383587f7968a3a", + # 17aeb5cb5 = + lazy imports for light actors, master port held (S24 L2.1); no loss path. + "17aeb5cb51635b9ba6ef4ca10db95762695c7e03", + # 44d41efd0 = + megatron_utils runtime patches moved out of the package init (S24 L2.2); no loss path. + "44d41efd0a0428cb12b4ee34422de6f5e1ccb9d9", ) @@ -138,10 +146,16 @@ def miles_fork_lineage_from(commit: str) -> frozenset[str]: # every SHA. PUBLIC on ghcr.io/michaellchung since 2026-10-07 (anonymous # pull); a private image needs SKYPILOT_DOCKER_{USERNAME,PASSWORD,SERVER} or # --rl-image-private (yeto.launcher.registry_login_for; read:packages token). -# Tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. +# Tag 44d41ef-2fa8801; (S24 L2.2: Miles s24-miles-import2 and the lazy-modelopt +# Megatron-LM patch in the layer; previous 17aeb5c-2fa8801 @sha256:b304f857...; +# build record openspec/changes/rl-infra-spec/evidence/ports-image/ +# 2026-10-14-44d41ef-2fa8801/). Earlier: tag 17aeb5c-2fa8801; (S24 L2.1-L2.4: Miles s24-miles-import, peft 0.20.0 and the +# lazy-modelopt Megatron-Bridge __init__ in the layer; previous 386ae3b-2fa8801 +# @sha256:bcd5ffd3...; build record openspec/changes/rl-infra-spec/evidence/ +# ports-image/2026-10-14-17aeb5c-2fa8801/). Earlier: tag 386ae3b-2fa8801; (S23 P0-1.11 GPU run: Miles s23-sao-streaming-roles; previous 64b591a-2fa8801 @sha256:fa2413be...; build record openspec/changes/rl-infra-spec/evidence/ports-image/2026-10-14-386ae3b-2fa8801/). Earlier: (S19 #1: critic family + algo-supplement metrics baked in; previous ddce209-2fa8801 @sha256:9c252c38...) (S18 ARU-3: Miles s18-agentic-suspend = 2f7871fb2 + --agentic-suspend-between-turns; previous 2f7871f-2fa8801 @sha256:62b4f164..., efbbc63-2fa8801 @sha256:a7990076..., 8bc5223-4e4148f @sha256:4aeafd77...) ghcr.io/michaellchung (public package); linux/amd64 only. MILES_NEXT_IMAGE = ( "docker:ghcr.io/michaellchung/yeto-miles-ports@sha256:" - "bcd5ffd3c51be496f3d10cc8c42e6010ecd902e1b449b0e42b1f71c2f40ab1ad" + "5b9dc2a9e6e315d73dabfbc6faae428b59cc01aabec6189645a5c8c80d6e3721" ) MILES_NEXT_IMAGE_MANIFEST = "/opt/yeto/image-manifest.json" # Nebius VM images whose /var/lib/docker already holds a docker image's diff --git a/yeto/rl/adapters/miles/policy_age.py b/yeto/rl/adapters/miles/policy_age.py index a93785417..d52886feb 100644 --- a/yeto/rl/adapters/miles/policy_age.py +++ b/yeto/rl/adapters/miles/policy_age.py @@ -33,7 +33,12 @@ from yeto.rl.engine.policy_age import STAGE_NAMES, PolicyAgeError, PolicyAgeSupport, validate_limit -SUPPORT = PolicyAgeSupport(backend="miles", stage=3, max_policy_age=1) +# rl-offpolicy-correction-passthrough L6.9: the cap was 1 (the only limit +# measured on GPU). Decoupled PPO needs A_hard >= 2. Miles itself has no upper +# bound (--max-weight-staleness and --agentic-suspend-max-rounds take any int); +# 4 bounds the rollout compute a dropped group can waste. Sub-agent decision, +# 2026-10-14; limits above 1 are not yet run on GPU (L6.10). +SUPPORT = PolicyAgeSupport(backend="miles", stage=3, max_policy_age=4) PARTIAL_ROLLOUT_FLAGS = ("--partial-rollout",) AGENTIC_GENERATE = "miles.rollout.generate_hub.agentic_tool_call.generate" diff --git a/yeto/rl/adapters/miles/publish.py b/yeto/rl/adapters/miles/publish.py index 30b422ea9..6832db65c 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -342,6 +342,14 @@ def _default_flatten() -> Callable[[Any], list[dict[str, Any]]]: return flatten_inference_engine_checksums +def engine_checksum_digest(body: dict[str, Any]) -> str: + """One engine's flattened weight checksum body as one sha256 (key-sorted JSON): + the value a convertible node compares after loading a published version.""" + return hashlib.sha256( + json.dumps(body, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") + ).hexdigest() + + async def _maybe_await(value: Any) -> Any: """The trainer hook may be sync (tests, fakes) or a coroutine (real Miles policy state, which must not block the running loop).""" @@ -373,6 +381,8 @@ def __init__( self._runner = runner or LoopRunner() self.last_engine_checksums: dict[str, dict[str, Any]] | None = None self._published_once = False + # (engines, ids, members) stamped by the last load_version (convertible node) + self._loaded_engines: tuple[list[Any], list[str], frozenset[str]] | None = None # (token, per-engine checksum body) of the last full publication whose # engines all reported the same body: the payload reference (3.5). self._reference: tuple[str, dict[str, Any]] | None = None @@ -469,24 +479,9 @@ def publish(self, state: TrainableState, *, token_rollout_id: int | None = None) ) return PublicationResult(manifest=manifest, members=frozenset(members)) - async def _publish( - self, version: int, token: str - ) -> tuple[frozenset[str], dict[str, Any] | None]: - update_weights = self._update_weights_to(version) - offload_rollout = bool(getattr(self._args, "offload_rollout", False)) - try: - # Upstream train.py: weights come back before every update except - # the first (engines start resident); KV comes back after it. - if offload_rollout and self._published_once: - await self._controller.onload_weights() - await update_weights(self._args, self._actor, self._executor, self._controller) - if offload_rollout: - await self._controller.onload_kv() - except Exception as exc: - raise PublicationError(f"upstream update_weights failed: {exc}", - cause=PublicationCause.UPDATE_FAILED) from exc - self._published_once = True - + async def _stamp_version(self, version: int, label: str) -> tuple[list[Any], list[str], frozenset[str]]: + """Under the controller's update lock, stamp ``version`` on every running + engine and read it back. Returns the engines, their ids and the members.""" info = await self._controller.start_update_weights() ok = False try: @@ -499,7 +494,7 @@ async def _publish( failed, cause = await _stamp_engine_version(engines, ids, version) if failed: raise PublicationError( - f"{len(failed)}/{len(engines)} rollout engines did not acknowledge {token}: " + f"{len(failed)}/{len(engines)} rollout engines did not acknowledge {label}: " f"{sorted(failed)}", frozenset(failed), cause=cause, engine_ids=failed, ) @@ -511,6 +506,26 @@ async def _publish( ) else: await self._controller.abort_update_weights() + return engines, ids, members + + async def _publish( + self, version: int, token: str + ) -> tuple[frozenset[str], dict[str, Any] | None]: + update_weights = self._update_weights_to(version) + offload_rollout = bool(getattr(self._args, "offload_rollout", False)) + try: + # Upstream train.py: weights come back before every update except + # the first (engines start resident); KV comes back after it. + if offload_rollout and self._published_once: + await self._controller.onload_weights() + await update_weights(self._args, self._actor, self._executor, self._controller) + if offload_rollout: + await self._controller.onload_kv() + except Exception as exc: + raise PublicationError(f"upstream update_weights failed: {exc}", + cause=PublicationCause.UPDATE_FAILED) from exc + self._published_once = True + engines, ids, members = await self._stamp_version(version, token) checksums = None if self._verify_checksums: @@ -528,6 +543,61 @@ async def _publish( return members, checksums + # -- convertible node (rl-trainer-start-stop L7.2) -------------------------- + @property + def published_checksum(self) -> str | None: + """Digest of the weight checksum every engine reported at the last full + publication (None: no publication yet, checksums off, or engines differed).""" + return None if self._reference is None else engine_checksum_digest(self._reference[1]) + + def load_version(self, version: int) -> None: + """Load the trainer's weights of published ``version`` into this node's + engines and stamp the version. + + For a convertible node going back to rollout: the caller has offloaded the + trainer (``update_weights`` reads its host backups) and given the engines + their weight memory back; KV onload and admission stay with the caller. + Read the result with :meth:`served_checksum` once the KV cache is back.""" + self._runner.run(self._load_version(version)) + + async def _load_version(self, version: int) -> None: + update_weights = self._update_weights_to(version) + try: + await update_weights(self._args, self._actor, self._executor, self._controller) + except Exception as exc: + raise PublicationError(f"upstream update_weights failed: {exc}", + cause=PublicationCause.UPDATE_FAILED) from exc + self._loaded_engines = await self._stamp_version(version, f"version {version}") + + def served_checksum(self) -> str: + """Digest of the weight checksum the engines of the last :meth:`load_version` + serve; all engines must report the same one. + + Needs the KV cache back: in LoRA mode the adapter is first loaded into + each engine's LoRA pool by one generate (as for a member publication), + so the checksum covers it (G-L7.4: a generate with the KV cache released + crashed the engine scheduler).""" + if self._loaded_engines is None: + raise PublicationError("served_checksum needs a load_version first") + return self._runner.run(self._served_checksum(*self._loaded_engines)) + + async def _served_checksum(self, engines: list[Any], ids: list[str], + members: frozenset[str]) -> str: + if lora_mode(self._args): + await self._warm_up_lora(engines, ids, [], sorted(members)) + raw = await self._controller.check_weights(action="checksum") + try: + bodies = (self._flatten or _default_flatten())(raw) + except AssertionError as exc: + raise PublicationError(f"engine weight checksum failed: {exc}", + cause=PublicationCause.PAYLOAD_MISMATCH) from exc + digests = {engine_checksum_digest(dict(b)) for b in bodies} + if len(digests) != 1: + raise PublicationError( + f"{len(bodies)} engines reported {len(digests)} different weight checksums", + cause=PublicationCause.PAYLOAD_MISMATCH) + return digests.pop() + # -- member publication (3.5) --------------------------------------------- def publish_members( self, diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index 9b49dd193..f5490f47f 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -30,6 +30,11 @@ from collections.abc import Mapping from types import SimpleNamespace +from yeto.rl.engine.offpolicy_correction import ( + GATE_DROPPED_LENGTHS, + GATE_GRANULARITY, + TRUNCATED_BY_TURN, +) from yeto.rl.engine.policy_versions import PublishedPolicies from yeto.rl.engine.ports import GroupMetadata, RolloutBatchHandle, over_sampling_fields from .traits import MILES_ABORT_MECHANISM @@ -361,6 +366,10 @@ def __init__(self, name: str = DEFAULT_SINK_ACTOR, *, timeout_s: float = 30.0) - def sink_spec(self) -> str: return f"ray:{self.name}" + def discard(self) -> int: + """Drop every pushed record (a sampler island has no consumer); returns the count.""" + return len(self._ray.get(self._actor.take_all.remote())) + def set_policy_token(self, token: str) -> None: self._ray.get(self._actor.set_token.remote(token)) @@ -460,6 +469,7 @@ def handle_from_metadata( aborted_in_flight_tokens=payload.get("aborted_in_flight_tokens"), aborted_in_flight_unknown_groups=payload.get("aborted_in_flight_unknown_groups"), abort_mechanism=MILES_ABORT_MECHANISM, + backend="miles", tool_wait_seconds=( float(payload["tool_wait_seconds"]) if "tool_wait_seconds" in payload else None ), @@ -507,6 +517,8 @@ def handle_from_metadata( "resumed_groups", "resumed_done_groups", "suspend_expired_trajectories", "policy_age_exceeded_trajectories", "suspended_trajectories", "suspended_env_seconds", "suspend_retried_turns", + # rl-offpolicy-correction-passthrough D5/D6: the driver pops these into ``offpolicy_gate`` + GATE_DROPPED_LENGTHS, GATE_GRANULARITY, TRUNCATED_BY_TURN, ) @@ -1025,6 +1037,11 @@ def _harness_admission(self, method: str, members: frozenset[str]) -> None: _resolve(_call(source, method, sorted(members))) + def set_admission(self, open_: bool) -> None: + """Convertible node (rl-trainer-start-stop L7.2): close (False) or open (True) + harness admission to this pool's running members (no-op without a board).""" + self._harness_admission("open_admission" if open_ else "close_admission", self.members()) + def _offload_after_rollout(self) -> None: """Upstream ``train.py`` serial colocated branch (--offload-rollout).""" diff --git a/yeto/rl/adapters/miles/round_cut.py b/yeto/rl/adapters/miles/round_cut.py index 488802221..0197eadb2 100644 --- a/yeto/rl/adapters/miles/round_cut.py +++ b/yeto/rl/adapters/miles/round_cut.py @@ -32,7 +32,7 @@ from typing import Any from yeto.rl.engine import recovery -from yeto.rl.engine.recovery import CUT_KEY, MANUAL, PERIODIC, ConsistentCut +from yeto.rl.engine.recovery import CUT_KEY, FINAL_TRIGGERS, MANUAL, PERIODIC, ConsistentCut from yeto.rl.engine.resume import ROUND_CUTS # under the checkpoint store (not under the state dir) ISLAND = "local" # round cuts run without an outer syncer: one island, the local one @@ -72,6 +72,13 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 self.pending_lr_check: dict[str, Any] | None = None # set by resume(): where the rollout data source continues (the cut's cursor) self.resume_data_cursor: dict[str, int] | None = None + # rl-trainer-start-stop D4: the durable sample buffer's consumed watermark, + # written into the cut record (None: the run reads no sample buffer). + # buffer_watermark(next_rollout_id) commits and returns the position after the + # last trained round; buffer_resume(watermark) continues there after a restore + # (buffer_rollout.wire_sample_buffer sets both). + self.buffer_watermark: Any = None # callable(int) -> int | None + self.buffer_resume: Any = None # callable(int) -> int # ---------------------------------------------------------------- pointer @property @@ -109,15 +116,21 @@ def due(self, rollout_id: int, *, final: bool = False) -> bool: return should_cut(int(rollout_id), every=self.every, final=final) - def save(self, driver: Any, *, rollout_id: int, final: bool = False) -> dict[str, Any] | None: + def save(self, driver: Any, *, rollout_id: int, final: bool = False, + trigger: str | None = None) -> dict[str, Any] | None: """Cut the trainer at the safe point before rollout ``rollout_id`` (= the published version), point the state dir at it, sync the store. Returns the pointer, or None when nothing was cut (not due, no publication yet, already cut at this rollout, - island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it).""" + island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it). + + ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for + a stop the trainer scheduler ordered); a periodic cut takes none.""" import time from yeto.rl.engine.resume import directory_bytes + if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): + raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") ctl = self.controller if ctl.recovery_required or driver.published_version is None: return None @@ -161,9 +174,11 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False) -> dict[str "final": bool(final), } cut = ConsistentCut( - checkpoint_id=cut_id, trigger=MANUAL if final else PERIODIC, + checkpoint_id=cut_id, trigger=(trigger or MANUAL) if final else PERIODIC, outer_step=info["outer_version"], island_steps={ISLAND: info["local_step"]}, - completed_groups=consumed_n, completed_groups_digest=consumed_sha) + completed_groups=consumed_n, completed_groups_digest=consumed_sha, + buffer_watermark=(self.buffer_watermark(int(rollout_id)) if callable(self.buffer_watermark) + else None)) saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) info[CUT_KEY] = saved[CUT_KEY] synced = saved["synced"] @@ -205,6 +220,8 @@ def resume(self, driver: Any) -> dict[str, Any] | None: info = self.pointer() if info is None: + if callable(self.buffer_resume): + self.buffer_resume(0) return None started = time.monotonic() rid = int(info["next_rollout_id"]) @@ -250,6 +267,12 @@ def resume(self, driver: Any) -> dict[str, Any] | None: if in_flight is not None: checks["in_flight"] = in_flight self.resume_data_cursor = _cut_data_cursor(manifest) + if callable(self.buffer_resume): + watermark = (info.get(CUT_KEY) or {}).get("buffer_watermark") + if watermark is None: + raise RoundCutError(f"round cut {info['cut_id']} has no sample buffer watermark " + "but this trainer reads the sample buffer") + checks["buffer_watermark"] = int(self.buffer_resume(int(watermark))) driver.local_step = int(info["local_step"]) self.resumed = dict(info) if info.get("lr_at_next_round") is not None: diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 3b301eca0..33e72de61 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -183,6 +183,16 @@ shared_flag("--sglang-page-size", type=int, default=None), shared_flag("--sglang-max-running-requests", type=int, default=None), shared_flag("--sglang-chunked-prefill-size", type=int, default=None), + shared_flag( + "--rl-sampler-snapshot", + action="store_true", + help=( + "ports, fixed partition: start the sampler (rollout) SGLang engines with " + "--enable-memory-saver --enable-weights-cpu-backup (GPU snapshot mode); " + "off by default. Refused unless sampler_snapshot_compat.json records a " + "pass for the pinned SGLang commit and the NCCL weight-update path" + ), + ), ) _AGENT = ( @@ -254,6 +264,9 @@ island_flag("--learner-generation", type=int, default=0), island_flag("--reward-sha256", required=True), island_flag("--source-sha256", required=True), + # rl-identity-tiering L9.5: the run code pack digest; the island refuses to start + # when its workdir is not that pack. + island_flag("--code-pack-digest", default=None), island_flag("--global-rounds", type=int, required=True), island_flag( "--parameter-mode", diff --git a/yeto/rl/adapters/miles/sampler_island.py b/yeto/rl/adapters/miles/sampler_island.py new file mode 100644 index 000000000..727de4200 --- /dev/null +++ b/yeto/rl/adapters/miles/sampler_island.py @@ -0,0 +1,220 @@ +"""Sampler and trainer islands over the sample buffer on Miles (L4.6). + +* :func:`run_sampler_island`: the sampler island's entry. Miles runs with + ``--debug-rollout-only`` (rollout components only: no trainer model, no + trainer GPU); the rollout function writes every group to the buffer + (``BufferWriterRolloutFn``); :class:`~yeto.rl.engine.sampler_loop.SamplerLoop` + runs the rounds and loads each policy release with :class:`EngineReleaseLoader`. +* :func:`wire_policy_releases`: on a trainer island, every successful publish + is also released into the store (:class:`ReleasingPublisher`): the trainer + exports an HF checkpoint (``actor.export_hf``) into the release directory. + +Engine weight load (the order of the publisher, ``publish.MilesPublisher``): +under the controller's update lock, every engine loads the checkpoint +(``update_weights_from_disk`` with the weight version), then the plain policy +version is stamped and read back on every engine; any failure aborts the lock. +The rollout executor then takes the version as its current one. + +This module imports Miles only inside :func:`run_sampler_island`. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from yeto.rl.engine.policy_release import PolicyFollower, PolicyRelease, PolicyReleaseStore +from yeto.rl.engine.policy_versions import engine_weight_version +from yeto.rl.engine.sample_roles import SAMPLER, TRAINER + +from .buffer_rollout import SampleBufferRolloutError, executor_read_cursor, open_buffer +from .publish import _engine_id, _maybe_await, _stamp_engine_version + +SAMPLER_TRAIN_PARALLEL_CONFIG = {"dp_size": 1} +BASE_POLICY_VERSION = 0 # the engines start with the base weights: policy version 0 + + +def _buffer_config(miles_args: Any, role: str) -> Any: + config = getattr(miles_args, "yeto_rl_sample_buffer", None) + return config if config is not None and config.role == role else None + + +def release_store(config: Any, environ: Any = None) -> PolicyReleaseStore: + from yeto.rl.engine.resume import store_for + + return PolicyReleaseStore(store_for(config.store, environ=environ)) + + +# ---------------------------------------------------------------------- sampler side +class EngineReleaseLoader: + """Loads a release into every rollout engine (``PolicyFollower`` load hook). + + ``run`` drives a coroutine on the island's event loop (``LoopRunner.run``).""" + + def __init__(self, controller: Any, executor: Any, run: Callable[[Any], Any]) -> None: + self.controller = controller + self.executor = executor + self.run = run + + def __call__(self, release: PolicyRelease, directory: Path) -> None: + self.run(self._load(str(directory), int(release.version))) + + def stamp(self, version: int) -> None: + """Stamp ``version`` on the engines without loading weights (the base policy).""" + self.run(self._load(None, int(version))) + + async def _load(self, path: str | None, version: int) -> None: + info = await self.controller.start_update_weights() + ok = False + try: + engines = list(info.rollout_engines) + if not engines: + raise SampleBufferRolloutError("no rollout engines to load the policy into") + if path is not None: + stamp = engine_weight_version(version) + await asyncio.gather(*(e.update_weights_from_disk(model_path=path, weight_version=stamp) + for e in engines)) + ids = [_engine_id(e, i) for i, e in enumerate(engines)] + failed, cause = await _stamp_engine_version(engines, ids, version) + if failed: + raise SampleBufferRolloutError( + f"{len(failed)}/{len(engines)} rollout engines did not take policy " + f"v{version} ({cause.value}): {sorted(failed)}") + ok = True + finally: + if ok: + await self.controller.end_update_weights( + snapshot_cell_id_to_hashes=info.snapshot_cell_id_to_hashes) + else: + await self.controller.abort_update_weights() + await _maybe_await(self.executor.set_weight_version(int(version))) + + +def run_sampler_island(miles_args: Any, *, environ: Any = None) -> int: + """Run the sampler island to its data end; returns the rounds generated.""" + from miles.ray.placement_group import create_rollout_components + from miles.utils.async_utils import Disposer + from miles.utils.orchestration_utils import init_orchestration_script + + from yeto.rl.engine.events import append_tape_event + from yeto.rl.engine.sampler_loop import SamplerLoop + + from . import LoopRunner + from .entry import connect_island_ray + from .rollout import RayMetadataSink, _release + + config = _buffer_config(miles_args, SAMPLER) + if config is None: + raise SampleBufferRolloutError("run_sampler_island needs --rl-sample-role sampler") + if not getattr(miles_args, "debug_rollout_only", False): + raise SampleBufferRolloutError("a sampler island runs Miles with --debug-rollout-only") + + def emit(event: str, **fields: Any) -> None: + append_tape_event(miles_args, {"event": event, **fields}) + + connect_island_ray(miles_args=miles_args) + # The rollout process's metadata hook pushes each round's group records to this + # named actor (on a trainer island the driver creates it and consumes them). + metadata = RayMetadataSink() + runner = LoopRunner() + disposer = Disposer() + + async def init() -> tuple[Any, Any]: + await disposer.__aenter__() + init_orchestration_script(miles_args, disposer=disposer) + controller, executor, _ = await create_rollout_components(miles_args) + disposer.add(controller, executor) + # Miles post-processes each round for the trainer's data-parallel layout, which a + # trainer actor sets; a sampler has none. The buffer trainer is DP=1 (full mode). + await executor.set_train_parallel_config(SAMPLER_TRAIN_PARALLEL_CONFIG) + return controller, executor + + async def generate(rollout_id: int) -> None: + await controller.prepare_rollout(rollout_id) + _release(await executor.get(rollout_id)) # the rollout function wrote the groups + metadata.discard() # the groups are in the buffer; nobody here reads the records + + error: BaseException | None = None + try: + controller, executor = runner.run(init()) + loader = EngineReleaseLoader(controller, executor, runner.run) + loader.stamp(BASE_POLICY_VERSION) # until the first release is loaded + follower = PolicyFollower(release_store(config, environ), loader, + version=BASE_POLICY_VERSION) + loop = SamplerLoop( + follower=follower, buffer=open_buffer(environ), emit=emit, + generate=lambda rollout_id: runner.run(generate(rollout_id)), + num_rollouts=int(miles_args.num_rollout), + max_backlog_groups=int(config.max_backlog_groups), + min_round_s=config.min_round_s) + emit("rl_sampler_start", store=config.store, num_rollouts=int(miles_args.num_rollout), + max_backlog_groups=int(config.max_backlog_groups)) + return loop.run() + except BaseException as exc: + error = exc + raise + finally: + try: + runner.run(disposer.__aexit__(None if error is None else type(error), error, + None if error is None else error.__traceback__)) + finally: + runner.close() + + +# ---------------------------------------------------------------------- trainer side +class ReleasingPublisher: + """A publisher that also releases each published policy for the samplers. + + ``export(version, directory)`` writes the weights; ``read_position()`` is the + trainer's buffer read position (None before its reader exists). Every other + attribute is the wrapped publisher's.""" + + def __init__(self, inner: Any, releases: PolicyReleaseStore, + export: Callable[[int, Path], None], + read_position: Callable[[], int | None], emit: Callable[..., None]) -> None: + self.inner = inner + self.releases = releases + self.export = export + self.read_position = read_position + self.emit = emit + + def publish(self, state: Any, **kwargs: Any) -> Any: + result = self.inner.publish(state, **kwargs) + version = int(state.policy_version) + position = self.read_position() + if position is None: # nothing read yet in this process: keep the last known position + latest = self.releases.latest() + position = None if latest is None else latest.buffer_selected + release = self.releases.release(version, state.policy_tensor_hash(), + lambda directory: self.export(version, directory), + buffer_selected=position) + if release is not None: + self.emit("rl_policy_released", seq=release.seq, policy_version=version, + policy_hash=release.policy_hash, buffer_selected=release.buffer_selected) + return result + + def __getattr__(self, name: str) -> Any: + return getattr(self.inner, name) + + +def wire_policy_releases(driver: Any, *, miles_args: Any, actor_model: Any, runner: Any, + environ: Any = None) -> ReleasingPublisher | None: + """A trainer island (``--rl-sample-role trainer``) releases every publish.""" + config = _buffer_config(miles_args, TRAINER) + if config is None: + return None + in_executor = getattr(driver.rollout, "_in_executor", None) + if not callable(in_executor): + raise SampleBufferRolloutError(f"rollout pool {type(driver.rollout).__name__} cannot " + "run code in the rollout executor") + + def export(version: int, directory: Path) -> None: + runner.run(actor_model.export_hf(rollout_id=int(version), path=str(directory))) + + publisher = ReleasingPublisher(driver.publisher, release_store(config, environ), export, + read_position=lambda: in_executor(executor_read_cursor), + emit=driver.emit) + driver.publisher = publisher + return publisher diff --git a/yeto/rl/adapters/miles/sampler_snapshot.py b/yeto/rl/adapters/miles/sampler_snapshot.py new file mode 100644 index 000000000..48203c455 --- /dev/null +++ b/yeto/rl/adapters/miles/sampler_snapshot.py @@ -0,0 +1,152 @@ +"""Opt-in SGLang snapshot mode for sampler engines (rl-trainer-cold-start-cost L8.2, design D5). + +Snapshot mode starts the rollout (sampler) SGLang engines with +``enable_memory_saver`` and ``enable_weights_cpu_backup``, the two flags a GPU +memory snapshot needs. The trainer is never affected. + +The mode is gated: it is refused at launch unless ``sampler_snapshot_compat.json`` +(shipped with the package, so with the image) holds a ``pass`` record for the +pinned SGLang commit AND the weight-update path the island uses. A record is +written only after a GPU check compared text and weight checksums with the flags +on and off across RL weight updates (L8.1: ``tensor``; L8.2: ``distributed``). + +Miles does not expose ``--sglang-enable-memory-saver`` (it derives the field from +``--offload-rollout``), so the flags reach the engines as per-group overrides in +an inline ``--sglang-config`` that mirrors Miles' own default single-group layout. +""" + +from __future__ import annotations + +import base64 +import json +from dataclasses import dataclass +from importlib import resources +from typing import Any, Mapping, Sequence + +from .pins import SGLANG_NEXT_COMMIT + +COMPAT_FILE = "sampler_snapshot_compat.json" +SNAPSHOT_SERVER_ARGS = {"enable_memory_saver": True, "enable_weights_cpu_backup": True} +# Weight-update path -> how the sampler engines receive trainer weights. +UPDATE_PATH_TENSOR = "tensor" # colocated: update_weights_from_tensor (CUDA IPC) +UPDATE_PATH_DISTRIBUTED = "distributed" # fixed partition: NCCL update_weights_from_distributed + + +class SamplerSnapshotRefused(ValueError): + """Snapshot mode was requested but cannot be enabled; the message names why.""" + + +@dataclass(frozen=True) +class CompatRecord: + sglang_commit: str + update_path: str + result: str + check: str + evidence: str + + +def load_compat_records(text: str | None = None) -> tuple[CompatRecord, ...]: + """Parse the shipped compatibility file (or ``text``, for tests).""" + + if text is None: + text = resources.files(__package__).joinpath(COMPAT_FILE).read_text(encoding="utf-8") + raw = json.loads(text) + if raw.get("schema") != 1: + raise ValueError(f"{COMPAT_FILE}: unsupported schema {raw.get('schema')!r}") + return tuple( + CompatRecord( + sglang_commit=str(r["sglang_commit"]), + update_path=str(r["update_path"]), + result=str(r["result"]), + check=str(r["check"]), + evidence=str(r.get("evidence", "")), + ) + for r in raw.get("records", ()) + ) + + +def require_recorded_pass( + update_path: str, + *, + sglang_commit: str = SGLANG_NEXT_COMMIT, + records: Sequence[CompatRecord] | None = None, +) -> CompatRecord: + """Return the pass record for (commit, path) or refuse with the missing check named.""" + + if records is None: + records = load_compat_records() + for record in records: + if record.sglang_commit == sglang_commit and record.update_path == update_path: + if record.result == "pass": + return record + raise SamplerSnapshotRefused( + f"sampler snapshot mode refused: the snapshot compatibility check " + f"({record.check}) recorded {record.result!r}, not 'pass', for SGLang " + f"{sglang_commit[:9]} on the {update_path!r} weight-update path" + ) + raise SamplerSnapshotRefused( + "sampler snapshot mode refused: no recorded pass of the SGLang snapshot " + "compatibility check (flags on vs off across RL weight updates, design D5) for " + f"SGLang {sglang_commit[:9]} on the {update_path!r} weight-update path " + f"(see yeto/rl/adapters/miles/{COMPAT_FILE})" + ) + + +def sglang_config_arg(rollout_gpus: int) -> str: + """Inline ``--sglang-config`` value: Miles' default single group plus the snapshot flags. + + Same layout Miles builds without a config (one ``default`` model, one + ``regular`` group over all rollout GPUs), so only the overrides differ. + """ + + config = { + "sglang": [ + { + "name": "default", + "server_groups": [ + { + "worker_type": "regular", + "num_gpus": int(rollout_gpus), + "overrides": dict(SNAPSHOT_SERVER_ARGS), + } + ], + } + ] + } + payload = json.dumps(config, sort_keys=True, separators=(",", ":")) # JSON is valid YAML + return "base64:" + base64.b64encode(payload.encode()).decode() + + +def snapshot_argv( + request: Any, + *, + extra_argv: Sequence[str] = (), + records: Sequence[CompatRecord] | None = None, + sglang_commit: str = SGLANG_NEXT_COMMIT, +) -> list[str]: + """Miles argv enabling snapshot mode for a placement ``request``; refuses unsupported shapes.""" + + if request.kind == "colocated": + raise SamplerSnapshotRefused( + "sampler snapshot mode refused: it applies to dedicated sampler engines " + "(fixed partition); colocated engines share GPUs with the trainer" + ) + if getattr(request, "placement_map_arg", None) is not None: + raise SamplerSnapshotRefused( + "sampler snapshot mode refused: not supported with an explicit role -> bundle " + "placement map (standby / named rollout cells)" + ) + if "--sglang-config" in extra_argv: + raise SamplerSnapshotRefused( + "sampler snapshot mode refused: it writes --sglang-config itself; " + "drop --sglang-config from the extra Miles argv" + ) + require_recorded_pass(UPDATE_PATH_DISTRIBUTED, sglang_commit=sglang_commit, records=records) + return ["--sglang-config", sglang_config_arg(request.rollout_gpus)] + + +def decode_sglang_config_arg(value: str) -> Mapping[str, Any]: + """Inverse of :func:`sglang_config_arg` (tests and provenance readers).""" + + assert value.startswith("base64:"), value + return json.loads(base64.b64decode(value[len("base64:"):]).decode()) diff --git a/yeto/rl/adapters/miles/sampler_snapshot_compat.json b/yeto/rl/adapters/miles/sampler_snapshot_compat.json new file mode 100644 index 000000000..9b0156a8c --- /dev/null +++ b/yeto/rl/adapters/miles/sampler_snapshot_compat.json @@ -0,0 +1,31 @@ +{ + "schema": 1, + "flags": [ + "enable_memory_saver", + "enable_weights_cpu_backup" + ], + "records": [ + { + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "update_path": "tensor", + "result": "pass", + "check": "G-L8.1", + "date": "2026-10-14", + "gpu": "NVIDIA H100 80GB HBM3", + "model": "Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca", + "evidence": "infra-drafts/SESSION24/gpu/L8.1-review.md; s1-runs/s24-l8.1/verdict.json (modal call fc-01M4XGJ1BNVHVYK51WC6RGQAYE)", + "summary": "flags on vs off: generated text and full-model check_weights checksum equal after two update_weights_from_tensor updates; release/resume after each update keeps the updated weights" + }, + { + "sglang_commit": "2fa880182eefdb31e64d8ea70317be27f9f32121", + "update_path": "distributed", + "result": "pass", + "check": "G-L8.2", + "date": "2026-10-14", + "gpu": "2x NVIDIA H100 80GB HBM3", + "model": "Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca", + "evidence": "infra-drafts/SESSION24/gpu/L8.2-review.md; s1-runs/s24-l8.2/verdict.json (modal call fc-01M4XKWDD28W8KF0T784FSQ2GT)", + "summary": "trainer GPU0 -> engine GPU1 NCCL update_weights_from_distributed, flags on vs off: text and full-model checksum equal after two updates; release/resume (GPU1 35154 -> 2200 -> 34878 MiB) then a new NCCL group and a second update; resume keeps the updated weights" + } + ] +} diff --git a/yeto/rl/adapters/miles/state.py b/yeto/rl/adapters/miles/state.py index 1bfd3b45c..175296cf6 100644 --- a/yeto/rl/adapters/miles/state.py +++ b/yeto/rl/adapters/miles/state.py @@ -18,7 +18,14 @@ from yeto.rl.engine.policy_digest import WeightsChanged from yeto.rl.engine.trainable_state import LAYOUT_LORA, TrainableState, require_supported_layout from . import LoopRunner -from .state_plugin import APPLY_STATE, EXPORT_DIGEST, EXPORT_STATE, OPTIMIZER_MODES, WEIGHTS_VERSION +from .state_plugin import ( + APPLY_STATE, + EXPORT_DIGEST, + EXPORT_FULL_DIGEST, + EXPORT_STATE, + OPTIMIZER_MODES, + WEIGHTS_VERSION, +) class PolicyStateError(RuntimeError): @@ -285,6 +292,99 @@ def _apply_result(self, lora: Any, results: list[Any]) -> None: self.policy_version = lora.policy_version + +class MilesFullParameterLocalPolicyState(MilesPolicyState): + """``PolicyState`` of a no-sync full-parameter island (L4.6 sampler/trainer). + + The island has no outer sync, so nothing outside the trainer reads the + weights: the rollout engines get them through Miles ``update_weights``. + The state is only a trainer-resident handle whose policy hash is a content + digest of every model parameter (each rank hashes its shard on the trainer, + see ``state_plugin.export_full_digest``). There is no LoRA export and no + apply. The layout hash is the full-parameter identity bound by run_miles. + """ + + def __init__(self, *, actor_model: Any, base_model_revision: str, config_hash: str, + expected_layout_hash: str | None, policy_version: int = 0, + runner: LoopRunner | None = None) -> None: + from yeto.rl.engine.policy_digest import RESIDENT_LAYOUTS + from yeto.rl.engine.trainable_state import LAYOUT_FRAGMENTS + + if expected_layout_hash is None: + raise PolicyStateError("a full-parameter island needs its bound layout hash") + super().__init__(actor_model=actor_model, base_model_revision=base_model_revision, + config_hash=config_hash, expected_layout_hash=expected_layout_hash, + policy_version=policy_version, runner=runner) + self.layout = require_supported_layout(LAYOUT_FRAGMENTS, RESIDENT_LAYOUTS) + + def export(self, *, policy_version: int | None = None): + return self.export_digest(policy_version=policy_version) + + async def aexport(self, *, policy_version: int | None = None): + return self.export_digest(policy_version=policy_version) + + def export_digest(self, *, policy_version: int | None = None): + from yeto.rl.engine.policy_digest import TrainerResidentState + + version = self._export_version(policy_version) + digest, weights = self._full_digest(self._plugin(EXPORT_FULL_DIGEST, {"policy_version": version}), + version) + return TrainerResidentState( + layout=self.layout, + base_model_revision=self._revision, + config_hash=self._config_hash, + layout_hash=self.layout_hash, + policy_version=version, + digest=digest, + materialize_fn=_no_full_parameter_export, + weights_mark=weights, + holds_fn=self.check_holds, + ) + + def _current_digest(self, version: int): + digest, _ = self._full_digest(self._plugin(EXPORT_FULL_DIGEST, {"policy_version": version}), version) + return digest + + def _full_digest(self, results: list[Any], version: int): + """Combine the per-rank shard digests (``run_plugin`` order) into one.""" + + import hashlib + + from yeto.rl.engine.policy_digest import POLICY_TENSOR_PREFIX, PolicyDigest + + if any(not isinstance(r, dict) or r.get("policy_version") != version for r in results): + raise PolicyStateError("full-parameter digest: a rank returned no digest for this version") + marks = [r["weights"] for r in results if r.get("weights")] + if len(marks) != 1: + raise PolicyStateError(f"expected exactly one main-rank weights version, got {len(marks)}") + policy = hashlib.sha256(POLICY_TENSOR_PREFIX) + for part in (self._revision, self._config_hash, self.layout_hash): + policy.update(part.encode("ascii")) + for rank, result in enumerate(results): + policy.update(f"{rank}:{result['shard_hash']}".encode("ascii")) + tensor_hash = policy.hexdigest() + payload_bytes = sum(int(r["shard_bytes"]) for r in results) + digest = PolicyDigest( + policy_tensor_hash=tensor_hash, + payload_hash=hashlib.sha256(f"full-parameter-payload:{tensor_hash}".encode("ascii")).hexdigest(), + payload_bytes=payload_bytes, + specs=(), + hash_seconds=max(float(r.get("hash_seconds", 0.0)) for r in results), + ) + print(f"[rl] full-parameter digest v{version}: ranks={len(results)} bytes={payload_bytes} " + f"hash={digest.hash_seconds:.1f}s", flush=True) + return digest, dict(marks[0]) + + def apply(self, state: Any, *, optimizer: str, local_step: int) -> None: + raise PolicyStateError("a no-sync full-parameter island never applies an outside policy") + + async def aapply(self, state: Any, *, optimizer: str, local_step: int) -> None: + self.apply(state, optimizer=optimizer, local_step=local_step) + + +def _no_full_parameter_export(_policy_version: int): + raise PolicyStateError("full-parameter tensors stay in the trainer (no-sync island)") + def _describe_tensors(tensors: Any, limit: int = 400) -> str: items = [] for name in sorted(tensors)[:limit]: diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 60f92f6f6..9ee39b0af 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -40,6 +40,8 @@ # rl-publish-fastpath: export + hash on the trainer, return only the digests. EXPORT_DIGEST = f"{_PLUGIN_MODULE}.export_digest" WEIGHTS_VERSION = f"{_PLUGIN_MODULE}.weights_version" +# A no-sync full-parameter island (L4.6): per-rank digest of every model parameter. +EXPORT_FULL_DIGEST = f"{_PLUGIN_MODULE}.export_full_digest" # rl-publish-fastpath D3: version of the weights this trainer process holds. # Kept in each trainer process (module state), bumped on every path that writes @@ -656,6 +658,56 @@ def export_digest( "export_seconds": export_seconds, "weights": current_weights_version()} +FULL_SHARD_PREFIX = b"yeto-rl-full-parameter-shard-v1\0" + + +def export_full_digest(actor: Any, *, policy_version: int) -> dict[str, Any]: + """Plugin of a no-sync full-parameter island: hash this rank's model shard. + + Every rank returns its own shard digest (name, dtype, shape and raw bytes of + each local parameter, sorted by name); the caller combines the ranks in + ``run_plugin`` order. No tensor leaves the trainer and no collective runs. + The main rank also reports its weights version (rl-publish-fastpath D3). + """ + + import hashlib + import time + + import torch + + install_grad_norm_recorder() + started = time.monotonic() + digest = hashlib.sha256(FULL_SHARD_PREFIX) + total = 0 + count = 0 + with trainer_resident(actor): + named: dict[str, Any] = {} + for index, chunk in enumerate(actor.model): + for name, param in chunk.named_parameters(): + named[f"{index}.{name}"] = param + for name in sorted(named): + value = named[name].detach() + if value.device.type != "cpu": + value = value.to("cpu") + value = value.contiguous() + encoded = name.encode("utf-8") + digest.update(len(encoded).to_bytes(4, "little")) + digest.update(encoded) + digest.update(f"{value.dtype}|{list(value.shape)}".encode("ascii")) + raw = value.view(torch.uint8) if value.dim() else value.reshape(1).view(torch.uint8) + digest.update(memoryview(raw.numpy()).cast("B")) + total += raw.numel() + count += 1 + return { + "policy_version": int(policy_version), + "shard_hash": digest.hexdigest(), + "shard_bytes": total, + "shard_tensors": count, + "hash_seconds": round(time.monotonic() - started, 3), + "weights": current_weights_version() if _is_main_rank(actor) else None, + } + + def is_native_flash_next(actor: Any) -> bool: """Qwen3.8-Flash-Next (qwen4_exp) trains through Miles' native LoRA plugin. diff --git a/yeto/rl/adapters/verl/fully_async_ports.py b/yeto/rl/adapters/verl/fully_async_ports.py index b339348a2..0f59ff8b2 100644 --- a/yeto/rl/adapters/verl/fully_async_ports.py +++ b/yeto/rl/adapters/verl/fully_async_ports.py @@ -136,6 +136,7 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non completed=len(groups), aborted=0, payload=None, submitted_groups=pulled, abort_mechanism=DISCARD_MECHANISM if collector.discarded else None, carry_over=carry, + backend="verl", # Same field and meaning as Miles: groups aborted in flight and NOT # trained. verl's partial rollout continues interrupted requests # (counted as resumed_trajectories in carry_over, like Miles' diff --git a/yeto/rl/adapters/verl/fully_async_round.py b/yeto/rl/adapters/verl/fully_async_round.py index 7791b3a3e..d428bd0f2 100644 --- a/yeto/rl/adapters/verl/fully_async_round.py +++ b/yeto/rl/adapters/verl/fully_async_round.py @@ -29,6 +29,13 @@ from dataclasses import dataclass, field from typing import Any +from yeto.rl.engine.offpolicy_correction import ( + GATE_DROPPED_LENGTHS, + GATE_GRANULARITY, + GATE_GROUP_DROPPED, + GATE_UNKNOWN_VERSION_DROPPED, + GRANULARITY_TRAJECTORY, +) from yeto.rl.engine.overlap import ROLLOUT_OVERLAP_MODE from .fully_async_translate import ( @@ -285,9 +292,17 @@ def discard_tally(self) -> dict[str, int]: "unknown_groups": sum(1 for v in self.discarded if v.unknown)} def carry_over_fields(self) -> dict[str, Any]: - """The round's ``rl_rollout_carry_over`` fields (Miles names, carry_over.CARRY_FIELDS).""" + """The round's ``rl_rollout_carry_over`` fields (Miles names, carry_over.CARRY_FIELDS). + + Also the ``offpolicy_gate`` inputs the driver turns into its own event + (rl-offpolicy-correction-passthrough D5/D6): verl has no drop by version + at consumption, so this collector is it (a group more than ``limit`` + versions behind is dropped whole). verl gives versions per trajectory, + so the granularity is ``trajectory``; truncation by turn is not reported. + """ carried = [v for v in self.kept if v.versions and min(v.versions) < self.current_outer] over_age = [v for v in self.discarded if not v.unknown] + unknown = [v for v in self.discarded if v.unknown] return { "execution_mode": ROLLOUT_OVERLAP_MODE, "max_policy_age": self.limit, @@ -297,15 +312,18 @@ def carry_over_fields(self) -> dict[str, Any]: "over_age_discarded_groups": len(over_age), "over_age_discarded_trajectories": sum(v.meta.n for v in over_age), "over_age_discarded_tokens": sum(sum(v.meta.lengths) for v in over_age), - "unknown_version_discarded_groups": sum(1 for v in self.discarded if v.unknown), + "unknown_version_discarded_groups": len(unknown), "cross_version_tokens": sum(v.older_tokens for v in self.kept), "trained_response_tokens": sum(sum(v.meta.lengths) for v in self.kept), "resumed_trajectories": sum(v.resumed for v in self.kept), "segment_token_mismatches": sum(v.segment_mismatches for v in self.kept), "discard_reasons": sorted({v.reason for v in self.discarded if v.reason})[:5], + GATE_DROPPED_LENGTHS: { + GATE_GROUP_DROPPED: [n for v in over_age for n in v.meta.lengths], + GATE_UNKNOWN_VERSION_DROPPED: [n for v in unknown for n in v.meta.lengths]}, + GATE_GRANULARITY: GRANULARITY_TRAJECTORY, } - def group_metadata(collector: RoundCollector, rollout_id: int, published: Mapping[int, str]): """Kept groups as :class:`GroupMetadata`; the policy token names the oldest version (the driver checks it against the hashes it published within the window).""" diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 03ce5a310..9147d5965 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -25,6 +25,7 @@ import sys import time from pathlib import Path +from types import SimpleNamespace from . import config as vconf from .pins import VERL_COMMIT, VERL_ROOT, expected_versions @@ -287,13 +288,16 @@ def main(argv=None) -> int: algo = algorithm_options(spec_dict) chat = json.loads(args.apply_chat_template_kwargs) if args.apply_chat_template_kwargs else {} max_prompt = 512 + limit = int(args.rl_max_policy_age or 0) run = vconf.VerlRunConfig( model_path=model_path, train_file=train_file, val_file=val_file, out_dir=str(work / "out"), lora_rank=args.lora_r, lora_alpha=args.lora_r, lr=args.inner_lr, groups_per_round=args.groups_per_round, samples_per_group=args.samples_per_group, max_prompt_len=max_prompt, max_response_len=args.rollout_max_response_len, seed=args.seed, total_rounds=args.global_rounds, correction=algo["correction"], - tis_lower=algo.get("tis_lower", 0.0), tis_upper=algo["tis_upper"], + # limit > 0: the correction is mapped by offpolicy_correction in + # fully_async_plan (lower > 0 warns there, D3); the sync path keeps its check. + tis_lower=0.0 if limit else algo.get("tis_lower", 0.0), tis_upper=algo["tis_upper"], gpu_memory_utilization=float(os.environ.get("YETO_VERL_GPU_MEM", "0.5")), reward_path=str(Path(__file__).with_name("reward_fn.py")), reward_name=function_for(args.reward_function), chat_template_kwargs=chat, @@ -310,10 +314,10 @@ def main(argv=None) -> int: sync = "strict" if sync != "none" and not args.syncer: raise SystemExit("--syncer is required unless --rl-single-island-no-sync") - limit = int(args.rl_max_policy_age or 0) asserted_keys = vconf.ASSERTED_KEYS if limit: # agentic-rollout-utilization 6.4b: fully_async path overrides, asserted_keys = fully_async_plan(limit, overrides, spec_dict, args, sync) + append_correction_record(args, spec_dict) test = json.loads(TEST_FILE.read_text()) if TEST_FILE.is_file() else {} exit_after = (test.get("exit_after_version") or {}).get(str(lid)) env_exit = os.environ.get("YETO_VERL_TEST_EXIT_AFTER", "") # TEST ONLY "ISLAND:VERSION[,...]" @@ -369,12 +373,18 @@ def main(argv=None) -> int: def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, args, sync: str): """Limit > 0: fully_async overrides + asserted keys; the algorithm contract is checked here (bounded staleness needs ``execution.max_policy_staleness >= limit`` - and a TIS correction) so a wrong spec fails before verl starts.""" + and a TIS correction) so a wrong spec fails before verl starts. + + The ``algorithm.rollout_correction`` overrides come from the correction + mapping (rl-offpolicy-correction-passthrough D2/D3): decoupled PPO with + ``bypass_mode=False`` written over the fully_async yaml default; a + forbidden setting is refused here.""" from yeto.rl.engine.algorithm import AlgorithmSpec from yeto.rl.engine.execution_profile import check_algorithm_contract from .fully_async_round import (FULLY_ASYNC_ASSERTED_KEYS, execution_profile_for, fully_async_run_overrides) + from .offpolicy_correction import apply_to_overrides, correction_mapping from .policy_age import SUPPORT SUPPORT.check(limit) @@ -384,6 +394,7 @@ def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, a check_algorithm_contract(profile, spec) out = fully_async_run_overrides(overrides, limit, groups_per_round=args.groups_per_round, rounds=args.global_rounds) + out = apply_to_overrides(out, correction_mapping(spec)) from yeto.hw.catalog import device_family if device_family() == "ascend": @@ -394,6 +405,28 @@ def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, a return out, tuple(vconf.ASSERTED_KEYS) + FULLY_ASYNC_ASSERTED_KEYS +def append_correction_record(args, spec_dict: dict | None) -> None: + """rl-offpolicy-correction-passthrough D8 (as Miles' island entry): the island + tape records the full correction group, the verl values and ``prox_anchor``; + each field verl cannot map gets a ``correction_config_warning``.""" + from yeto.rl.engine.algorithm import AlgorithmSpec + from yeto.rl.engine.events import append_tape_event + from yeto.rl.engine.offpolicy_correction import CORRECTION_RECORD_EVENT, correction_record + + from .offpolicy_correction import BACKEND, correction_mapping + + spec = AlgorithmSpec.from_dict(spec_dict) if spec_dict else AlgorithmSpec() + if spec.correction.is_default(): + return # default tapes unchanged + mapping = correction_mapping(spec) + tape = SimpleNamespace(yeto_rl_event_tape=args.event_tape, + yeto_rl_learner_id=args.learner_id) + append_tape_event(tape, {"event": CORRECTION_RECORD_EVENT, + **correction_record(spec.correction, mapping, backend=BACKEND)}) + for warning in mapping.warnings: + append_tape_event(tape, warning.event()) + + def _override_value(overrides, key): for item in overrides: if item.startswith(key + "="): diff --git a/yeto/rl/adapters/verl/offpolicy_correction.py b/yeto/rl/adapters/verl/offpolicy_correction.py new file mode 100644 index 000000000..3d57b12f1 --- /dev/null +++ b/yeto/rl/adapters/verl/offpolicy_correction.py @@ -0,0 +1,175 @@ +"""verl mapping of the neutral correction group +(change ``rl-offpolicy-correction-passthrough`` L6.4, L6.6, L6.7, L6.11). + +:func:`correction_mapping` turns ``AlgorithmSpec.correction`` into Hydra +overrides under ``algorithm.rollout_correction`` (verl fork ``acad9875``, +``verl/trainer/ppo/rollout_corr_helper.py``). Mapping table (design D3): + +========================== ===================================================== +neutral verl +========================== ===================================================== +tis, token / sequence ``rollout_is=token|sequence``, ``rollout_is_threshold=upper`` +ratio_mask=icepop ``rollout_is_threshold='lower_upper'`` (weights outside + the range set to 0; ``loss_mask`` unchanged) +ratio_mask=mis ``rollout_rs=token_k1|seq_mean_k1``, + ``rollout_rs_threshold='lower_upper'`` (rejected tokens + leave ``response_mask``: the loss denominator shrinks) +batch_normalize ``rollout_is_batch_normalize=True`` +========================== ===================================================== + +Rows verl cannot express give a :class:`CorrectionWarning` (phase 1 runs on): +``lower > 0`` on plain TIS (verl clamps the upper side only), level +``sequence`` with the MIS mask (verl has the geometric mean only), ``veto``, +``opsm``, a custom function, and the MIS mask (denominator). + +Decoupled PPO (design D2): ``bypass_mode=False`` makes verl recompute +``old_log_probs`` as pi_prox; the fully_async trainer computes it on the +weights of the first local step of the current sync cycle +(``fully_async_trainer.py`` ``_compute_old_log_prob``), hence +``prox_anchor=first_local_step``. ``fully_async_ppo_trainer.yaml`` sets +``bypass_mode: True``, so every mapping writes ``bypass_mode`` explicitly and +:data:`FORBIDDEN` refuses ``bypass_mode=True``. + +``actor.use_rollout_log_probs`` (yaml: True) is not forbidden: no Python code +in the fork reads it (``grep`` over ``verl/`` at ``acad9875``: only the config +dataclasses declare it), so it does not change pi_prox (L6.6). +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from yeto.rl.engine.algorithm import AlgorithmSpec +from yeto.rl.engine.offpolicy_correction import ( + PROX_FIRST_LOCAL_STEP, + CorrectionMapping, + CorrectionWarning, + ForbiddenSetting, + refuse_forbidden, +) + +BACKEND = "verl" +PROX_ANCHOR = PROX_FIRST_LOCAL_STEP +PREFIX = "algorithm.rollout_correction." +BYPASS_KEY = PREFIX + "bypass_mode" + +# verl's own default (rollout_correction.yaml) when the spec gives no upper bound. +DEFAULT_UPPER = 2.0 +# verl IcePop needs lower > 0 (_parse_rollout_is_threshold); a missing or zero +# lower runs with this bound, which keeps every positive weight in practice. +ICEPOP_MIN_LOWER = 1e-8 + +FORBIDDEN = ( + ForbiddenSetting(BACKEND, f"{BYPASS_KEY}=True", + "it sets old_log_probs = rollout_log_probs, so pi_prox is lost"), +) + + +def _num(value: float) -> str: + return repr(float(value)) + + +def _range(lower: float, upper: float) -> str: + return f"{_num(lower)}_{_num(upper)}" + + +def _override_value(value: Any) -> str: + if value is None: + return "null" + if isinstance(value, bool): + return str(value) + if isinstance(value, float): + return _num(value) + if isinstance(value, str) and "_" in value: + return f"'{value}'" # "lower_upper" stays a string for Hydra + return str(value) + + +def _mapped(spec: AlgorithmSpec) -> tuple[dict[str, Any], list[CorrectionWarning]]: + c = spec.correction + values: dict[str, Any] = {"bypass_mode": bool(c.use_rollout_logprobs), + "rollout_is": None, "rollout_rs": None} + warnings: list[CorrectionWarning] = [] + + def warn(field: str, reason: str, ran_with: Any = None) -> None: + warnings.append(CorrectionWarning(field, BACKEND, reason, ran_with)) + + if c.method == "opsm": + warn("method", "verl has no OPSM; the run has no correction", "none") + if c.method == "custom": + warn("function", "verl takes no custom correction function; the run has no correction", + "none") + if c.method != "tis": + return values, warnings + + upper = DEFAULT_UPPER if c.upper is None else float(c.upper) + lower = 0.0 if c.lower is None else float(c.lower) + if c.veto is not None: + warn("veto", "verl has no per-sequence veto", None) + + if c.ratio_mask == "mis": + if lower <= 0.0: + warn("lower", "verl rejection sampling needs lower > 0; the k1 rule uses 1/upper", + 1.0 / upper) + lower = 1.0 / upper + values["rollout_rs"] = "token_k1" if c.level == "token" else "seq_mean_k1" + values["rollout_rs_threshold"] = _range(lower, upper) + if c.level == "sequence": + warn("level", "verl rejection sampling has no product-of-ratios rule; the " + "geometric mean (seq_mean_k1) runs", "geometric") + warn("ratio_mask", "verl removes rejected tokens from response_mask, so the loss " + "denominator shrinks; the neutral meaning is 'denominator unchanged'", + "denominator reduced") + if c.batch_normalize: + warn("batch_normalize", "verl normalizes IS weights only; the MIS mask maps to " + "rejection sampling without weights", False) + return values, warnings + + values["rollout_is"] = c.level # token | sequence (geometric needs the MIS mask) + if c.ratio_mask == "icepop": + if lower <= 0.0: + warn("lower", "verl IcePop needs lower > 0", ICEPOP_MIN_LOWER) + lower = ICEPOP_MIN_LOWER + values["rollout_is_threshold"] = _range(lower, upper) + else: + values["rollout_is_threshold"] = upper + if lower > 0.0: + warn("lower", "verl TIS clamps the upper side only (weights below lower pass " + "through)", 0.0) + values["rollout_is_batch_normalize"] = bool(c.batch_normalize) + return values, warnings + + +def correction_mapping(spec: AlgorithmSpec) -> CorrectionMapping: + """The verl settings of ``spec.correction`` with warnings and ``prox_anchor``.""" + + values, warnings = _mapped(spec) + keyed = {PREFIX + k: v for k, v in values.items()} + overrides = tuple(f"{k}={_override_value(v)}" for k, v in keyed.items()) + return CorrectionMapping(backend=BACKEND, prox_anchor=PROX_ANCHOR, values=keyed, + warnings=tuple(warnings), overrides=overrides) + + +def _normalized(override: str) -> str: + key, sep, value = override.lstrip("+~").partition("=") + if sep and value.strip().strip("'\"").lower() in ("true", "false"): + value = value.strip().strip("'\"").capitalize() + return f"{key}{sep}{value}" + + +def refuse_forbidden_overrides(overrides: Iterable[str]) -> None: + """Refuse Hydra overrides that hold a :data:`FORBIDDEN` assignment (names it).""" + + refuse_forbidden((_normalized(o) for o in overrides), FORBIDDEN) + + +def apply_to_overrides(overrides: Iterable[str], mapping: CorrectionMapping) -> list[str]: + """``overrides`` with every ``algorithm.rollout_correction.*`` entry replaced by + the mapping's (one value per key), then checked against :data:`FORBIDDEN`.""" + + out = [o for o in overrides if not o.lstrip("+~").startswith(PREFIX)] + out += list(mapping.overrides) + refuse_forbidden_overrides(out) + return out + diff --git a/yeto/rl/adapters/verl/policy_age.py b/yeto/rl/adapters/verl/policy_age.py index d5824c364..755c6ba68 100644 --- a/yeto/rl/adapters/verl/policy_age.py +++ b/yeto/rl/adapters/verl/policy_age.py @@ -17,9 +17,11 @@ from yeto.rl.engine.policy_age import PolicyAgeSupport, validate_limit # 6.4b (10-09): stage 2 through the fully_async path (rollouter + trainer on -# separate GPUs). The largest limit is 1 -- the only one the pre-launch review -# (rl-verl-backend 6.6) tests; sub-agent decision, raise after a GPU pass. -SUPPORT = PolicyAgeSupport(backend="verl", stage=2, max_policy_age=1) +# separate GPUs). rl-offpolicy-correction-passthrough L6.9: the cap was 1; +# decoupled PPO needs A_hard >= 2. verl has no upper bound of its own +# (staleness_threshold = N, fully_async_translate.staleness_threshold_for); +# 4 matches the Miles cap. Limits above 1 are not yet run on GPU (L6.10). +SUPPORT = PolicyAgeSupport(backend="verl", stage=2, max_policy_age=4) def policy_age_overrides(limit: int, *, groups_per_round: int | None = None) -> tuple[str, ...]: diff --git a/yeto/rl/adapters/verl/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index bed469f12..dad0f2801 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -253,6 +253,7 @@ def generate(self, rollout_id: int, *, expected_policy_version: str | None = Non rollout_id=rollout_id, policy_version=version, policy_hash=digest, groups=tuple(metas), completed=len(metas), aborted=0, payload=batch, batch_summary=summary, data_cursor=self.data_cursor(), + backend="verl", ) def abort(self) -> None: diff --git a/yeto/rl/algos/grpo_knobs.py b/yeto/rl/algos/grpo_knobs.py index f9c36a99e..93ea6a33d 100644 --- a/yeto/rl/algos/grpo_knobs.py +++ b/yeto/rl/algos/grpo_knobs.py @@ -258,16 +258,14 @@ def _reject_plugins(s) -> str | None: required = required_plugins(s) listed = {p.path: p.sha256 for p in s.plugins} missing = [r.path for r in required if r.path not in listed] - stale = [r.path for r in required if r.path in listed and listed[r.path] != r.sha256] - problems = [] + # rl-identity-tiering L9.2: a stale source hash is a startup warning + # (AlgorithmSpec.verify_plugins), not a rejection. if missing: - problems.append( + return ( f"plugins must list the code of every reward-pipeline / sample-filter module " f"{missing} (grpo_knobs.with_pipeline_plugins adds them)" ) - if stale: - problems.append(f"plugins {stale}: source hash differs from the current module source") - return "; ".join(problems) or None + return None register_rejection("grpo_knobs_reward_pipeline", _reject_pipeline) diff --git a/yeto/rl/algos/mismatch_correction.py b/yeto/rl/algos/mismatch_correction.py index 768eb8da8..2f53962dc 100644 --- a/yeto/rl/algos/mismatch_correction.py +++ b/yeto/rl/algos/mismatch_correction.py @@ -11,7 +11,7 @@ mechanism spec masks ================== =============================================== ========= mismatch_observe method=custom, function=:data:`OBSERVE_PATH` no - tis (P0 built-in) method=tis, tis_clip, tis_clip_low no + tis (P0 built-in) method=tis, upper, lower no icepop method=custom, function=:data:`ICEPOP_PATH` yes opsm_trainer OPSM, opsm_old_logprob_source=trainer yes opsm_rollout OPSM, opsm_old_logprob_source=rollout yes @@ -105,14 +105,16 @@ def selected_corrections(spec: AlgorithmSpec) -> tuple[str, ...]: out = [] c = spec.correction fn = _function(spec) - if c.method == "tis": + if c.method == "tis" and c.ratio_mask == "none": out.append("tis") if fn == OBSERVE_PATH: out.append("mismatch_observe") - if fn == ICEPOP_PATH: + if fn == ICEPOP_PATH or c.ratio_mask == "icepop": out.append("icepop") if fn == MIS_PATH: out.append("mis_mask" if c.mis_mode == "mask" else "mis") + if c.ratio_mask == "mis": # neutral MIS ratio mask (rl-offpolicy-correction D3) + out.append("mis_mask") if _uses_opsm(spec): out.append(f"opsm_{c.opsm_old_logprob_source}") return tuple(out) @@ -205,18 +207,18 @@ def _positive(path: str, value: Any): def _reject_bounds(spec: AlgorithmSpec) -> str | None: c = spec.correction fn = _function(spec) - if fn == ICEPOP_PATH and (c.tis_clip is None or c.tis_clip_low is None): - missing = [n for n in ("tis_clip", "tis_clip_low") if getattr(c, n) is None] + if fn == ICEPOP_PATH and (c.upper is None or c.lower is None): + missing = [n for n in ("upper", "lower") if getattr(c, n) is None] return ( f"IcePop needs explicit interval bounds: missing " f"{', '.join('correction.' + m for m in missing)} (no engine default is inherited)" ) - if (c.method == "tis" or fn == ICEPOP_PATH) and c.tis_clip is not None \ - and c.tis_clip_low is not None and not c.tis_clip_low < c.tis_clip: + if (c.method == "tis" or fn == ICEPOP_PATH) and c.upper is not None \ + and c.lower is not None and not c.lower < c.upper: what = "IcePop interval" if fn == ICEPOP_PATH else "TIS clip" return ( - f"{what}: correction.tis_clip_low={c.tis_clip_low} must be smaller than " - f"correction.tis_clip={c.tis_clip}" + f"{what}: correction.lower={c.lower} must be smaller than " + f"correction.upper={c.upper}" ) return None @@ -226,8 +228,8 @@ def _reject_observe(spec: AlgorithmSpec) -> str | None: if _function(spec) != OBSERVE_PATH: return None problems = [] - if c.tis_clip is not None or c.tis_clip_low is not None: - problems.append("correction.tis_clip/tis_clip_low (observe-only weights are 1)") + if c.upper is not None or c.lower is not None: + problems.append("correction.upper/lower (observe-only weights are 1)") if _uses_opsm(spec): problems.append("OPSM (observe-only cannot be combined with any correction)") if not c.mismatch_metrics: @@ -295,8 +297,8 @@ def _reject_mis(spec: AlgorithmSpec) -> str | None: ) if c.mis_batch_normalize and c.mis_level == "geometric": problems.append("correction.mis_batch_normalize supports mis_level token|sequence only") - if c.tis_clip is not None or c.tis_clip_low is not None: - problems.append("MIS uses mis_lower_bound/mis_upper_bound, not correction.tis_clip*") + if c.upper is not None or c.lower is not None: + problems.append("MIS uses mis_lower_bound/mis_upper_bound, not correction.lower/upper") if problems: return "MIS: " + "; ".join(problems) return None @@ -322,9 +324,11 @@ def mis_config(spec: AlgorithmSpec) -> dict[str, Any]: Empty unless the MIS function is selected. """ + c = spec.correction + if c.ratio_mask == "mis": + return neutral_mis_config(c) if _function(spec) != MIS_PATH: return {} - c = spec.correction return { "rs_level": c.mis_level, "rs_lower_bound": None, @@ -339,6 +343,29 @@ def mis_config(spec: AlgorithmSpec) -> dict[str, Any]: } +def neutral_mis_config(c: Any) -> dict[str, Any]: + """MIS attributes for the neutral ``ratio_mask='mis'`` (rl-offpolicy-correction D3). + + Mask mode over ``[lower, upper]`` at ``level``. ``veto`` needs Miles' + rejection-sampling branch (``use_rs``); its RS range is the same interval, + so RS adds only the veto (vendor/miles_mis.py, read 2026-10-14, not run on GPU). + """ + + veto = c.veto is not None + return { + "rs_level": c.level, + "rs_lower_bound": c.lower if veto else None, + "rs_upper_bound": c.upper if veto else None, + "rs_veto_threshold": c.veto, + "tis_batch_normalize": bool(c.batch_normalize), + "tis_level": c.level, + "tis_lower_bound": c.lower, + "tis_mode": "mask", + "tis_upper_bound": c.upper, + "use_rs": veto, + } + + register_runtime_attrs("mismatch_mis", mis_config) diff --git a/yeto/rl/elastic_client.py b/yeto/rl/elastic_client.py index b20997f3d..5e660be06 100644 --- a/yeto/rl/elastic_client.py +++ b/yeto/rl/elastic_client.py @@ -575,6 +575,9 @@ def __init__(self, config: ElasticClientConfig, key: bytes, *, self.ack: JoinAck | None = None self.errors: list[str] = [] self.latest_base: ElasticBase | None = None + # L5.6: (Unix time it became visible to the island, seconds from header to last byte) + # of ``latest_base``; set together with it under ``_base_cv``. + self.latest_base_timing: tuple[float, float] | None = None self._verdict_cv = threading.Condition() self._verdicts: list[SampleVerdict] = [] self._index_lock = threading.Lock() # one SAMPLE_INDEX in flight: replies are in order @@ -646,8 +649,16 @@ def _send(self, msg: Any) -> None: def _reader(self) -> None: while not self._stop.is_set(): + header_at = 0.0 + + def limit(msg_type: int) -> int: + nonlocal header_at + header_at = time.monotonic() # header read; payload follows + return _frame_limit(msg_type) + try: - t, payload = read_frame(self.sock, _frame_limit) + t, payload = read_frame(self.sock, limit) + recv_seconds = time.monotonic() - header_at except (OSError, ConnectionError, ValueError) as exc: if not self._stop.is_set(): self.errors.append(f"connection: {exc}") @@ -685,6 +696,7 @@ def _reader(self) -> None: self._base_seq += 1 if self.latest_base is None or msg.outer_version >= self.latest_base.outer_version: self.latest_base = msg + self.latest_base_timing = (time.time(), recv_seconds) self._base_cv.notify_all() self._on_event({"event": "rl_elastic_received", "type": t, **_fields(msg)}) diff --git a/yeto/rl/engine/algorithm.py b/yeto/rl/engine/algorithm.py index 4ef3d3cd8..68743915d 100644 --- a/yeto/rl/engine/algorithm.py +++ b/yeto/rl/engine/algorithm.py @@ -42,6 +42,8 @@ ALGORITHM_SPEC_SCHEMA = "yeto-rl-algorithm-spec-v1" ALGORITHM_SPEC_SCHEMA_V2 = "yeto-rl-algorithm-spec-v2" +# rl-identity-tiering D1: domain of AlgorithmSpec.semantic_sha256(). +ALGORITHM_SEMANTIC_SCHEMA = "yeto-rl-algorithm-semantic-v1" ALGORITHM_SPEC_SCHEMAS = (ALGORITHM_SPEC_SCHEMA, ALGORITHM_SPEC_SCHEMA_V2) # Neutral dynamic-sampling filter names (yeto-framework-decoupling 4.4, design # D6). The backend adapter binds each name to its implementation (Miles: @@ -116,6 +118,18 @@ def normalize_filter_name(name: str | None) -> str | None: KL_PLACEMENTS = ("none", "reward", "loss") KL_ESTIMATORS = ("k1", "k2", "k3", "low_var_kl") CORRECTION_METHODS = ["none", "tis", "opsm", "custom"] +# rl-offpolicy-correction-passthrough D1/D3: neutral correction fields. +CORRECTION_LEVELS = ("token", "sequence", "geometric") +# Optional ratio masks on top of method='tis' (default none, design A2). +RATIO_MASKS = ("none", "icepop", "mis") +# The correction group as every pre-D1 hash saw a default spec (design D8: +# phase 1 hashes this constant instead of the real group). +PHASE1_CORRECTION_IDENTITY = ( + ("function", None), ("method", "none"), ("mismatch_metrics", False), ("opsm_delta", None), + ("tis_clip", None), ("tis_clip_low", None), ("use_rollout_logprobs", False), +) +# Pre-D1 names of the bounds, still accepted on input (constructor and JSON). +CORRECTION_FIELD_ALIASES = {"tis_clip": "upper", "tis_clip_low": "lower"} DYNAMIC_SAMPLING_FILTERS = { BOUNDED_NONZERO_STD_FILTER, STOCK_NONZERO_STD_FILTER, SECRLENV_GROUP_CHECK_FILTER, } @@ -232,6 +246,22 @@ def _plugin(path: str, value: Any) -> "PluginRef | None": raise AlgorithmSpecError(f"{path} must be a plugin reference {{path, sha256}}, got {value!r}") +def _strip_plugin_hashes(value: Any) -> Any: + """Copy of canonical JSON ``value`` with ``sha256`` dropped from plugin refs. + + A plugin ref is the only canonical object whose keys are ``path`` plus an + optional ``sha256`` (:meth:`PluginRef.to_dict`). + """ + + if isinstance(value, Mapping): + if "path" in value and set(value) <= {"path", "sha256"}: + return {"path": value["path"]} + return {k: _strip_plugin_hashes(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_strip_plugin_hashes(v) for v in value] + return value + + # -------------------------------------------------------------------------- # PluginRef (design D7) # -------------------------------------------------------------------------- @@ -259,10 +289,17 @@ def plugin_source_sha256(path: str) -> str: @dataclass(frozen=True) class PluginRef: - """A plugin callable identified by dotted path + its module source SHA256.""" + """A plugin callable identified by dotted path and, optionally, the SHA256 + of its module source. + + rl-identity-tiering migration step 2: ``sha256`` may be absent, and a + present ``sha256`` that differs from the runtime source is a warning, not + a reject (:meth:`verify`). The path stays part of the semantic hash; the + source hash does not (:meth:`AlgorithmSpec.semantic_sha256`). + """ path: str - sha256: str + sha256: str | None = None def __post_init__(self) -> None: if not isinstance(self.path, str) or not self.path or " " in self.path: @@ -273,6 +310,8 @@ def __post_init__(self) -> None: f"{sorted(PLUGIN_NAMESPACES)}" ) digest = self.sha256 + if digest is None: + return if not isinstance(digest, str) or len(digest) != 64 or any( c not in "0123456789abcdef" for c in digest.lower() ): @@ -291,19 +330,28 @@ def from_path(cls, path: str) -> "PluginRef": @classmethod def from_dict(cls, payload: Mapping[str, Any], *, field_path: str = "plugin") -> "PluginRef": - if not isinstance(payload, Mapping) or set(payload) != {"path", "sha256"}: - raise AlgorithmSpecError(f"{field_path} must be an object with exactly path and sha256") - return cls(path=payload["path"], sha256=payload["sha256"]) + if (not isinstance(payload, Mapping) or "path" not in payload + or not set(payload) <= {"path", "sha256"}): + raise AlgorithmSpecError( + f"{field_path} must be an object with path and an optional sha256" + ) + return cls(path=payload["path"], sha256=payload.get("sha256")) def to_dict(self) -> dict[str, str]: - return {"path": self.path, "sha256": self.sha256} + # No sha256 key when absent, so every pinned ref keeps its old canonical form. + return {"path": self.path, **({"sha256": self.sha256} if self.sha256 else {})} + + def verify(self, *, import_callable: bool = True) -> str | None: + """Startup check: the module must be importable (else reject). - def verify(self, *, import_callable: bool = True) -> None: - """Startup check: importable and the source hash matches (else reject).""" + Returns a warning message when the pinned source hash differs from + the runtime source, else None. The caller records the warning. + """ actual = plugin_source_sha256(self.path) - if actual != self.sha256: - raise PluginIdentityError( + warning = None + if self.sha256 is not None and actual != self.sha256: + warning = ( f"plugin {self.path!r} source hash mismatch: spec {self.sha256}, " f"runtime {actual}" ) @@ -315,6 +363,7 @@ def verify(self, *, import_callable: bool = True) -> None: raise PluginIdentityError(f"plugin {self.path!r}: import failed: {exc}") from exc if not callable(getattr(loaded, name, None)): raise PluginIdentityError(f"plugin {self.path!r}: {name!r} is not a callable") + return warning # -------------------------------------------------------------------------- @@ -787,10 +836,28 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class CorrectionSpec(_Group): + """Off-policy / mismatch correction (rl-offpolicy-correction-passthrough D1). + + The neutral fields are written in full because the backends have + different defaults: ``level`` (token | sequence | geometric), ``lower`` and + ``upper`` (the importance-weight bounds; ``upper`` is the truncation bound + ``C``), ``veto``, ``batch_normalize`` and ``ratio_mask`` (none | icepop | + mis, on top of ``method='tis'``). Adapters map them to backend settings; + yeto never computes the weights. ``tis_clip`` / ``tis_clip_low`` are the + pre-D1 names of ``upper`` / ``lower``: accepted on input, read-only after. + + Phase 1 (design D8): the group does not enter ``AlgorithmSpec.sha256()``; + the run manifest and the island tapes record it instead. + """ + GROUP = "correction" method: str = "none" # none | tis | opsm | custom (custom tis function) - tis_clip: float | None = None - tis_clip_low: float | None = None + level: str = "token" + lower: float | None = None + upper: float | None = None + veto: float | None = None + batch_normalize: bool = False + ratio_mask: str = "none" opsm_delta: float | None = None function: PluginRef | None = None # --custom-tis-function-path use_rollout_logprobs: bool = False @@ -800,22 +867,43 @@ class CorrectionSpec(_Group): def __post_init__(self) -> None: s = object.__setattr__ s(self, "method", _choice("correction.method", self.method, CORRECTION_METHODS)) - s(self, "tis_clip", _number("correction.tis_clip", self.tis_clip, low=0.0, low_open=True)) - s(self, "tis_clip_low", _number("correction.tis_clip_low", self.tis_clip_low, low=0.0)) + s(self, "level", _choice("correction.level", self.level, CORRECTION_LEVELS)) + s(self, "ratio_mask", _choice("correction.ratio_mask", self.ratio_mask, RATIO_MASKS)) + s(self, "upper", _number("correction.upper", self.upper, low=0.0, low_open=True)) + s(self, "lower", _number("correction.lower", self.lower, low=0.0)) + s(self, "veto", _number("correction.veto", self.veto, low=0.0, low_open=True)) s(self, "opsm_delta", _number("correction.opsm_delta", self.opsm_delta, low=0.0, low_open=True)) s(self, "function", _plugin("correction.function", self.function)) + _boolean("correction.batch_normalize", self.batch_normalize) _boolean("correction.use_rollout_logprobs", self.use_rollout_logprobs) _boolean("correction.mismatch_metrics", self.mismatch_metrics) if self.method == "custom" and self.function is None: raise AlgorithmSpecError("correction.function is required for correction.method='custom'") if self.method != "custom" and self.function is not None: raise AlgorithmSpecError("correction.function requires correction.method='custom'") - if self.method not in ("tis", "custom") and ( - self.tis_clip is not None or self.tis_clip_low is not None - ): + weighting = self.method in ("tis", "custom") + if not weighting and (self.upper is not None or self.lower is not None): + raise AlgorithmSpecError( + "correction.upper/lower (tis_clip/tis_clip_low) require correction.method " + "tis or custom" + ) + if not weighting: + given = [name for name, default in (("level", "token"), ("veto", None), + ("batch_normalize", False)) + if getattr(self, name) != default] + if given: + raise AlgorithmSpecError( + f"{', '.join('correction.' + n for n in given)} require correction.method " + "tis or custom" + ) + if self.ratio_mask != "none" and self.method != "tis": raise AlgorithmSpecError( - "correction.tis_clip/tis_clip_low require correction.method tis or custom" + f"correction.ratio_mask={self.ratio_mask!r} requires correction.method='tis'" + ) + if self.level == "geometric" and self.ratio_mask != "mis" and self.method == "tis": + raise AlgorithmSpecError( + "correction.level='geometric' needs correction.ratio_mask='mis'" ) if self.method == "none" and self.opsm_delta is not None: raise AlgorithmSpecError( @@ -824,17 +912,54 @@ def __post_init__(self) -> None: ) if self.method == "opsm" and self.opsm_delta is None: raise AlgorithmSpecError("correction.opsm_delta is required for correction.method='opsm'") - if self.method == "tis" and (self.tis_clip is None or self.tis_clip_low is None): + if self.method == "tis" and (self.upper is None or self.lower is None): raise AlgorithmSpecError( - "correction.tis_clip and correction.tis_clip_low are required for " - "correction.method='tis' (no engine default is inherited)" + "correction.upper and correction.lower (tis_clip, tis_clip_low) are required " + "for correction.method='tis' (no engine default is inherited)" ) - if self.mismatch_metrics and self.method not in ("tis", "custom"): + if self.mismatch_metrics and not weighting: raise AlgorithmSpecError( "correction.mismatch_metrics needs a tis/custom correction function" ) self._normalize_ext() + @property + def tis_clip(self) -> float | None: + return self.upper + + @property + def tis_clip_low(self) -> float | None: + return self.lower + + @classmethod + def from_dict(cls, payload: Any): + if isinstance(payload, Mapping): + payload = _rename_correction_aliases(dict(payload)) + return super().from_dict(payload) + + +def _rename_correction_aliases(values: dict[str, Any]) -> dict[str, Any]: + for alias, name in CORRECTION_FIELD_ALIASES.items(): + if alias not in values: + continue + value = values.pop(alias) + if name in values and values[name] != value: + raise AlgorithmSpecError( + f"correction.{alias} is the old name of correction.{name}; give one of them") + values[name] = value + return values + + +_correction_init = CorrectionSpec.__init__ + + +def _correction_init_with_aliases(self, *args: Any, **kwargs: Any) -> None: + _correction_init(self, *args, **_rename_correction_aliases(kwargs)) + + +# The pre-D1 keyword names stay accepted; dataclasses.replace() sees only fields. +CorrectionSpec.__init__ = _correction_init_with_aliases # type: ignore[method-assign] + @dataclass(frozen=True) class SamplingSpec(_Group): @@ -1267,12 +1392,55 @@ def effective_default_at(self, path: str) -> Any: return self.default_at(path) def canonical_json(self) -> str: + """The full spec (correction group included): manifests, tapes, islands.""" return json.dumps( self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False ) + def identity_dict(self) -> dict[str, Any]: + """The payload of :meth:`sha256` (rl-offpolicy-correction-passthrough D8). + + Phase 1 leaves the correction group out of the hash: it is replaced + by the pre-D1 default group, so every spec whose correction group is + the default keeps its hash. A custom correction function stays in the + identity as a plugin reference (its source hash, design D8). Phase 2 + (L6b.1) puts the group back under a hash migration. + """ + + if self.correction.is_default(): + view = self + else: + plugins = list(self.plugins) + function = self.correction.function + if function is not None and function.path not in {p.path for p in plugins}: + plugins.append(function) + view = self.replace(correction=CorrectionSpec(), plugins=plugins) + out = view.to_dict() + if "correction" in out: + out["correction"] = dict(PHASE1_CORRECTION_IDENTITY) + return out + def sha256(self) -> str: - return hashlib.sha256(self.canonical_json().encode("utf-8")).hexdigest() + payload = json.dumps( + self.identity_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def semantic_dict(self) -> dict[str, Any]: + """The numerics-relevant part of the spec (rl-identity-tiering D1). + + The full v2 form with every plugin reduced to its path: a plugin + source change alone does not change it; any configuration change does. + """ + + return {"schema": ALGORITHM_SEMANTIC_SCHEMA, + "spec": _strip_plugin_hashes(self.structured_dict())} + + def semantic_sha256(self) -> str: + payload = json.dumps( + self.semantic_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() @property def schema(self) -> str: @@ -1324,9 +1492,12 @@ def referenced_plugins(self) -> tuple[PluginRef, ...]: refs.setdefault(value.path, value) return tuple(refs[k] for k in sorted(refs)) - def verify_plugins(self, *, import_callable: bool = True) -> None: - for ref in self.referenced_plugins(): - ref.verify(import_callable=import_callable) + def verify_plugins(self, *, import_callable: bool = True) -> tuple[str, ...]: + """Verify every referenced plugin; returns the source hash warnings.""" + + warnings = (ref.verify(import_callable=import_callable) + for ref in self.referenced_plugins()) + return tuple(w for w in warnings if w is not None) # -- derived (design D4/D6) ---------------------------------------------- def required_mechanisms(self) -> frozenset[tuple[str, str]]: diff --git a/yeto/rl/engine/backends.py b/yeto/rl/engine/backends.py index 59b23f6d3..243e07c98 100644 --- a/yeto/rl/engine/backends.py +++ b/yeto/rl/engine/backends.py @@ -56,6 +56,8 @@ class BackendEntry: # identity BackendIdentity of this adapter (phase 5) # image Modal-built engine image recipe (verl: ``verl-build:``) # policy_age SUPPORT: PolicyAgeSupport (agentic-rollout-utilization 2.1/6.3) + # offpolicy_correction correction_mapping(spec) + FORBIDDEN (rl-offpolicy-correction + # -passthrough; neutral side: yeto.rl.engine.offpolicy_correction) # overlay source patch applied to the engine checkout (launcher setup, run manifest) roles: Mapping[str, str] = field(default_factory=dict) @@ -79,6 +81,7 @@ class BackendEntry: "binding": f"{_MILES}.binding", "identity": f"{_MILES}.identity", "policy_age": f"{_MILES}.policy_age", + "offpolicy_correction": f"{_MILES}.offpolicy_correction", "overlay": f"{_MILES}.overlay", }), # rl-verl-backend: first step FSDP2 + vLLM, one GPU per island (roles the @@ -94,6 +97,7 @@ class BackendEntry: "identity": f"{_VERL}.identity", "image": f"{_VERL}.image", "policy_age": f"{_VERL}.policy_age", + "offpolicy_correction": f"{_VERL}.offpolicy_correction", }), } diff --git a/yeto/rl/engine/bridges.py b/yeto/rl/engine/bridges.py index 324ade1d7..b1f955d32 100644 --- a/yeto/rl/engine/bridges.py +++ b/yeto/rl/engine/bridges.py @@ -55,7 +55,9 @@ from . import progress as legacy_progress from .pause_audit import PAUSABLE_PHASE +from .sync_timing import SyncTiming, unix_now from .backend_identity import config_island_contract_sha256 +from .policy_digest import is_resident from .driver import IslandDriver, SyncBoundary, SyncStart from .leave_reason import sendable from .trainable_state import TrainableState @@ -85,7 +87,10 @@ def _local_state(driver: Any, version: int) -> Any: state = resident(policy_version=version) if callable(resident) else None if state is not None: return state - return TrainableState.from_lora(_at_version(_lora(driver.export_local()), version)) + exported = driver.export_local() + if is_resident(exported): # a full-parameter trainer exports no tensors at all + return exported.with_version(version) + return TrainableState.from_lora(_at_version(_lora(exported), version)) # --------------------------------------------------------------------------- @@ -104,6 +109,9 @@ def __init__(self, num_rollout: int) -> None: # rl-resume-from-checkpoint: our own stop after this many rounds (exclusive # rollout id; None = run to num_rollout). A final cut is kept at the stop. self.stop_after: int | None = None + # rl-trainer-start-stop D3: trigger of the final cut (None = MANUAL); the + # trainer lifecycle sets SCHEDULED for a stop the trainer scheduler ordered. + self.final_trigger: str | None = None def start(self, driver: IslandDriver) -> SyncStart: resumed = self.round_cuts.resume(driver) if self.round_cuts is not None else None @@ -118,8 +126,9 @@ def at_safe_point(self, driver: IslandDriver, *, rollout_id: int, final: bool = if self.round_cuts is None: return try: - info = (self.round_cuts.save(driver, rollout_id=rollout_id, final=True) if final - else self.round_cuts.save(driver, rollout_id=rollout_id)) + info = (self.round_cuts.save(driver, rollout_id=rollout_id, final=True, + **({"trigger": self.final_trigger} if self.final_trigger else {})) + if final else self.round_cuts.save(driver, rollout_id=rollout_id)) except Exception as error: # noqa: BLE001 driver.emit("rl_round_cut", rollout_id=rollout_id, ok=False, error=repr(error)[:2000]) return @@ -293,6 +302,7 @@ def __init__( self.bridge: StrictRlBridge | None = None self.current: CanonicalLoraState | None = None self.permit = None + self._submit_ts: float | None = None # L5.6: when the last PUSH was handed over def _apply(self, driver: IslandDriver, state: CanonicalLoraState) -> TrainableState: return driver.apply_policy( @@ -346,12 +356,18 @@ def _submit(self, driver, *, rollout_id, stats) -> None: # The strict delta is taken against the exact base version. local = _at_version(_lora(driver.export_local()), rollout_id) self.bridge.submit_local_state(self.permit, self.current, local, stats) + self._submit_ts = unix_now() self.bridge.release_current(rollout_id) self.current = None def _await(self, driver, *, rollout_id) -> CanonicalLoraState: driver.phase("wait_global", policy_version=rollout_id + 1) - return self.bridge.wait_for_global_policy(rollout_id + 1) + current = self.bridge.wait_for_global_policy(rollout_id + 1) + # The strict client receives broadcasts in the background, so the + # receive time is not separated: the wait includes it. + timing = SyncTiming(self._submit_ts, unix_now()) + driver.emit("rl_sync_timing", rollout_id=rollout_id, outer_sync="strict", **timing.fields()) + return current def _commit(self, driver, current: CanonicalLoraState) -> SyncBoundary: state = self._apply(driver, current) @@ -483,6 +499,15 @@ def _client_event(self, ev: dict) -> None: # 0.27: the launcher reads this as "left the elastic pool" (exit 7) driver.emit("elastic_left_pool", reason="rejoin_exhausted", failures=ev["failures"]) + def _base_timing(self, submit_ts: float, base) -> SyncTiming: + """L5.6 timing of ``base``; the client records receive time only + for its latest base, so an older one has no global_ready_ts.""" + timing = getattr(self.client, "latest_base_timing", None) + if timing is None or getattr(self.client, "latest_base", None) is not base: + return SyncTiming(submit_ts, None) + ready_ts, recv_seconds = timing + return SyncTiming(submit_ts, ready_ts, recv_seconds) + def _wait(self, newer_than): base = self.client.wait_base(newer_than=newer_than, timeout_s=self.base_wait_s) if base is None: @@ -615,10 +640,15 @@ def boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: self.client.raise_if_finished() self.client.delta_tensor(base_version=sent_on, c_tokens=int(stats.action_tokens), c_steps=self.config.local_optimizer_steps, update=update) + submit_ts = unix_now() if self.progress is not None: self.progress.commit_round(stats) driver.phase("wait_global", base_version=sent_on) - state = self._apply(driver, self._wait(sent_on), rollout_id + 1) + base = self._wait(sent_on) + driver.emit("rl_sync_timing", rollout_id=rollout_id, outer_sync="elastic", + base_version=int(base.outer_version), + **self._base_timing(submit_ts, base).fields()) + state = self._apply(driver, base, rollout_id + 1) self._maybe_pause_link(driver) except ElasticFinished as done: return self._finished(driver, done, rollout_id=rollout_id, sent_on=sent_on) @@ -734,6 +764,9 @@ def apply_policy(self, state: TrainableState, *, optimizer: str, local_step: int def phase(self, name: str, **fields: Any) -> None: self.driver.phase(f"critic_{name}", **fields) + def emit(self, event: str, **fields: Any) -> None: + self.driver.emit(event, channel="critic", **fields) + class DualStrictAvgSync: """Strict-avg of actor (LoRA, actor syncer) + critic (full parameters, critic @@ -931,6 +964,8 @@ def __init__( self.action_tokens = 0 self.finished = False self._driver: IslandDriver | None = None + # Wall seconds blocked on the syncer in the current boundary (L5.6). + self._remote_wait_seconds = 0.0 # legacy-compatible helpers ------------------------------------------ def _append_event(self, event: dict[str, Any]) -> None: @@ -1066,7 +1101,7 @@ def _commit_final_cut(self, manifest, final, *, stats, final_payload_bytes_recei return SyncBoundary(self._snapshot_state(self.snapshot), stop=True) def _finish(self, *, policy_version, stats) -> SyncBoundary: - manifest, final = self.bridge.wait_for_final_cut(policy_version=policy_version) + manifest, final = self._remote_wait(self.bridge.wait_for_final_cut, policy_version=policy_version) return self._commit_final_cut( manifest, final, @@ -1093,18 +1128,31 @@ def boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: import time started = time.monotonic() + self._remote_wait_seconds = 0.0 result = self._boundary(driver, rollout_id=rollout_id, stats=stats) self._append_event( { "event": "rl_sync_hook", "rl/rollout_id": rollout_id, "sync/hook_seconds": time.monotonic() - started, - "sync/remote_quorum_wait_seconds": 0.0, + "sync/remote_quorum_wait_seconds": self._remote_wait_seconds, "sync/finalization": result.stop, } ) return result + def _remote_wait(self, call, *args, **kwargs): + """Run a bridge call that blocks on the syncer and add its wall time + to this boundary's ``sync/remote_quorum_wait_seconds`` (L5.6). A + normal decoupled round does not block, so the value stays 0 there.""" + import time + + started = time.monotonic() + try: + return call(*args, **kwargs) + finally: + self._remote_wait_seconds += time.monotonic() - started + def _boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: if ( self.current is None @@ -1123,8 +1171,8 @@ def _boundary(self, driver, *, rollout_id, stats) -> SyncBoundary: self._record_local_round(stats, rollout_id=rollout_id) return self._finish(policy_version=next_rollout_id, stats=stats) if getattr(self.args, "yeto_rl_learner_budget_steps", None) == self.optimizer_steps: - consolidation = self.bridge.consolidate_budget( - local, optimizer_steps=self.optimizer_steps, action_tokens=self.action_tokens + consolidation = self._remote_wait( + self.bridge.consolidate_budget, local, optimizer_steps=self.optimizer_steps, action_tokens=self.action_tokens ) stats = replace( stats, diff --git a/yeto/rl/engine/convertible.py b/yeto/rl/engine/convertible.py new file mode 100644 index 000000000..6ce4191d4 --- /dev/null +++ b/yeto/rl/engine/convertible.py @@ -0,0 +1,455 @@ +"""Convertible nodes: rollout by default, trainer when the scheduler says so. + +rl-trainer-start-stop design D7 (Q5), tasks L7.1–L7.3. Neutral core: node +classes, the conversion port, the P0-S cut rule for in-flight trajectories, the +rollout capacity of the node pool, and the scheduler that converts nodes. +No backend names here; the Miles/SGLang port is +:mod:`yeto.rl.adapters.miles.convertible`. + +Node classes (D7): + +* ``trainer``: on-demand, fast interconnect, trainer precision; +* ``rollout``: may be spot, another region, lower precision; +* ``convertible``: a trainer-class node that does rollout by default. A spot node + or a node with a rollout precision other than the trainer's is refused. + +Conversion rollout -> trainer (in this order): + +1. the rollout engine stops taking requests; the in-flight trajectories are cut + by the P0-S rule (:func:`apply_cut_rule`): finished segments go back to the + buffer, unfinished trajectories are dropped and regenerated (P0-S-Q1 ③); +2. the rollout engine releases its GPU memory; +3. the trainer onloads on the same GPUs. + +Conversion trainer -> rollout: the trainer offloads, the rollout engine takes its +memory back and loads the published version, the weight checksum is compared +with the published checksum, then the engine takes requests again. + +The scheduler never converts the last rollout-capable node, and it counts the +rollout capacity the conversion takes away. The idle time of a convertible node +is the time it does neither rollout nor training. +""" + +from __future__ import annotations + +import enum +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol + +from .trainer_lifecycle import BacklogStats, LifecycleState, SampleIndexEntry, Start, Stop, TrainerScheduler + +CONVERSION_EVENT = "node_conversion" +CONVERSION_DECISION_KIND = "convertible_scheduler_decision" + + +# --------------------------------------------------------------------------- +# L7.1 Node classes +# --------------------------------------------------------------------------- +class NodeClass(str, enum.Enum): + ROLLOUT = "rollout" + TRAINER = "trainer" + CONVERTIBLE = "convertible" + + +class NodeClassError(ValueError): + """A node configuration breaks the class rule; the message names the node.""" + + +@dataclass(frozen=True) +class NodeSpec: + """One node of the run. ``precision`` is the rollout precision of the node + (None: the trainer precision). ``capacity`` is its rollout capacity in any + common unit (e.g. trajectories per second); only ratios are used.""" + + name: str + node_class: NodeClass + spot: bool = False + precision: str | None = None + capacity: float = 1.0 + + def __post_init__(self) -> None: + if not self.name: + raise NodeClassError("node name is empty") + if self.capacity <= 0: + raise NodeClassError(f"node {self.name}: capacity must be > 0, got {self.capacity}") + + @property + def does_rollout(self) -> bool: + return self.node_class is not NodeClass.TRAINER + + +def validate_nodes(nodes: Iterable[NodeSpec], *, trainer_precision: str) -> None: + """Refuse a convertible node that is spot or has another rollout precision.""" + seen: set[str] = set() + for node in nodes: + if node.name in seen: + raise NodeClassError(f"node {node.name}: name is used twice") + seen.add(node.name) + if node.node_class is not NodeClass.CONVERTIBLE: + continue + if node.spot: + raise NodeClassError( + f"node {node.name}: a spot node cannot be convertible " + "(only trainer-class nodes: on-demand, trainer precision)") + if node.precision is not None and node.precision != trainer_precision: + raise NodeClassError( + f"node {node.name}: rollout precision {node.precision} is not the trainer " + f"precision {trainer_precision}; only trainer-class nodes can be convertible") + + +def parse_node_spec(text: str) -> NodeSpec: + """``NAME=CLASS[:spot][:PRECISION][:capacity=X]``, e.g. ``island1=convertible``, + ``island2=rollout:spot:fp8``.""" + name, sep, rest = text.partition("=") + if not sep: + raise NodeClassError(f"node spec {text!r}: expected NAME=CLASS[:spot][:PRECISION]") + head, *options = rest.split(":") + try: + node_class = NodeClass(head) + except ValueError: + raise NodeClassError(f"node {name}: class {head!r} is not one of " + f"{[c.value for c in NodeClass]}") from None + spot, precision, capacity = False, None, 1.0 + for option in options: + if option == "spot": + spot = True + elif option.startswith("capacity="): + capacity = float(option.removeprefix("capacity=")) + elif option: + precision = option + return NodeSpec(name.strip(), node_class, spot=spot, precision=precision, capacity=capacity) + + +# --------------------------------------------------------------------------- +# L7.2 P0-S cut rule for in-flight trajectories +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class InflightTrajectory: + """A trajectory on the node when it stops taking requests. ``finished`` are + its segments that ended before the stop (each a complete sample with its + behavior version); ``unfinished`` is true when a part was still generating.""" + + trajectory_id: str + finished: tuple[SampleIndexEntry, ...] = () + unfinished: bool = True + + +@dataclass(frozen=True) +class CutOutcome: + returned: tuple[str, ...] = () # sample ids appended to the buffer + dropped: tuple[str, ...] = () # trajectory ids to regenerate + + +class SampleSink(Protocol): + def append(self, entry: SampleIndexEntry) -> None: ... + + +def apply_cut_rule(inflight: Iterable[InflightTrajectory], buffer: SampleSink) -> CutOutcome: + """P0-S cut: finished segments go to the buffer; the unfinished part is dropped + and the sampler regenerates the trajectory (P0-S-Q1 ③). A trajectory whose + segments all finished is not dropped.""" + returned: list[str] = [] + dropped: list[str] = [] + for trajectory in inflight: + for entry in trajectory.finished: + buffer.append(entry) + returned.append(entry.sample_id) + if trajectory.unfinished: + dropped.append(trajectory.trajectory_id) + return CutOutcome(tuple(returned), tuple(dropped)) + + +# --------------------------------------------------------------------------- +# L7.2 Conversion port and node +# --------------------------------------------------------------------------- +class ConversionPort(Protocol): + """Mechanism on one convertible node. The node calls the verbs in the order of + the module docstring; the port does not decide anything.""" + + def stop_requests(self) -> list[InflightTrajectory]: ... + def release_rollout(self) -> None: ... + def onload_trainer(self) -> None: ... + def offload_trainer(self) -> None: ... + def resume_rollout(self, version: int) -> str | None: ... # weight checksum after load + def accept_requests(self) -> None: ... + + +class Role(str, enum.Enum): + ROLLOUT = "rollout" + TO_TRAINER = "to_trainer" + TRAINER = "trainer" + TO_ROLLOUT = "to_rollout" + FAILED = "failed" + + +class ConversionError(RuntimeError): + pass + + +class LastRolloutNodeError(ConversionError): + """The conversion would leave no node doing rollout.""" + + +@dataclass(frozen=True) +class ConversionReceipt: + node: str + direction: str # "to_trainer" | "to_rollout" + seconds: float + phases: Mapping[str, float] + cut: CutOutcome | None = None + version: int | None = None + weight_checksum: str | None = None + + def event(self) -> dict[str, Any]: + out: dict[str, Any] = {"event": CONVERSION_EVENT, "node": self.node, + "direction": self.direction, "seconds": self.seconds, + "phases": dict(self.phases)} + if self.cut is not None: + out["returned"] = len(self.cut.returned) + out["dropped"] = len(self.cut.dropped) + if self.version is not None: + out["version"] = self.version + out["weight_checksum"] = self.weight_checksum + return out + + +class ConvertibleNode: + """Role state of one convertible node, plus its idle time (L7.3): the time in + ``TRAINER`` role while it does not train. Rollout time is never idle.""" + + def __init__(self, spec: NodeSpec, port: ConversionPort, *, + clock: Callable[[], float] = time.monotonic) -> None: + if spec.node_class is not NodeClass.CONVERTIBLE: + raise NodeClassError(f"node {spec.name}: class {spec.node_class.value} is not convertible") + self.spec = spec + self.port = port + self._clock = clock + self.role = Role.ROLLOUT + self._training = False + self._idle_s = 0.0 + self._idle_from: float | None = None + + @property + def name(self) -> str: + return self.spec.name + + # ------------------------------------------------------------ idle (L7.3) + def set_training(self, active: bool) -> None: + """The trainer on this node trains (True) or waits for samples (False).""" + self._training = active + self._update_idle() + + def idle_seconds(self) -> float: + now = self._clock() + return self._idle_s + (now - self._idle_from if self._idle_from is not None else 0.0) + + def _update_idle(self) -> None: + idle = self.role is Role.TRAINER and not self._training + now = self._clock() + if idle and self._idle_from is None: + self._idle_from = now + elif not idle and self._idle_from is not None: + self._idle_s += now - self._idle_from + self._idle_from = None + + def _set_role(self, role: Role) -> None: + self.role = role + if role is not Role.TRAINER: + self._training = False + self._update_idle() + + # ------------------------------------------------------------ conversion + def to_trainer(self, buffer: SampleSink) -> ConversionReceipt: + self._require(Role.ROLLOUT, "to_trainer") + self._set_role(Role.TO_TRAINER) + phases: dict[str, float] = {} + t0 = self._clock() + try: + inflight = self._timed(phases, "stop_requests", self.port.stop_requests) + cut = apply_cut_rule(inflight or (), buffer) + self._timed(phases, "release_rollout", self.port.release_rollout) + self._timed(phases, "onload_trainer", self.port.onload_trainer) + except BaseException as exc: + self._set_role(Role.FAILED) + raise ConversionError(f"node {self.name}: rollout -> trainer failed: {exc}") from exc + self._set_role(Role.TRAINER) + return ConversionReceipt(self.name, "to_trainer", self._clock() - t0, phases, cut=cut) + + def to_rollout(self, version: int, *, expected_checksum: str | None = None) -> ConversionReceipt: + """Back to rollout on published ``version``. With ``expected_checksum`` + (the published checksum of that version), a different engine checksum + fails the conversion and the engine does not take requests.""" + self._require(Role.TRAINER, "to_rollout") + self._set_role(Role.TO_ROLLOUT) + phases: dict[str, float] = {} + t0 = self._clock() + try: + self._timed(phases, "offload_trainer", self.port.offload_trainer) + checksum = self._timed(phases, "resume_rollout", lambda: self.port.resume_rollout(version)) + if expected_checksum is not None and checksum != expected_checksum: + raise ConversionError(f"weight checksum {checksum} != published {expected_checksum} " + f"for version {version}") + self._timed(phases, "accept_requests", self.port.accept_requests) + except BaseException as exc: + self._set_role(Role.FAILED) + raise ConversionError(f"node {self.name}: trainer -> rollout failed: {exc}") from exc + self._set_role(Role.ROLLOUT) + return ConversionReceipt(self.name, "to_rollout", self._clock() - t0, phases, + version=version, weight_checksum=checksum) + + def _require(self, role: Role, verb: str) -> None: + if self.role is not role: + raise ConversionError(f"node {self.name}: {verb} needs role {role.value}, " + f"node is {self.role.value}") + + def _timed(self, phases: dict[str, float], name: str, fn: Callable[[], Any]) -> Any: + t = self._clock() + try: + return fn() + finally: + phases[name] = self._clock() - t + + +# --------------------------------------------------------------------------- +# L7.3 Rollout capacity of the pool +# --------------------------------------------------------------------------- +class NodePool: + """All nodes of the run. Fixed rollout and trainer nodes have a fixed role; + convertible nodes carry their :class:`ConvertibleNode` role.""" + + def __init__(self, specs: Iterable[NodeSpec], convertible: Iterable[ConvertibleNode] = ()) -> None: + self.specs = {s.name: s for s in specs} + self.convertible = {n.name: n for n in convertible} + for name, node in self.convertible.items(): + if self.specs.get(name, node.spec) != node.spec: + raise NodeClassError(f"node {name}: convertible node spec differs from the pool spec") + self.specs[name] = node.spec + + def doing_rollout(self) -> list[str]: + out = [] + for name, spec in self.specs.items(): + node = self.convertible.get(name) + if node is not None: + if node.role is Role.ROLLOUT: + out.append(name) + elif spec.node_class is NodeClass.ROLLOUT: + out.append(name) + return out + + def total_capacity(self) -> float: + """Rollout capacity with every rollout-capable node doing rollout.""" + return sum(s.capacity for s in self.specs.values() if s.does_rollout) + + def capacity(self, without: Iterable[str] = ()) -> float: + skip = set(without) + return sum(self.specs[n].capacity for n in self.doing_rollout() if n not in skip) + + def check_conversion(self, names: Sequence[str]) -> float: + """Rollout capacity lost by converting ``names`` to trainer. Raises + :class:`LastRolloutNodeError` if no node would be left doing rollout.""" + rolling = set(self.doing_rollout()) + for name in names: + if name not in self.convertible or name not in rolling: + raise ConversionError(f"node {name}: not a convertible node doing rollout") + if not rolling - set(names): + raise LastRolloutNodeError( + f"converting {sorted(names)} leaves no node doing rollout " + "(the last rollout-capable node is never converted)") + return self.capacity() - self.capacity(without=names) + + def in_role(self, role: Role) -> list[ConvertibleNode]: + return [n for n in self.convertible.values() if n.role is role] + + +# --------------------------------------------------------------------------- +# L7.3 Scheduler over convertible nodes +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class ConversionDecision: + to_trainer: tuple[str, ...] = () + to_rollout: tuple[str, ...] = () + lost_capacity: float = 0.0 # rollout capacity the decision takes away + capacity_after: float = 0.0 + note: str = "" + + +@dataclass +class ConversionScheduler: + """Turns the start/stop decisions of a :class:`TrainerScheduler` (e.g. S1) into + conversions of ``group_size`` convertible nodes (a multi-GPU trainer needs a + group). Start converts a group to trainer; Stop converts it back to rollout. + The ski-rental stop of S1 counts only the time the converted group waits in + ``TRAINER`` role with an empty backlog, i.e. the time it does neither. + + A start that would convert the last rollout-capable node holds instead. Each + decision carries the rollout capacity it takes away (negative: gives back) + and the capacity after it, and is journaled with the pool capacity.""" + + inner: TrainerScheduler + pool: NodePool + group_size: int = 1 + record: Callable[..., None] | None = None + events: list[dict[str, Any]] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.group_size < 1: + raise ValueError(f"group_size must be >= 1, got {self.group_size}") + + def lifecycle_state(self) -> LifecycleState: + if self.pool.in_role(Role.TO_TRAINER) or self.pool.in_role(Role.TO_ROLLOUT): + return LifecycleState.WARMING # a transition: the inner scheduler holds + return LifecycleState.TRAINING if self.pool.in_role(Role.TRAINER) else LifecycleState.STOPPED + + def decide(self, backlog: BacklogStats, + prices: Mapping[str, float] | None = None) -> ConversionDecision: + state = self.lifecycle_state() + decision = self.inner.decide(backlog, state, prices) + out = ConversionDecision(capacity_after=self.pool.capacity()) + if isinstance(decision, Start): + out = self._plan_to_trainer() + elif isinstance(decision, Stop): + names = tuple(n.name for n in self.pool.in_role(Role.TRAINER)) + back = sum(self.pool.specs[n].capacity for n in names) + out = ConversionDecision(to_rollout=names, lost_capacity=-back, + capacity_after=self.pool.capacity() + back, + note=f"stop:{decision.reason.name}") + if self.record is not None and (out.to_trainer or out.to_rollout or out.note): + self.record(CONVERSION_DECISION_KIND, tx_id=None, + inputs={"lifecycle_state": state.value, "batches": backlog.batches, + "samples": backlog.samples, "capacity": self.pool.capacity(), + "total_capacity": self.pool.total_capacity()}, + decision={"to_trainer": list(out.to_trainer), + "to_rollout": list(out.to_rollout), + "lost_capacity": out.lost_capacity, + "capacity_after": out.capacity_after, "note": out.note}) + return out + + def _plan_to_trainer(self) -> ConversionDecision: + candidates = [n.name for n in self.pool.in_role(Role.ROLLOUT)] + if len(candidates) < self.group_size: + return ConversionDecision(capacity_after=self.pool.capacity(), + note=f"hold:only {len(candidates)} convertible nodes doing rollout") + # Convert the smallest capacities first: the least rollout lost. + group = tuple(sorted(candidates, key=lambda n: (self.pool.specs[n].capacity, n))[:self.group_size]) + try: + lost = self.pool.check_conversion(group) + except LastRolloutNodeError: + return ConversionDecision(capacity_after=self.pool.capacity(), note="hold:last_rollout_node") + return ConversionDecision(to_trainer=group, lost_capacity=lost, + capacity_after=self.pool.capacity() - lost) + + def apply(self, decision: ConversionDecision, buffer: SampleSink, *, version: int | None = None, + expected_checksum: str | None = None) -> list[ConversionReceipt]: + """Run the conversions of ``decision`` and keep their events.""" + receipts: list[ConversionReceipt] = [] + for name in decision.to_trainer: + receipts.append(self.pool.convertible[name].to_trainer(buffer)) + for name in decision.to_rollout: + if version is None: + raise ConversionError("to_rollout needs the published version") + receipts.append(self.pool.convertible[name].to_rollout( + version, expected_checksum=expected_checksum)) + self.events.extend(r.event() for r in receipts) + return receipts diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index f44e29cc4..7b6c34236 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -86,6 +86,7 @@ train_blockers, ) from .leave_reason import leave_reason_for +from .offpolicy_correction import OFFPOLICY_GATE_EVENT, gate_from_carry_over from .overlap import EvalOverlap, EvalStarter, overlap_refusal from .ports import ( Placement, @@ -615,8 +616,20 @@ def handshake(self) -> None: runtime_fingerprint=self.capabilities.runtime_fingerprint, **({"eval_guard": dict(self.eval_guard_report)} if self.eval_guard_report else {}), **({"island_tasks": dict(self.island_tasks)} if self.island_tasks else {}), + **self._code_identity(), ) + def _code_identity(self) -> dict[str, str]: + """rl-identity-tiering L9.1 dual write; absent without a code digest.""" + + from yeto.provenance import CODE_PACK_MODE_SOURCE_TREE + + if not self.code_pack_digest: + return {} + mode = getattr(self, "code_pack_mode", None) or CODE_PACK_MODE_SOURCE_TREE + return {"semantic_sha256": self.algorithm.semantic_sha256(), + "code_pack_digest": self.code_pack_digest, "code_pack_mode": mode} + # -- helpers used by sync sessions ----------------------------------- def apply_policy( self, state: TrainableState, *, optimizer: str, local_step: int, @@ -1259,7 +1272,14 @@ def _emit_carry_over(self, rollout_id: int, batch: Any) -> None: fields = getattr(batch, "carry_over", None) if not fields: return + fields = dict(fields) groups = getattr(batch, "groups", ()) or () + gate = gate_from_carry_over(fields, kept=len(groups), + backend=getattr(batch, "backend", None) or "unknown") + if gate: # rl-offpolicy-correction-passthrough D6: its own event per round + self.emit(OFFPOLICY_GATE_EVENT, rollout_id=rollout_id, + policy_version=int(getattr(batch, "policy_version", rollout_id)), + t=self.clock(), **self._labels(), **gate) self.emit("rl_rollout_carry_over", rollout_id=rollout_id, policy_version=int(getattr(batch, "policy_version", rollout_id)), t=self.clock(), **{**self._labels(), **dict(fields)}, @@ -1300,6 +1320,13 @@ def _is_final_round(self, rollout_id: int) -> bool: # rl-per-island-tasks: this island's task ids and prompt-file SHA256, written # into rl_driver_start as ``island_tasks`` only when the island has its own list. island_tasks: Mapping[str, Any] | None = None + # rl-identity-tiering L9.1: this island's code identity, written into + # rl_driver_start with ``semantic_sha256`` only when set. + code_pack_digest: str | None = None + code_pack_mode: str | None = None # None: the L9.1 transition (source tree hash) + # rl-trainer-cold-start-cost L0: the island's StartPhases; the first + # optimizer step records ``first_step_done`` (None: not recorded). + start_phases: Any = None def _export_eval_version(self, rollout_id: int, *, final: bool) -> None: state = self.published_state @@ -1414,6 +1441,8 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, started = time.monotonic() receipt = self.trainer.train_step(batch) train_seconds = time.monotonic() - started + if self.start_phases is not None: + self.start_phases.mark_once("first_step_done") raw = self.trainer.step_metrics() metrics = raw if isinstance(raw, TrainStepMetrics) else TrainStepMetrics(**dict(raw)) rollout_fraction = getattr(batch, "cross_version_truncated_fraction", None) diff --git a/yeto/rl/engine/fake_lifecycle.py b/yeto/rl/engine/fake_lifecycle.py new file mode 100644 index 000000000..a73bd034c --- /dev/null +++ b/yeto/rl/engine/fake_lifecycle.py @@ -0,0 +1,209 @@ +"""CPU fakes for the trainer lifecycle (rl-trainer-start-stop L3.1). + +* :class:`FakeMembership` records JOIN / LEAVE on a tape; +* :class:`CuttableFakeTrainer` is the fake trainer with the ``CuttableTrainer`` + verbs ``save_cut`` / ``restore_cut`` (weights, optimizer moments, step on disk); +* :class:`FakeRoundCuts` is the round-cut holder ``LocalOnlySync`` uses + (``save`` at the final safe point, ``resume`` at start). It writes every cut + through :func:`recovery.save` and is also the lifecycle's :class:`CutLog`. + +Not a GPU substitute. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable +from dataclasses import replace +from pathlib import Path +from typing import Any + +import torch + +from yeto.rl.elastic_client import LeaveReason + +from . import recovery +from .fake import FakeEngine, FakeRolloutPool, FakeTrainerGroup +from .resume import CheckpointStore +from .trainer_lifecycle import CutRef, SampleIndexEntry + +POINTER = "lifecycle-cut.json" + + +class FakeMembership: + """``tape``: ``("JOIN", incarnation)`` / ``("LEAVE", incarnation, reason)``.""" + + def __init__(self, base_version: int | None = None) -> None: + self.base_version = base_version + self.tape: list[tuple] = [] + + def join(self, incarnation: int) -> int | None: + self.tape.append(("JOIN", int(incarnation))) + return self.base_version + + def leave(self, incarnation: int, reason: LeaveReason) -> None: + self.tape.append(("LEAVE", int(incarnation), LeaveReason(reason))) + + +class CuttableFakeTrainer(FakeTrainerGroup): + """The fake trainer with the cut verbs. A cut is one file per cut id.""" + + def layout(self) -> dict[str, int]: + return {"tp": 1, "pp": 1, "dp": 1} + + def actual_layout(self) -> dict[str, int]: + return self.layout() + + def save_cut(self, *, epoch: int, context: Any) -> str: + e = self.engine + cut_id, root = context["cut_id"], Path(context["root"]) + root.mkdir(parents=True, exist_ok=True) + e.calls.append(("save_cut", cut_id)) + torch.save({"epoch": int(epoch), "tensors": e.tensors, "moments": e.moments, + "step": e.scheduler_step}, root / f"{cut_id}.pt") + return cut_id + + def restore_cut(self, cut_id: str, *, epoch: int, root: str, expect: Any, + shared_filesystem: bool = True) -> Any: + e = self.engine + e.calls.append(("restore_cut", cut_id)) + payload = torch.load(Path(root) / f"{cut_id}.pt", weights_only=True) + if int(payload["epoch"]) != int(epoch): + raise ValueError(f"cut {cut_id} is at epoch {payload['epoch']}, not {epoch}") + if set(payload["tensors"]) != set(e.tensors): + raise ValueError("layout mismatch") + for name, value in payload["tensors"].items(): + e.tensors[name].copy_(value) + for name, value in payload["moments"].items(): + e.moments[name].copy_(value) + e.scheduler_step = int(payload["step"]) + return {"cut_id": cut_id, "local_step": e.scheduler_step} + + +def cuttable_fake_engine(**kwargs: Any) -> FakeEngine: + """A fake engine (a freshly acquired trainer) whose trainer has the cut verbs.""" + engine = FakeEngine(**kwargs) + engine.trainer = CuttableFakeTrainer(engine) + return engine + + +class FakeRoundCuts: + """The round-cut holder of ``LocalOnlySync`` for fakes, and a ``CutLog``. + + Cuts only at the end of an incarnation (``final=True``). The pointer holds the + cut record (:data:`recovery.CUT_KEY`) and the published version at the cut; a + later start restores that cut into its fresh trainer.""" + + def __init__(self, root: str | os.PathLike[str], *, + watermark: Callable[[], int | None] | None = None) -> None: + self.store = CheckpointStore(root) + self.watermark = watermark # the durable buffer's consumed watermark (D4) + self.root = Path(root) + self.saved: list[dict[str, Any]] = [] + self.resumed: dict[str, Any] | None = None + + @property + def pointer_path(self) -> Path: + return self.root / POINTER + + def pointer(self) -> dict[str, Any] | None: + path = self.pointer_path + return json.loads(path.read_text(encoding="utf-8")) if path.is_file() else None + + def latest(self) -> CutRef | None: + info = self.pointer() + if info is None: + return None + cut = recovery.ConsistentCut.from_record(info[recovery.CUT_KEY]) + return CutRef(cut_id=cut.checkpoint_id, trigger=cut.trigger, + policy_version=int(info["policy_version"]), + watermark=cut.buffer_watermark) + + def save(self, driver: Any, *, rollout_id: int, final: bool = False, + trigger: str | None = None) -> dict[str, Any] | None: + if not final: + return None + current = self.pointer() + if current is not None and int(current["policy_version"]) == int(rollout_id): + return None # nothing trained since this cut + cut_id = f"r{int(rollout_id):06d}-n{len(self.saved)}" + driver.trainer.save_cut(epoch=0, context={"cut_id": cut_id, "root": str(self.root)}) + cut = recovery.ConsistentCut(checkpoint_id=cut_id, trigger=trigger or recovery.MANUAL, + outer_step=int(rollout_id), island_steps={"0": int(driver.local_step)}, + buffer_watermark=None if self.watermark is None else self.watermark()) + info = {"cut_id": cut_id, "policy_version": int(rollout_id), "local_step": int(driver.local_step)} + saved = recovery.save(cut, store=self.store, write=lambda record: self._write(info, record)) + self.saved.append(saved) + return {**info, recovery.CUT_KEY: saved[recovery.CUT_KEY]} + + def _write(self, info: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]: + tmp = self.pointer_path.with_suffix(".tmp") + tmp.write_text(json.dumps({**info, recovery.CUT_KEY: record}, sort_keys=True), encoding="utf-8") + os.replace(tmp, self.pointer_path) + return {} + + def resume(self, driver: Any) -> dict[str, Any] | None: + info = self.pointer() + if info is None: + return None + driver.trainer.restore_cut(info["cut_id"], epoch=0, root=str(self.root), expect=None) + driver.local_step = int(info["local_step"]) + self.resumed = info + return {"next_rollout_id": int(info["policy_version"]), **info} + + +class BufferFedFakeRollout(FakeRolloutPool): + """The fake rollout pool of a trainer that reads the durable buffer: each round + selects one batch, and ``commit_round`` (called by the lifecycle after the + round's boundary) commits its cursor. ``trained`` lists each round's sample + ids and behavior versions.""" + + def __init__(self, engine: FakeEngine, buffer: Any, *, consumer: str = "trainer") -> None: + super().__init__(engine) + self.buffer = buffer + self.consumer = consumer + self.batch = engine.groups * engine.samples_per_group + self.trained: dict[int, list[tuple[str, int]]] = {} + self._cursor: dict[int, int] = {} + + def generate(self, rollout_id: int, *, expected_policy_version: str | None = None) -> Any: + selected = self.buffer.select(self.consumer, self.batch) + if len(selected) < self.batch: + raise RuntimeError(f"backlog has {len(selected)} samples, a batch needs {self.batch}") + batch = super().generate(rollout_id, expected_policy_version=expected_policy_version) + spg = self.engine.samples_per_group + groups = tuple(replace(g, sample_ids=tuple(s.entry.sample_id for s in selected[i * spg:(i + 1) * spg])) + for i, g in enumerate(batch.groups)) + self.trained[rollout_id] = [(s.entry.sample_id, s.entry.behavior_version) for s in selected] + self._cursor[rollout_id] = selected[-1].position + 1 + return replace(batch, groups=groups) + + def commit_round(self, rollout_id: int) -> None: + self.buffer.commit_consumed(self._cursor.pop(rollout_id)) + + +class FakeSampler: + """A sampler that appends to the durable buffer with the version it loaded. + ``load`` takes the newest published version (D5); while no trainer runs, + ``latest_published`` returns None and the sampler keeps its version.""" + + def __init__(self, buffer: Any, latest_published: Callable[[], int | None], *, name: str = "s0") -> None: + self.buffer = buffer + self.latest_published = latest_published + self.name = name + self.loaded_version = 0 + self.produced = 0 + + def load(self) -> int: + version = self.latest_published() + if version is not None and version > self.loaded_version: + self.loaded_version = int(version) + return self.loaded_version + + def produce(self, n: int) -> None: + for _ in range(n): + sample_id = f"{self.name}-{self.produced:05d}" + ref = self.buffer.write_payload(sample_id, f"{sample_id}@v{self.loaded_version}".encode()) + self.buffer.append(SampleIndexEntry(sample_id, self.loaded_version, ref, tokens=6)) + self.produced += 1 diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index e936a8f5d..f4b2141d2 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -35,6 +35,7 @@ from yeto.rl.elastic_client import LeaveReason from .driver import SyncBoundary, SyncStart +from .sync_timing import SyncTiming, unix_now if TYPE_CHECKING: from .driver import IslandDriver @@ -507,6 +508,7 @@ def __init__( self.finished = False self.closed = False self._driver: IslandDriver | None = None + self._timing = SyncTiming(None, None) # L5.6, set by each boundary @property def actor(self) -> RoleStream: @@ -519,6 +521,7 @@ def start(self, driver: IslandDriver) -> SyncStart: def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> SyncBoundary: del stats + self._timing = SyncTiming(None, None) self.terminal = self.run(self._boundary(rollout_id)) versions = self.pairing.fragment_versions() state = self.actor.resident_state(rollout_id + 1) @@ -530,6 +533,7 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync "sync/fragment_versions": list(versions), "sync/global_policy_hash": state.policy_tensor_hash(), "sync/inflight_pairs": len(self.pairing.inflight), + **self._timing.fields(), }, ) return SyncBoundary(state=state, stop=self.terminal) @@ -544,9 +548,15 @@ async def _boundary(self, rollout_id: int) -> bool: ) before = pairing.fragment_versions() pairing.submit_ready() + submit_ts = unix_now() await pairing.commit_staged_broadcasts() if pairing.profile.sweep: await self._wait_for_sweep(before) + # Broadcasts are staged in the background, so the receive time + # is not separated: the wait includes it. + self._timing = SyncTiming(submit_ts, unix_now()) + else: + self._timing = SyncTiming(submit_ts, None) # streaming: no global wait if pairing.any_finalizing: await pairing.finish_without_more_training() return True diff --git a/yeto/rl/engine/island_ledger.py b/yeto/rl/engine/island_ledger.py index 8b6f2dee5..58d97277e 100644 --- a/yeto/rl/engine/island_ledger.py +++ b/yeto/rl/engine/island_ledger.py @@ -108,6 +108,14 @@ class StalenessPolicy: # Algorithm families whose value/critic needs on-policy data (design Q3). on_policy_only: bool = False + @classmethod + def for_algorithm(cls, spec: Any, **kwargs: Any) -> "StalenessPolicy": + """``correction`` derived from the algorithm spec, not the fixed ``"tis"`` + (rl-offpolicy-correction-passthrough D7).""" + from .offpolicy_correction import staleness_correction + + return cls(correction=staleness_correction(spec), **kwargs) + def __post_init__(self) -> None: if self.correction not in IS_CORRECTIONS: raise LedgerError(f"correction {self.correction!r} is not one of {IS_CORRECTIONS}") diff --git a/yeto/rl/engine/merge_metrics.py b/yeto/rl/engine/merge_metrics.py new file mode 100644 index 000000000..31bbbc784 --- /dev/null +++ b/yeto/rl/engine/merge_metrics.py @@ -0,0 +1,98 @@ +"""Per-outer-step merge metrics from the elastic syncer tape (rl-trainer-start-stop D9, L5.5). + +The elastic syncer (``syncer/src/elastic.rs``) writes one JSONL record per +event, keyed by ``kind``. :class:`MergeMetricsFolder` reads those records in +order and returns one flat metrics dict per ``outer_step``: + +* ``merge/island//weight_share`` - the island's normalized merge weight + (``w = c_tokens**2 / c_steps``, carried deltas discounted by + ``carry_gamma**lag``), summed over its entries in the step. +* ``merge/carried_weight_share`` - the share of the step's weight that comes + from carried (late) deltas; ``merge/carried_fraction`` - carried deltas + over all deltas submitted in the round; ``merge/carried_discount_mean``. +* ``merge/rejected_fraction`` and ``merge/rejected_fraction/`` - + rejected deltas over all deltas submitted in the round. +* Timing from the tape clock (L5.7) when present: ``merge/step_seconds``, + ``merge/last_arrival_island``, ``merge/time_unix``. + +"Submitted in the round" counts every ``delta_accepted``, +``delta_carried_over`` and ``delta_rejected`` since the previous +``outer_step``. A carried delta is counted in the round where it arrived, +which is the step that merges it. Values that the tape does not hold stay +absent (never 0). +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any + + +class MergeMetricsFolder: + """Fold elastic tape records into one metrics dict per outer step.""" + + def __init__(self) -> None: + self._reset() + + def _reset(self) -> None: + self._accepted = 0 + self._carried_discounts: list[float] = [] + self._rejected: Counter[str] = Counter() + + def feed(self, record: dict[str, Any]) -> dict[str, Any] | None: + """Consume one tape record; return the metrics on ``outer_step``.""" + kind = record.get("kind") + if kind == "delta_accepted": + self._accepted += 1 + elif kind == "delta_carried_over": + discount = record.get("discount") + self._carried_discounts.append(float(discount) if discount is not None else float("nan")) + elif kind == "delta_rejected": + self._rejected[str(record.get("reason", "unknown"))] += 1 + elif kind == "outer_step": + metrics = self._step_metrics(record) + self._reset() + return metrics + return None + + def _step_metrics(self, rec: dict[str, Any]) -> dict[str, Any]: + base_version = int(rec.get("base_version", 0)) + metrics: dict[str, Any] = { + "global_step": base_version + 1, + "merge/base_version": base_version, + "merge/arrived": int(rec.get("arrived", 0)), + "merge/timed_out": int(bool(rec.get("timed_out", False))), + } + shares: dict[int, float] = {} + carried_share = 0.0 + carried_in_step = 0 + for island, entry_base, weight in rec.get("weights") or []: + shares[int(island)] = shares.get(int(island), 0.0) + float(weight) + if int(entry_base) < base_version: + carried_share += float(weight) + carried_in_step += 1 + for island, share in sorted(shares.items()): + metrics[f"merge/island/{island}/weight_share"] = share + metrics["merge/carried_weight_share"] = carried_share + metrics["merge/carried_in_step"] = carried_in_step + + rejected = self._rejected + if "rejected" in rec: # the step's own list wins over the counted events + rejected = Counter(str(reason) for _island, reason in rec["rejected"]) + carried = len(self._carried_discounts) + submitted = self._accepted + carried + sum(rejected.values()) + metrics["merge/submitted"] = submitted + if submitted: + metrics["merge/carried_fraction"] = carried / submitted + metrics["merge/rejected_fraction"] = sum(rejected.values()) / submitted + for reason, count in sorted(rejected.items()): + metrics[f"merge/rejected_fraction/{reason}"] = count / submitted + if self._carried_discounts: + metrics["merge/carried_discount_mean"] = sum(self._carried_discounts) / carried + + for source, name in (("step_seconds", "merge/step_seconds"), + ("last_arrival_island", "merge/last_arrival_island"), + ("time_unix", "merge/time_unix")): + if rec.get(source) is not None: + metrics[name] = rec[source] + return metrics diff --git a/yeto/rl/engine/offpolicy_correction.py b/yeto/rl/engine/offpolicy_correction.py new file mode 100644 index 000000000..114413149 --- /dev/null +++ b/yeto/rl/engine/offpolicy_correction.py @@ -0,0 +1,394 @@ +"""Off-policy correction pass-through: the backend-neutral part +(change ``rl-offpolicy-correction-passthrough``, phase 1). + +yeto does not compute importance weights. A backend adapter maps the neutral +``CorrectionSpec`` (design D1) to its own settings and returns a +:class:`CorrectionMapping`. This module holds what every adapter shares: + +* :class:`CorrectionMapping` / :class:`CorrectionWarning` -- the mapped + values, the proximal-policy anchor (``prox_anchor``, design D2) and the + fields the backend could not map (``correction_config_warning``, D3); +* :class:`ForbiddenSetting` / :func:`refuse_forbidden` -- backend settings + that break decoupled PPO (Miles ``--use-rollout-logprobs``, verl + ``bypass_mode=True``, ...). They are refused at launch with the setting + named, also in phase 1; +* :func:`correction_record` / :func:`compare_island_records` -- the manifest + and tape record (D8) and the cross-island comparison (D2, D8); +* :func:`start_version` / :func:`start_version_gap` -- the version gap from + the start version (D4); +* :func:`offpolicy_gate_event` -- the per-round drop and truncation event (D6); +* :func:`cross_version_bounds` / :func:`staleness_correction` -- the + observe-only cross-version IS takes its bounds and name from the spec (D7). + +Phase 1 warns on unmappable settings and on island differences; it does not +refuse them (phase 2, L6b.2, will). +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec + +CORRECTION_RECORD_SCHEMA = "yeto-offpolicy-correction-v1" +# Tape / manifest event names (design D2, D3, D6, D8). +CORRECTION_RECORD_EVENT = "rl_offpolicy_correction" +CORRECTION_CONFIG_WARNING = "correction_config_warning" +OFFPOLICY_GATE_EVENT = "offpolicy_gate" + +# Which weights the backend uses for pi_prox (design D2). +PROX_PER_STEP = "per_step" # recomputed on the weights before each update (Miles) +PROX_FIRST_LOCAL_STEP = "first_local_step" # first local step of the cycle (verl) +PROX_ANCHORS = (PROX_PER_STEP, PROX_FIRST_LOCAL_STEP) + +# offpolicy_gate reasons (design D6). Backends may add their own admission drops. +GATE_KEPT = "kept" +GATE_GROUP_DROPPED = "group_dropped" # gap over A_hard at consumption, whole group +GATE_UNKNOWN_VERSION_DROPPED = "unknown_version_dropped" +GATE_REASONS = (GATE_KEPT, GATE_GROUP_DROPPED, GATE_UNKNOWN_VERSION_DROPPED) +# Groups the backend dropped as stale inside its own buffer (Miles --fully-async +# ``stale_dropped_groups``), counted apart from the yeto gate drops. Present only +# when the backend reported the count; no lengths (the backend does not give them). +GATE_BACKEND_STALE_DROPPED = "backend_stale_dropped" +# Response-length bucket edges (tokens) of the dropped-trajectory histogram. +# Fixed in code so that rounds and runs stay comparable. +DROPPED_LENGTH_EDGES = (256, 1024, 4096, 16384) +GRANULARITY_TOKEN = "token" +GRANULARITY_TRAJECTORY = "trajectory" +GRANULARITIES = (GRANULARITY_TOKEN, GRANULARITY_TRAJECTORY) +# Round-metadata keys an adapter fills for the driver (removed from +# rl_rollout_carry_over): dropped lengths per reason, truncation per turn, +# and the granularity of the turn counts. +GATE_DROPPED_LENGTHS = "offpolicy_gate_dropped_lengths" +TRUNCATED_BY_TURN = "cross_version_truncated_by_turn" +GATE_GRANULARITY = "offpolicy_gate_granularity" + + +class CorrectionConfigError(ValueError): + """A backend setting that breaks the decoupled PPO objective; launch refused.""" + + +# -------------------------------------------------------------------------- +# mapping, warnings, forbidden settings +# -------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class CorrectionWarning: + """One field a backend cannot map as written (design D3 "warning" rows).""" + + field: str # neutral field, e.g. "lower" + backend: str + reason: str + ran_with: Any = None # the value the backend actually runs with + + def to_dict(self) -> dict[str, Any]: + return {"field": self.field, "backend": self.backend, "reason": self.reason, + "ran_with": self.ran_with} + + def event(self) -> dict[str, Any]: + return {"event": CORRECTION_CONFIG_WARNING, "level": "warning", **self.to_dict()} + + +@dataclass(frozen=True) +class CorrectionMapping: + """What an adapter made of the correction group. + + ``values``: backend setting -> value (Miles: flag -> argv value or True; + verl: Hydra key -> value). ``argv`` / ``runtime_attrs`` / ``overrides`` + are the forms the adapter passes on (empty when unused by the backend). + """ + + backend: str + prox_anchor: str + values: Mapping[str, Any] = field(default_factory=dict) + warnings: tuple[CorrectionWarning, ...] = () + argv: tuple[str, ...] = () + runtime_attrs: Mapping[str, Any] = field(default_factory=dict) + overrides: tuple[str, ...] = () + + def __post_init__(self) -> None: + if self.prox_anchor not in PROX_ANCHORS: + raise ValueError(f"prox_anchor must be one of {PROX_ANCHORS}, got {self.prox_anchor!r}") + + +@dataclass(frozen=True) +class ForbiddenSetting: + """A backend setting refused at launch because it breaks decoupled PPO.""" + + backend: str + setting: str # Miles flag, verl Hydra assignment ("key=value") + reason: str + + def message(self) -> str: + return (f"{self.setting} is refused on backend {self.backend!r}: {self.reason} " + "(rl-offpolicy-correction-passthrough D2)") + + +def refuse_forbidden(present: Iterable[str], forbidden: Iterable[ForbiddenSetting]) -> None: + """Raise :class:`CorrectionConfigError` naming the first forbidden setting in ``present``.""" + + table = {f.setting: f for f in forbidden} + for setting in present: + if setting in table: + raise CorrectionConfigError(table[setting].message()) + + +# -------------------------------------------------------------------------- +# manifest / tape record and the cross-island comparison (D2, D8) +# -------------------------------------------------------------------------- + + +def correction_record(correction: CorrectionSpec, mapping: CorrectionMapping | None, + *, backend: str) -> dict[str, Any]: + """The full group, the backend values and ``prox_anchor`` (manifest and tape). + + ``mapping`` None: the backend has no correction mapping yet; the record + says so (``mapped`` None) and the caller warns. + """ + + return { + "schema": CORRECTION_RECORD_SCHEMA, + "backend": backend, + "group": correction.to_dict(), + "mapped": None if mapping is None else dict(mapping.values), + "prox_anchor": None if mapping is None else mapping.prox_anchor, + "warnings": [] if mapping is None else [w.to_dict() for w in mapping.warnings], + } + + +def compare_island_records(records: Mapping[str, Mapping[str, Any]]) -> dict[str, Any] | None: + """One ``correction_config_warning`` listing every field that differs + between islands (group fields and ``prox_anchor``), or None. + + ``records``: island name -> :func:`correction_record`. Phase 1 warns only. + """ + + if len(records) < 2: + return None + islands = sorted(records) + keys = sorted({k for r in records.values() for k in (r.get("group") or {})}) + differences: dict[str, dict[str, Any]] = {} + for key in keys: + values = {i: (records[i].get("group") or {}).get(key) for i in islands} + if len({_comparable(v) for v in values.values()}) > 1: + differences[key] = values + anchors = {i: records[i].get("prox_anchor") for i in islands} + if len(set(anchors.values())) > 1: + differences["prox_anchor"] = anchors + if not differences: + return None + return {"event": CORRECTION_CONFIG_WARNING, "level": "warning", "scope": "cross_island", + "reason": "islands run different correction settings", "differences": differences} + + +def backend_mapping(spec: AlgorithmSpec, backend: str) -> CorrectionMapping | None: + """The ``offpolicy_correction`` role of ``backend`` applied to ``spec``. + + Adapter interface (role ``offpolicy_correction`` in + :mod:`yeto.rl.engine.backends`): a module with ``correction_mapping(spec) + -> CorrectionMapping`` and ``FORBIDDEN`` (a tuple of + :class:`ForbiddenSetting`). None when the backend has no such role yet. + """ + + from . import backends + + try: + module = backends.module("offpolicy_correction", backend) + except backends.UnknownBackend: + return None + return module.correction_mapping(spec) + + +def run_correction_manifest(spec: AlgorithmSpec, + island_backends: Mapping[str, str]) -> dict[str, Any]: + """The run manifest's ``offpolicy_correction`` section (design D8, L6.5). + + ``island_backends``: island name -> backend. Returns the record of every + island and the ``correction_config_warning`` events: one per field a + backend cannot map, one per backend without a mapping, and at most one + for differences between islands. Phase 1: warnings only. + """ + + records: dict[str, dict[str, Any]] = {} + warnings: list[dict[str, Any]] = [] + mappings: dict[str, CorrectionMapping | None] = {} + for island, backend in sorted(island_backends.items()): + if backend not in mappings: + mapping = mappings[backend] = backend_mapping(spec, backend) + if mapping is None and not spec.correction.is_default(): + warnings.append(CorrectionWarning( + "*", backend, "the backend adapter has no correction mapping yet; the " + "group is recorded but not passed on").event()) + elif mapping is not None: + warnings.extend(w.event() for w in mapping.warnings) + records[island] = correction_record(spec.correction, mappings[backend], backend=backend) + cross = compare_island_records(records) + if cross is not None: + warnings.append(cross) + return {"islands": records, "warnings": warnings} + + +def _comparable(value: Any) -> Any: + if isinstance(value, Mapping): + return tuple(sorted((k, _comparable(v)) for k, v in value.items())) + if isinstance(value, list): + return tuple(_comparable(v) for v in value) + return value + + +# -------------------------------------------------------------------------- +# version gap from the start version (D4) +# -------------------------------------------------------------------------- + + +def start_version(token_versions: Sequence[int | None]) -> int | None: + """The version of the first response token whose version is known.""" + + return next((v for v in token_versions if v is not None), None) + + +def start_version_gap(trainer_version: int, token_versions: Sequence[int | None]) -> int | None: + """Current trainer version minus the trajectory start version (design D4).""" + + start = start_version(token_versions) + return None if start is None else trainer_version - start + + +def oldest_is_start(token_versions: Sequence[int | None]) -> bool | None: + """L6.12: whether the oldest token version equals the start version. + + True for a continued trajectory (versions never decrease along the + response). None when no version is known. + """ + + known = [v for v in token_versions if v is not None] + return None if not known else min(known) == known[0] + + +def segments_start_check(records: Iterable[Mapping[str, Any]], + key: str = "policy_versions") -> dict[str, int]: + """L6.12 on a tape: counts of records whose oldest segment version is (not) + the first one. ``key`` holds ``[[version, start, end), ...]`` segments + (``rl_trajectory_reward.policy_versions``).""" + + out = {"checked": 0, "oldest_is_start": 0, "oldest_before_start": 0} + for record in records: + segments = record.get(key) + if not segments or not isinstance(segments[0], (list, tuple)): + continue + versions = [int(s[0]) for s in sorted(segments, key=lambda s: int(s[1]))] + out["checked"] += 1 + out["oldest_is_start" if oldest_is_start(versions) else "oldest_before_start"] += 1 + return out + + +# -------------------------------------------------------------------------- +# offpolicy_gate (D6) +# -------------------------------------------------------------------------- + + +def length_bucket(length: int) -> str: + low = 0 + for edge in DROPPED_LENGTH_EDGES: + if length < edge: + return f"{low}-{edge - 1}" + low = edge + return f"{low}+" + + +def length_histogram(lengths: Iterable[int]) -> dict[str, int]: + """Fixed-bucket histogram; every bucket is present (zeros included).""" + + edges = (0, *DROPPED_LENGTH_EDGES) + out = {f"{lo}-{hi - 1}": 0 for lo, hi in zip(edges, edges[1:])} + out[f"{DROPPED_LENGTH_EDGES[-1]}+"] = 0 + for length in lengths: + out[length_bucket(int(length))] += 1 + return out + + +def offpolicy_gate_event(*, backend: str, counts: Mapping[str, int], + dropped_lengths: Mapping[str, Sequence[int]], + truncation_by_turn: Mapping[int, tuple[int, int]] | None, + granularity: str, **extra: Any) -> dict[str, Any]: + """The ``offpolicy_gate`` fields of one round (design D6). + + ``counts``: reason -> groups. ``dropped_lengths``: reason -> response + lengths of the dropped trajectories. ``truncation_by_turn``: turn index + -> (truncated, total); None when the backend reported nothing. + ``granularity`` says whether the turn counts are tokens or trajectories. + """ + + if granularity not in GRANULARITIES: + raise ValueError(f"granularity must be one of {GRANULARITIES}, got {granularity!r}") + turns = None + if truncation_by_turn is not None: + turns = {str(turn): {"truncated": int(t), "total": int(n), + "ratio": (t / n) if n else None} + for turn, (t, n) in sorted(truncation_by_turn.items())} + return { + "backend": backend, + "counts": {k: int(v) for k, v in sorted(counts.items())}, + "dropped_length_edges": list(DROPPED_LENGTH_EDGES), + "dropped_length_histogram": {reason: length_histogram(lengths) + for reason, lengths in sorted(dropped_lengths.items())}, + "truncation_by_turn": turns, + "granularity": granularity, + **extra, + } + + +def gate_from_carry_over(fields: dict[str, Any], *, kept: int, backend: str + ) -> dict[str, Any] | None: + """Pop the gate keys from a round's carry-over fields and build the + ``offpolicy_gate`` event fields; None when the adapter gave none.""" + + lengths = fields.pop(GATE_DROPPED_LENGTHS, None) + turns = fields.pop(TRUNCATED_BY_TURN, None) + granularity = fields.pop(GATE_GRANULARITY, GRANULARITY_TOKEN) + if lengths is None: + return None + counts = {GATE_KEPT: kept, + GATE_GROUP_DROPPED: int(fields.get("over_age_discarded_groups", 0) or 0), + GATE_UNKNOWN_VERSION_DROPPED: + int(fields.get("unknown_version_discarded_groups", 0) or 0)} + backend_stale = fields.get("stale_dropped_groups") + if isinstance(backend_stale, int): # None = the backend did not expose it + counts[GATE_BACKEND_STALE_DROPPED] = backend_stale + if turns is not None: + turns = {int(k): tuple(v) for k, v in turns.items()} + return offpolicy_gate_event(backend=backend, counts=counts, dropped_lengths=lengths, + truncation_by_turn=turns, granularity=granularity) + + +# -------------------------------------------------------------------------- +# cross_version_is: observe-only, bounds and name from the spec (D7) +# -------------------------------------------------------------------------- + +# Bounds the observe-only estimate used before D7; kept when the spec has none. +OBSERVE_DEFAULT_BOUNDS = (0.0, 2.0) + + +def cross_version_bounds(correction: CorrectionSpec) -> tuple[float, float]: + """``(clip_low, clip_high)`` for :func:`version_segments.cross_version_is`. + + The spec's ``lower`` / ``upper``; :data:`OBSERVE_DEFAULT_BOUNDS` for a + missing bound. ``cross_version_is`` needs ``low <= 1 <= high``, so the + bounds are clamped to that and the estimate stays observe-only. + """ + + low = OBSERVE_DEFAULT_BOUNDS[0] if correction.lower is None else float(correction.lower) + high = OBSERVE_DEFAULT_BOUNDS[1] if correction.upper is None else float(correction.upper) + return min(low, 1.0), max(high, 1.0) + + +def staleness_correction(spec: AlgorithmSpec, default: str = "tis") -> str: + """The ``StalenessPolicy.correction`` name derived from the spec (D7).""" + + from yeto.rl.algos.mismatch_correction import selected_corrections + + named = [m for m in selected_corrections(spec) if m != "mismatch_observe"] + return named[0] if named else default diff --git a/yeto/rl/engine/policy_release.py b/yeto/rl/engine/policy_release.py new file mode 100644 index 000000000..8d4495cb6 --- /dev/null +++ b/yeto/rl/engine/policy_release.py @@ -0,0 +1,159 @@ +"""Policy releases: trainer -> sampler weight hand-off through the store (L4.6, D5). + +A trainer island whose samplers run elsewhere releases every published policy +version into the checkpoint store; a sampler island loads the newest release +into its inference engines. The store is the one both islands already mount for +the sample buffer (local directory, bucket mount or Modal Volume), so the +hand-off needs no new network path. Layout under ``policy-releases/``: + +* ``r-v/``: the weights in the engine's load format, written by + the engine adapter, then ``RELEASE.json`` in it (the release record). ``seq`` + counts the releases of the store (1, 2, ...); +* ``LATEST.json``: a copy of the newest release record, written last. It is the + commit point: a reader never sees a release whose weights are incomplete. + +The record carries the policy version (the engine weight version the sampler +stamps), the version-independent policy hash and the trainer's buffer read +position when it released (``buffer_selected``; the sampler's backlog is the +buffer length minus this position). + +A trainer that publishes the same version and hash again (a restart) releases +nothing. A trainer restored from an older cut publishes an older version with +other weights (a new lineage): that is a new release with a higher ``seq``. + +:class:`PolicyFollower` is the sampler side: it loads a release when its ``seq`` +differs from the one its engines serve. When the trainer stops, no new release +appears and the sampler keeps the last version. +""" + +from __future__ import annotations + +import json +import re +import shutil +import time +from collections.abc import Callable +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +from .resume import CheckpointStore, _write_json_atomic + +SCHEMA = "yeto.rl.policy-release/v1" +RELEASE_DIR = "policy-releases" +LATEST = "LATEST.json" +RECORD = "RELEASE.json" +_RELEASE_NAME = re.compile(r"r(\d+)-v\d+") + + +class PolicyReleaseError(RuntimeError): + pass + + +@dataclass(frozen=True) +class PolicyRelease: + seq: int + version: int + policy_hash: str + path: str # relative to the release root: "r-v" + buffer_selected: int | None = None + released_at: float = 0.0 + schema: str = SCHEMA + + @classmethod + def from_record(cls, record: dict[str, Any]) -> "PolicyRelease": + if record.get("schema") != SCHEMA: + raise PolicyReleaseError(f"release schema {record.get('schema')!r} != {SCHEMA}") + return cls(**{k: record[k] for k in record if k in cls.__dataclass_fields__}) + + +class PolicyReleaseStore: + """Writes (trainer) and reads (sampler) the releases in one store.""" + + def __init__(self, store: CheckpointStore, *, keep: int = 2, + clock: Callable[[], float] = time.time) -> None: + if int(keep) < 1: + raise ValueError("keep must be >= 1") + self.store = store + self.keep = int(keep) + self.clock = clock + self.root = Path(store.root) / RELEASE_DIR + + def resolve(self, release: PolicyRelease) -> Path: + """The release's weight directory as this island mounts the store.""" + return self.root / release.path + + def latest(self) -> PolicyRelease | None: + self.store.reload() + path = self.root / LATEST + if not path.is_file(): + return None + return PolicyRelease.from_record(json.loads(path.read_text(encoding="utf-8"))) + + def release(self, version: int, policy_hash: str, write: Callable[[Path], None], *, + buffer_selected: int | None = None) -> PolicyRelease | None: + """Write the weights with ``write(directory)``, then the record, then LATEST. + + Returns None (writes nothing) when the latest release is this version and hash.""" + version = int(version) + if not policy_hash: + raise ValueError("policy_hash must be non-empty") + latest = self.latest() + if latest is not None and (latest.version, latest.policy_hash) == (version, policy_hash): + return None + seq = 1 if latest is None else latest.seq + 1 + name = f"r{seq}-v{version}" + directory = self.root / name + if directory.exists(): # a partial release of an earlier attempt (no LATEST names it) + shutil.rmtree(directory) + directory.mkdir(parents=True) + write(directory) + record = PolicyRelease(seq=seq, version=version, policy_hash=str(policy_hash), path=name, + buffer_selected=None if buffer_selected is None else int(buffer_selected), + released_at=round(self.clock(), 3)) + _write_json_atomic(directory / RECORD, asdict(record)) + self.store.commit() # the weights and the record are durable before LATEST names them + _write_json_atomic(self.root / LATEST, asdict(record)) + self.store.commit() + self._prune() + return record + + def _prune(self) -> list[str]: + """Delete all but the newest ``keep`` releases (a sampler that is loading + one of them fails the load and retries with the next LATEST).""" + named = sorted((int(m.group(1)), p.name) for p in self.root.iterdir() + if p.is_dir() and (m := _RELEASE_NAME.fullmatch(p.name))) + removed = [name for _, name in named[:-self.keep]] + for name in removed: + shutil.rmtree(self.root / name, ignore_errors=True) + if removed: + self.store.commit() + return removed + + +class PolicyFollower: + """Sampler side: load each new release into the engines. + + ``load(release, directory)`` loads the weights and stamps ``release.version`` + on every engine (raises on any failure; the served release is then unchanged). + ``version`` is the version the engines serve before any release (the base + policy, 0).""" + + def __init__(self, releases: PolicyReleaseStore, + load: Callable[[PolicyRelease, Path], None], *, version: int = 0) -> None: + self.releases = releases + self._load = load + self.version = int(version) + self.release: PolicyRelease | None = None # the release the engines serve + self.last_seen: PolicyRelease | None = None # the newest release read + + def refresh(self) -> PolicyRelease | None: + """The release loaded by this call, or None when there is no new one.""" + latest = self.releases.latest() + self.last_seen = latest + if latest is None or (self.release is not None and latest.seq == self.release.seq): + return None + self._load(latest, self.releases.resolve(latest)) + self.version = latest.version + self.release = latest + return latest diff --git a/yeto/rl/engine/ports.py b/yeto/rl/engine/ports.py index 402830a70..9399d2bcf 100644 --- a/yeto/rl/engine/ports.py +++ b/yeto/rl/engine/ports.py @@ -158,6 +158,10 @@ class RolloutBatchHandle: # (:data:`OVER_SAMPLING_KEYS`: submit batch sizes, refill calls, filtered / # kept / surplus groups, groups in flight at the end). None = not reported. over_sampling: Mapping[str, int] | None = field(default=None, compare=False) + # rl-offpolicy-correction-passthrough L6.11: the engine backend that built + # this batch ("miles", "verl", ...), set by the adapter; it labels the + # driver's ``offpolicy_gate`` event. None = not reported ("unknown"). + backend: str | None = field(default=None, compare=False) def mismatched_groups(self, expected_token: str) -> tuple[GroupMetadata, ...]: return tuple(g for g in self.groups if g.policy_token != expected_token) diff --git a/yeto/rl/engine/recovery.py b/yeto/rl/engine/recovery.py index af6859801..d79cd9114 100644 --- a/yeto/rl/engine/recovery.py +++ b/yeto/rl/engine/recovery.py @@ -35,11 +35,17 @@ # Trigger of each save path: InFlightSaver.save (reclaim, via handle_notice) -> PREEMPT; # InFlightSaver.rehearse and RoundCutCheckpoint.save(final=False) -> PERIODIC; -# RoundCutCheckpoint.save(final=True) (last round or our own stop) -> MANUAL. +# RoundCutCheckpoint.save(final=True) (last round or our own stop) -> MANUAL; +# RoundCutCheckpoint.save(final=True, trigger=SCHEDULED) (a stop the trainer +# scheduler ordered, rl-trainer-start-stop D3) -> SCHEDULED. +# Resume treats every trigger alike: a committed cut is complete. PREEMPT = "preempt" PERIODIC = "periodic" MANUAL = "manual" -TRIGGERS = (PREEMPT, PERIODIC, MANUAL) +SCHEDULED = "scheduled" +TRIGGERS = (PREEMPT, PERIODIC, MANUAL, SCHEDULED) +# Triggers of a cut taken at the end of an incarnation (not a periodic one). +FINAL_TRIGGERS = (PREEMPT, MANUAL, SCHEDULED) @dataclass(frozen=True) @@ -49,7 +55,10 @@ class ConsistentCut: ``outer_step`` is the outer (syncer / published) version, ``island_steps`` the local training step of each island at the cut, ``completed_groups`` the number of prompt groups already consumed and ``completed_groups_digest`` their sorted - sha256 (``resume.consumed_digest``). ``checkpoint_id`` names the payload.""" + sha256 (``resume.consumed_digest``). ``checkpoint_id`` names the payload. + ``buffer_watermark`` is the consumed position of the durable sample buffer + (rl-trainer-start-stop D4; None = the run does not read a sample buffer, and + the record then has no such key).""" checkpoint_id: str trigger: str @@ -57,19 +66,25 @@ class ConsistentCut: island_steps: Mapping[str, int] = field(default_factory=dict) completed_groups: int | None = None completed_groups_digest: str | None = None + buffer_watermark: int | None = None def __post_init__(self) -> None: if self.trigger not in TRIGGERS: raise ValueError(f"cut trigger {self.trigger!r} is not one of {TRIGGERS}") if not self.checkpoint_id: raise ValueError("a consistent cut needs a checkpoint id") + if self.buffer_watermark is not None and int(self.buffer_watermark) < 0: + raise ValueError("buffer watermark must be >= 0") def to_record(self) -> dict[str, Any]: - return {"schema": CUT_SCHEMA, "checkpoint_id": self.checkpoint_id, "trigger": self.trigger, - "outer_step": self.outer_step, - "island_steps": {str(k): int(v) for k, v in sorted(self.island_steps.items())}, - "completed_groups": self.completed_groups, - "completed_groups_digest": self.completed_groups_digest} + record = {"schema": CUT_SCHEMA, "checkpoint_id": self.checkpoint_id, "trigger": self.trigger, + "outer_step": self.outer_step, + "island_steps": {str(k): int(v) for k, v in sorted(self.island_steps.items())}, + "completed_groups": self.completed_groups, + "completed_groups_digest": self.completed_groups_digest} + if self.buffer_watermark is not None: + record["buffer_watermark"] = int(self.buffer_watermark) + return record @classmethod def from_record(cls, record: Mapping[str, Any]) -> ConsistentCut: @@ -79,7 +94,8 @@ def from_record(cls, record: Mapping[str, Any]) -> ConsistentCut: outer_step=record.get("outer_step"), island_steps={str(k): int(v) for k, v in (record.get("island_steps") or {}).items()}, completed_groups=record.get("completed_groups"), - completed_groups_digest=record.get("completed_groups_digest")) + completed_groups_digest=record.get("completed_groups_digest"), + buffer_watermark=record.get("buffer_watermark")) Writer = Callable[[dict[str, Any]], Mapping[str, Any]] diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 438f67aa5..b17f98b81 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -27,6 +27,9 @@ from pathlib import Path from typing import Any +from . import sample_roles +from .sample_roles import SampleBufferConfig + PARAMETER_MODES = frozenset({"lora", "full"}) RECIPE_GENERIC = "generic" RECIPE_QWEN3_5 = "qwen3_5" @@ -392,6 +395,9 @@ class ServingConfig: page_size: Any = None max_running_requests: Any = None chunked_prefill_size: Any = None + # rl-trainer-cold-start-cost L8.2: opt-in SGLang snapshot flags on the + # sampler (rollout) engines; refused without a recorded compatibility pass. + sampler_snapshot: bool = False @dataclass(frozen=True) @@ -456,6 +462,9 @@ class RLRunConfig: # and the load-window length for ``timeline.load_windows`` (None = no hook). edge_costs_path: str | None = None elastic_window_s: float | None = None + # rl-trainer-start-stop L4.6 (--rl-sample-role sampler|trainer): the island + # writes or reads the durable sample buffer. None = role both (unchanged). + sample_buffer: SampleBufferConfig | None = None # -------------------------------------------------------------------------- @@ -950,6 +959,7 @@ def resolve_rl_run_config( page_size=getattr(args, "sglang_page_size", None), max_running_requests=getattr(args, "sglang_max_running_requests", None), chunked_prefill_size=getattr(args, "sglang_chunked_prefill_size", None), + sampler_snapshot=bool(getattr(args, "rl_sampler_snapshot", False)), ), ports=NetworkPorts( rollout_engine_base_port=getattr(args, "rollout_engine_base_port", None), @@ -992,6 +1002,7 @@ def resolve_rl_run_config( edge_costs_path=getattr(args, "rl_edge_costs_path", None) or None, elastic_window_s=(float(args.rl_elastic_window_s) if getattr(args, "rl_elastic_window_s", None) else None), + sample_buffer=sample_roles.resolve(args, groups_per_round=args.groups_per_round), ) diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 55c73d208..872e4b465 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -529,6 +529,15 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "there after every commit point and restores from it when its state dir " "is empty (machine replaced). Without it the state stays on node0's local " "disk (warning on a multi-node island)"), + # rl-identity-tiering L9.4-L9.6 + cli_flag("--rl-code-pack", choices=["off", "git-patch"], default="off", + help="git-patch: pack the local git tree (commit + dirty patch with untracked " + "files) once at the first start, store it in --rl-checkpoint-store and ship it " + "as every island's workdir; resumed and replacement islands get the same pack"), + cli_flag("--rl-code-pack-upgrade", default=None, metavar="REASON", + help="with --rl-code-pack git-patch on an existing run: switch the run to a pack " + "of the current local tree; refused unless the algorithm semantic hash is " + "unchanged; the manifest records the old -> new pack and REASON"), # rl-resume-from-checkpoint (S17 C4) shared_flag("--rl-cut-every", type=int, default=None, help="with --rl-checkpoint-store: keep a round cut every N rounds (plus the " @@ -542,6 +551,24 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. shared_flag("--rl-stop-after-rounds", type=int, default=None, help="with --rl-checkpoint-store and no --rl-elastic: stop after N rounds " "with a final cut (our own stop; the next launch resumes)"), + # rl-trainer-start-stop L4.6: sampler / trainer islands over the durable sample buffer + shared_flag("--rl-sample-role", choices=("both", "sampler", "trainer"), default="both", + help="both (default): the island generates and trains as before. sampler: the " + "island only generates and writes each group to the sample buffer in " + "--rl-checkpoint-store (no trainer model, no trainer GPU) and loads every policy " + "the trainer releases. trainer: the island trains on groups read from the buffer " + "(needs --rl-max-policy-age >= 1) and releases every published policy. Both need " + "--rl-single-island-no-sync and the same --rl-checkpoint-store; full-parameter " + "mode only"), + shared_flag("--rl-sampler-max-backlog-rounds", type=int, default=None, metavar="N", + help="--rl-sample-role sampler: wait while the buffer holds N rounds the trainer " + "has not read (as of its last release); default 4"), + shared_flag("--rl-sampler-min-round-seconds", type=float, default=None, metavar="S", + help="--rl-sample-role sampler: start one round at most every S seconds (a " + "throttle; a check run uses it to make the trainer idle); default no throttle"), + shared_flag("--rl-sample-wait-s", type=float, default=None, metavar="S", + help="--rl-sample-role trainer: max seconds to wait for one full round from " + "the buffer (default 3600)"), ) _SCHEDULING = ( diff --git a/yeto/rl/engine/sample_buffer.py b/yeto/rl/engine/sample_buffer.py new file mode 100644 index 000000000..70955e8a7 --- /dev/null +++ b/yeto/rl/engine/sample_buffer.py @@ -0,0 +1,290 @@ +"""Durable sample buffer: an append-only index next to the sample payloads. + +rl-trainer-start-stop L4.1/L4.3, design D4. The buffer keeps samples across +trainer stops and head restarts. It lives in the existing cut storage (a +:class:`~yeto.rl.engine.resume.CheckpointStore`: local directory or bucket mount) +under ``sample-buffer/``: + +* ``payloads/`` holds one sample payload; it is written (tmp + rename) + before its index line, so the index never names a missing payload; +* ``index.jsonl`` holds one JSON line per sample. The line number is the sample's + position. A torn last line (crash during a write) is ignored by readers and cut + off by the next writer. + +The consumed watermark is the position before which every sample is consumed. +The buffer does not persist it: the trainer writes it into the cut record +(``ConsistentCut.buffer_watermark``), and the next start calls :meth:`reset_to` +with the watermark of the cut it restores. So a sample is consumed at most once +per cut lineage, and samples consumed after the last cut are consumed again. + +Arrival notification is by polling the index (``poll`` / ``start_polling``). + +Limits of this first implementation: one consumer; concurrent appenders must +share one host (``flock`` on the index); the version gate on selection is the +consumer's job (P0-S ``max_policy_age``). +""" + +from __future__ import annotations + +import fcntl +import json +import os +import re +import threading +import time +from collections import Counter +from collections.abc import Callable +from pathlib import Path + +from .resume import CheckpointStore +from .trainer_lifecycle import BacklogStats, SampleIndexEntry, Selected + +SCHEMA = "yeto.rl.sample-buffer/v1" +BUFFER_DIR = "sample-buffer" +INDEX = "index.jsonl" +PAYLOADS = "payloads" +_SAFE_NAME = re.compile(r"[^A-Za-z0-9._-]") + + +class SampleBufferError(RuntimeError): + pass + + +class FileSampleBuffer: + """:class:`~yeto.rl.engine.trainer_lifecycle.DurableSampleBuffer` over a store. + + ``batch_size``: samples per training batch (``BacklogStats.batches``); + ``arrival_window_s``: the window of ``BacklogStats.arrival_rate``.""" + + def __init__(self, store: CheckpointStore, *, batch_size: int, arrival_window_s: float = 60.0, + clock: Callable[[], float] = time.time) -> None: + if batch_size < 1: + raise ValueError("batch_size must be positive") + self.store = store + self.root = Path(store.root) / BUFFER_DIR + (self.root / PAYLOADS).mkdir(parents=True, exist_ok=True) + self.index_path = self.root / INDEX + self.index_path.touch(exist_ok=True) + self.batch_size = int(batch_size) + self.arrival_window_s = float(arrival_window_s) + self.clock = clock + self._lock = threading.RLock() + self._entries: list[SampleIndexEntry] = [] + self._ended = False + self._offset = 0 # bytes of complete index lines read + self._consumer: str | None = None + self._read = 0 # next position select() returns + self._watermark = 0 # positions below are consumed (committed) + self._callbacks: list[Callable[[], None]] = [] + self._polled = 0 # entries seen by the last poll + self._poller: threading.Thread | None = None + self._stop_polling = threading.Event() + self.refresh() + self._polled = len(self._entries) + + # ---------------------------------------------------------------- writing + def write_payload(self, sample_id: str, data: bytes) -> str: + """Store one payload; returns the ``payload_ref`` for its index entry.""" + name = _SAFE_NAME.sub("_", sample_id) + ".bin" + path = self.root / PAYLOADS / name + tmp = path.with_suffix(".bin.tmp") + with open(tmp, "wb") as fh: + fh.write(data) + fh.flush() + os.fsync(fh.fileno()) + os.replace(tmp, path) + return f"{PAYLOADS}/{name}" + + def read_payload(self, payload_ref: str) -> bytes: + return (self.root / payload_ref).read_bytes() + + def append(self, entry: SampleIndexEntry) -> None: + if entry.payload_ref and not (self.root / entry.payload_ref).is_file(): + raise SampleBufferError(f"sample {entry.sample_id}: payload {entry.payload_ref} is missing") + if entry.appended_at == 0.0: + entry = SampleIndexEntry(entry.sample_id, entry.behavior_version, entry.payload_ref, + entry.tokens, self.clock()) + self._write_line(lambda pos: {"schema": SCHEMA, "pos": pos, "sample_id": entry.sample_id, + "behavior_version": int(entry.behavior_version), + "payload_ref": entry.payload_ref, "tokens": int(entry.tokens), + "appended_at": float(entry.appended_at)}) + + def mark_data_end(self) -> None: + """The samplers will append nothing more.""" + self._write_line(lambda pos: {"schema": SCHEMA, "pos": pos, "end": True}) + + def _write_line(self, make: Callable[[int], dict]) -> None: + with self._lock, open(self.index_path, "r+b") as fh: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) + try: + self._read_new(fh) + if self._ended: + raise SampleBufferError("append after the data end") + fh.truncate(self._offset) # drop a torn last line + fh.seek(self._offset) + line = (json.dumps(make(len(self._entries)), sort_keys=True) + "\n").encode() + fh.write(line) + fh.flush() + os.fsync(fh.fileno()) + self._read_new(fh) + finally: + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + self.store.commit() + + # ---------------------------------------------------------------- reading + def refresh(self) -> int: + """Read index lines other writers appended; returns the number of samples.""" + with self._lock: + self.store.reload() + with open(self.index_path, "rb") as fh: + self._read_new(fh) + return len(self._entries) + + def _read_new(self, fh) -> None: + fh.seek(self._offset) + data = fh.read() + while True: + end = data.find(b"\n") + if end < 0: + return # nothing more, or a torn last line + raw, data = data[:end], data[end + 1:] + record = json.loads(raw) + if record.get("schema") != SCHEMA or int(record["pos"]) != len(self._entries): + raise SampleBufferError(f"{self.index_path}: bad line at byte {self._offset}: {raw[:200]!r}") + self._offset += end + 1 + if record.get("end"): + self._ended = True + continue + self._entries.append(SampleIndexEntry( + sample_id=record["sample_id"], behavior_version=int(record["behavior_version"]), + payload_ref=record["payload_ref"], tokens=int(record["tokens"]), + appended_at=float(record["appended_at"]))) + + def __len__(self) -> int: + with self._lock: + return len(self._entries) + + @property + def data_ended(self) -> bool: + with self._lock: + return self._ended + + # ---------------------------------------------------------------- consuming + @property + def watermark(self) -> int: + with self._lock: + return self._watermark + + @property + def read_cursor(self) -> int: + with self._lock: + return self._read + + def reset_to(self, watermark: int) -> None: + """Continue at ``watermark`` (from the restored cut; 0 for a first start): + samples from there are selected again.""" + with self._lock: + self.refresh() + if not 0 <= int(watermark) <= len(self._entries): + raise SampleBufferError(f"watermark {watermark} outside the index (0..{len(self._entries)})") + self._read = self._watermark = int(watermark) + + def select(self, consumer: str, limit: int) -> list[Selected]: + with self._lock: + if self._consumer is None: + self._consumer = consumer + elif consumer != self._consumer: + raise SampleBufferError(f"buffer is read by {self._consumer!r}, not {consumer!r}") + self.refresh() + stop = min(len(self._entries), self._read + max(0, int(limit))) + out = [Selected(pos, self._entries[pos]) for pos in range(self._read, stop)] + self._read = stop + return out + + def commit_consumed(self, cursor: int) -> None: + """Every sample below ``cursor`` is trained; the next cut records it.""" + with self._lock: + if not self._watermark <= int(cursor) <= self._read: + raise SampleBufferError( + f"consumed cursor {cursor} outside {self._watermark}..{self._read} (watermark..selected)") + self._watermark = int(cursor) + + def backlog(self, current_version: int, *, start: int | None = None) -> BacklogStats: + """Samples from ``start`` on; default: the samples not yet selected here. + A head that does not read the buffer passes the trainer's committed + watermark (:class:`CommittedBacklog`).""" + with self._lock: + self.refresh() + pending = self._entries[self._read if start is None else max(0, int(start)):] + versions = Counter(e.behavior_version for e in pending) + oldest = min(versions) if versions else None + since = self.clock() - self.arrival_window_s + recent = sum(1 for e in self._entries if e.appended_at >= since) + return BacklogStats( + batches=len(pending) / self.batch_size, samples=len(pending), oldest_version=oldest, + version_histogram=dict(sorted(versions.items())), + arrival_rate=recent / self.arrival_window_s, data_ended=self._ended, + max_version_gap=None if oldest is None else int(current_version) - oldest) + + # ---------------------------------------------------------------- arrival + def on_arrival(self, callback: Callable[[], None]) -> None: + with self._lock: + self._callbacks.append(callback) + + def poll(self) -> int: + """Read the index once; call every arrival callback when samples arrived + since the last poll. Returns the number of new samples.""" + with self._lock: + new = self.refresh() - self._polled + self._polled += max(new, 0) + callbacks = list(self._callbacks) if new > 0 else [] + for callback in callbacks: + callback() + return max(new, 0) + + def start_polling(self, interval_s: float) -> None: + if interval_s <= 0: + raise ValueError("interval_s must be positive") + if self._poller is not None: + raise SampleBufferError("already polling") + self._stop_polling.clear() + + def loop() -> None: + while not self._stop_polling.wait(interval_s): + self.poll() + + self._poller = threading.Thread(target=loop, name="sample-buffer-poll", daemon=True) + self._poller.start() + + def stop_polling(self) -> None: + self._stop_polling.set() + if self._poller is not None: + self._poller.join() + self._poller = None + + +class CommittedBacklog: + """The head's view of a buffer the trainer reads elsewhere (rl-trainer-start-stop + L5.3c, D4). The backlog starts at the consumed watermark the trainer committed + into its newest cut (``watermark()``, e.g. ``StoreCutLog.watermark``), not at a + cursor of this process. While a trainer runs, the samples it selected after its + newest cut still count as backlog (at most the rounds between two cuts).""" + + def __init__(self, buffer: FileSampleBuffer, watermark: Callable[[], int | None]) -> None: + self.buffer = buffer + self._watermark = watermark + + def backlog(self, current_version: int) -> BacklogStats: + return self.buffer.backlog(current_version, start=self._watermark() or 0) + + def append(self, entry: SampleIndexEntry) -> None: + self.buffer.append(entry) + + def select(self, consumer: str, limit: int) -> list[Selected]: + raise SampleBufferError("the head does not select: the trainer island reads the buffer") + + def commit_consumed(self, cursor: int) -> None: + raise SampleBufferError("the head does not commit: the trainer island writes the watermark") + + def on_arrival(self, callback: Callable[[], None]) -> None: + self.buffer.on_arrival(callback) diff --git a/yeto/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py new file mode 100644 index 000000000..0dd52f08e --- /dev/null +++ b/yeto/rl/engine/sample_roles.py @@ -0,0 +1,211 @@ +"""Island sample roles (rl-trainer-start-stop L4.6). + +An island has one sample role: + +* ``both`` (default): the island generates and trains in one loop, as before; +* ``sampler``: the island only generates and writes every finished group to the + durable sample buffer. It builds no trainer model and takes no trainer GPU. It + loads each policy the trainer releases (:mod:`.policy_release`) and keeps + generating with the last one while the trainer is stopped, up to a backlog cap; +* ``trainer``: the island trains on groups read from the buffer and releases + every published policy for the samplers. + +The sampler and the trainer mount the same checkpoint store (the buffer and the +releases live there). They are either separate single-island launches +(``--rl-sample-role``), or one launch of two islands, a *store-coupled pair* +(``--rl-island-sample-role 0:sampler --rl-island-sample-role 1:trainer``, L5.4): +there is no syncer, and each island runs as its own one-learner, no-sync island +(learner id 0 of 1) whose role comes from its island args (:func:`island_args`). +This module is engine-neutral; the Miles mapping is in +``adapters.miles.buffer_rollout``. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any + +BOTH, SAMPLER, TRAINER = "both", "sampler", "trainer" +SAMPLE_ROLES = (BOTH, SAMPLER, TRAINER) + +# groups a sampler may have ahead of the trainer's last released read position, +# in units of one round (rollout batch), when --rl-sampler-max-backlog-rounds is unset +DEFAULT_MAX_BACKLOG_ROUNDS = 4 + + +@dataclass(frozen=True) +class SampleBufferConfig: + """``role`` is sampler or trainer (``both`` resolves to no config). + + ``store``: the checkpoint store path that holds the buffer and the releases; + ``max_backlog_groups``: sampler admission cap (unconsumed groups); + ``wait_s``: trainer, max seconds to wait for one full round; + ``min_round_s``: sampler, min seconds from one round start to the next (a + throttle, so a test can make the sampler slower than the trainer).""" + + role: str + store: str + max_backlog_groups: int | None = None + wait_s: float | None = None + min_round_s: float | None = None + + +def sample_role(args: Any) -> str: + role = getattr(args, "rl_sample_role", None) or BOTH + if role not in SAMPLE_ROLES: + raise ValueError(f"--rl-sample-role {role!r}: expected one of {SAMPLE_ROLES}") + return role + + +def island_roles(args: Any, num_islands: int) -> dict[int, str]: + """``--rl-island-sample-role ISLAND:ROLE`` entries as {island: role}; {} without + them. Raises ValueError on a malformed, out-of-range or repeated entry.""" + roles: dict[int, str] = {} + for item in getattr(args, "rl_island_sample_role", None) or (): + island, sep, role = str(item).partition(":") + if not sep or not island.strip().isdigit(): + raise ValueError(f"--rl-island-sample-role takes ISLAND:ROLE, got {item!r}") + m = int(island) + if not 0 <= m < num_islands: + raise ValueError(f"--rl-island-sample-role {item!r}: --gpu has {num_islands} island(s)") + if role not in (SAMPLER, TRAINER): + raise ValueError(f"--rl-island-sample-role {item!r}: ROLE must be {SAMPLER} or {TRAINER}") + if m in roles: + raise ValueError(f"--rl-island-sample-role names island {m} twice") + roles[m] = role + return roles + + +def is_pair(args: Any) -> bool: + """Whether the launch is a store-coupled sampler/trainer pair.""" + return bool(getattr(args, "rl_island_sample_role", None)) + + +def island_args(args: Any, island: int, num_islands: int) -> Any: + """``args`` as island ``island`` of a pair sees it: a copy with its + ``rl_sample_role``, keeping only its own role's option (the sampler's backlog + cap, the trainer's wait); ``args`` itself outside a pair.""" + roles = island_roles(args, num_islands) + if not roles: + return args + out = copy.copy(args) + out.rl_sample_role = roles[island] + if out.rl_sample_role == SAMPLER: + out.rl_sample_wait_s = None + else: + out.rl_sampler_max_backlog_rounds = None + out.rl_sampler_min_round_seconds = None + return out + + +def island_learner(args: Any, island: int, num_learners: int) -> tuple[int, int]: + """``(learner id, number of learners)`` the island command carries. A pair + island is a one-learner no-sync island: (0, 1). The launcher keeps ``island`` + for the cluster / Modal function name, the tape and the volumes.""" + return (0, 1) if is_pair(args) else (island, num_learners) + + +def buffer_full_parameter(args: Any) -> bool: + """``--tuning full`` on sampler/trainer islands (single or pair, no-sync): the + island trains every parameter without the full-parameter sync, since a buffer + island has no syncer; the policy reaches the sampler as a store release.""" + return (getattr(args, "tuning", "lora") == "full" + and bool(getattr(args, "rl_single_island_no_sync", False)) + and (is_pair(args) or sample_role(args) != BOTH)) + + +def check_pair(args: Any, num_islands: int) -> None: + """Launch checks of a store-coupled pair (nothing outside a pair): exactly two + islands, one sampler and one trainer, ``--rl-single-island-no-sync`` (no + syncer), no ``--rl-sample-role`` and no external learner; then every island's + own role checks on its island args.""" + roles = island_roles(args, num_islands) + if not roles: + return + if num_islands != 2 or sorted(roles.values()) != [SAMPLER, TRAINER] or len(roles) != 2: + raise ValueError("--rl-island-sample-role: a pair is two --gpu islands, one " + f"{SAMPLER} and one {TRAINER} (got {num_islands} island(s), roles " + f"{dict(sorted(roles.items()))})") + if sample_role(args) != BOTH: + raise ValueError("--rl-island-sample-role replaces --rl-sample-role; give only one") + if int(getattr(args, "external_learners", 0) or 0): + raise ValueError("--rl-island-sample-role does not take external learners") + if not getattr(args, "rl_single_island_no_sync", False): + raise ValueError("--rl-island-sample-role needs --rl-single-island-no-sync: the pair " + "has no syncer; each island is a one-learner no-sync island over the " + "shared --rl-checkpoint-store") + for m in range(num_islands): + check_launch(island_args(args, m, num_islands)) + + +def check_launch(args: Any) -> None: + """Launch-side checks of one island's role, before any cloud work. ``both`` + checks nothing (a pair checks each island with :func:`check_pair`).""" + role = sample_role(args) + if role == BOTH: + for flag, dest in (("--rl-sampler-max-backlog-rounds", "rl_sampler_max_backlog_rounds"), + ("--rl-sampler-min-round-seconds", "rl_sampler_min_round_seconds"), + ("--rl-sample-wait-s", "rl_sample_wait_s")): + if getattr(args, dest, None) is not None: + raise ValueError(f"{flag} needs --rl-sample-role sampler or trainer") + return + if not getattr(args, "rl_single_island_no_sync", False): + raise ValueError(f"--rl-sample-role {role} needs --rl-single-island-no-sync: a sampler " + "and a trainer are separate no-sync islands over one store") + if not getattr(args, "rl_checkpoint_store", None): + raise ValueError(f"--rl-sample-role {role} needs --rl-checkpoint-store (the sample " + "buffer and the policy releases live in it)") + if getattr(args, "rl_elastic", False): + raise ValueError(f"--rl-sample-role {role} does not support --rl-elastic") + if role == TRAINER and not int(getattr(args, "rl_max_policy_age", 0) or 0) >= 1: + raise ValueError("--rl-sample-role trainer trains on samples of earlier versions: " + "it needs --rl-max-policy-age >= 1") + rounds = getattr(args, "rl_sampler_max_backlog_rounds", None) + if rounds is not None and (role != SAMPLER or int(rounds) < 1): + raise ValueError("--rl-sampler-max-backlog-rounds needs --rl-sample-role sampler " + "and a value >= 1") + min_round_s = getattr(args, "rl_sampler_min_round_seconds", None) + if min_round_s is not None and (role != SAMPLER or not float(min_round_s) > 0): + raise ValueError("--rl-sampler-min-round-seconds needs --rl-sample-role sampler " + "and a value > 0") + wait_s = getattr(args, "rl_sample_wait_s", None) + if wait_s is not None and (role != TRAINER or not float(wait_s) > 0): + raise ValueError("--rl-sample-wait-s needs --rl-sample-role trainer and a value > 0") + + +def resolve(args: Any, *, groups_per_round: int) -> SampleBufferConfig | None: + """The island's buffer config (island side: the store is ``--rl-resume-store``).""" + role = sample_role(args) + if role == BOTH: + return None + store = getattr(args, "rl_resume_store", None) + if not store: + raise ValueError(f"--rl-sample-role {role} needs the checkpoint store " + "(--rl-resume-store on the island)") + rounds = getattr(args, "rl_sampler_max_backlog_rounds", None) + wait_s = getattr(args, "rl_sample_wait_s", None) + min_round_s = getattr(args, "rl_sampler_min_round_seconds", None) + return SampleBufferConfig( + role=role, store=str(store), + max_backlog_groups=(int(groups_per_round) * int(rounds or DEFAULT_MAX_BACKLOG_ROUNDS) + if role == SAMPLER else None), + wait_s=None if wait_s is None else float(wait_s), + min_round_s=(float(min_round_s) if role == SAMPLER and min_round_s is not None + else None)) + + +def island_flags(args: Any) -> str: + """The learner flags the launcher forwards (empty for ``both``); ``args`` are + the island args (:func:`island_args`).""" + role = sample_role(args) + if role == BOTH: + return "" + flags = f" --rl-sample-role {role}" + if getattr(args, "rl_sampler_max_backlog_rounds", None) is not None: + flags += f" --rl-sampler-max-backlog-rounds {int(args.rl_sampler_max_backlog_rounds)}" + if getattr(args, "rl_sampler_min_round_seconds", None) is not None: + flags += f" --rl-sampler-min-round-seconds {float(args.rl_sampler_min_round_seconds)!r}" + if getattr(args, "rl_sample_wait_s", None) is not None: + flags += f" --rl-sample-wait-s {float(args.rl_sample_wait_s)!r}" + return flags diff --git a/yeto/rl/engine/sampler_loop.py b/yeto/rl/engine/sampler_loop.py new file mode 100644 index 000000000..56b825184 --- /dev/null +++ b/yeto/rl/engine/sampler_loop.py @@ -0,0 +1,111 @@ +"""The sampler island's loop (rl-trainer-start-stop L4.6). + +A sampler island generates rounds and writes their groups to the durable sample +buffer (the engine's rollout function does the write). Before each round it: + +1. loads the newest policy release, when there is one (:class:`PolicyFollower`); +2. waits while the backlog (groups in the buffer the trainer has not read, as of + its last release) is at the cap. A stopped trainer releases nothing, so the + sampler keeps the last version and stops at the cap; a newer release (the + trainer runs again) lifts the wait. + +After the last round it marks the data end, so the trainer stops after reading +the rest. The loop is engine-neutral: ``generate(rollout_id)`` runs one round. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any + +from .policy_release import PolicyFollower + +DEFAULT_POLL_S = 5.0 + + +class SamplerLoop: + def __init__(self, *, follower: PolicyFollower, buffer: Any, + generate: Callable[[int], None], num_rollouts: int, + max_backlog_groups: int, emit: Callable[..., None], + poll_s: float = DEFAULT_POLL_S, sleep: Callable[[float], None] = time.sleep, + should_stop: Callable[[], bool] = lambda: False, + min_round_s: float | None = None, + clock: Callable[[], float] = time.monotonic) -> None: + if int(max_backlog_groups) < 1: + raise ValueError("max_backlog_groups must be >= 1") + self.follower = follower + self.buffer = buffer + self.generate = generate + self.num_rollouts = int(num_rollouts) + self.max_backlog_groups = int(max_backlog_groups) + self.emit = emit + self.poll_s = float(poll_s) + self.sleep = sleep + self.should_stop = should_stop + self.min_round_s = None if min_round_s is None else float(min_round_s) + self.clock = clock + + def backlog(self) -> int: + """Groups in the buffer after the trainer's read position of its last release + (0 before any release: no trainer has read anything yet).""" + self.buffer.refresh() + seen = self.follower.last_seen + read = 0 if seen is None or seen.buffer_selected is None else int(seen.buffer_selected) + return max(0, len(self.buffer) - read) + + def _load_newest(self) -> None: + loaded = self.follower.refresh() + if loaded is not None: + self.emit("rl_sampler_policy_loaded", policy_version=loaded.version, + policy_hash=loaded.policy_hash, buffer_selected=loaded.buffer_selected) + + def _admit(self, rollout_id: int) -> bool: + """True when the round may start; False when asked to stop while waiting.""" + waiting_since = None + while True: + self._load_newest() + backlog = self.backlog() + if backlog < self.max_backlog_groups: + if waiting_since is not None: + self.emit("rl_sampler_admitted", rollout_id=rollout_id, backlog_groups=backlog, + waited_s=round(time.monotonic() - waiting_since, 3), + policy_version=self.follower.version) + return True + if waiting_since is None: + waiting_since = time.monotonic() + self.emit("rl_sampler_backlog_full", rollout_id=rollout_id, backlog_groups=backlog, + max_backlog_groups=self.max_backlog_groups, + policy_version=self.follower.version) + if self.should_stop(): + return False + self.sleep(self.poll_s) + + def _throttle(self, started: float, *, last: bool) -> None: + """Keep ``min_round_s`` between two round starts (not after the last round).""" + if self.min_round_s is None or last: + return + rest = self.min_round_s - (self.clock() - started) + if rest > 0: + self.emit("rl_sampler_throttle", sleep_s=round(rest, 3)) + self.sleep(rest) + + def run(self, start: int = 0) -> int: + """Generate rounds ``start .. num_rollouts - 1``; returns the rounds generated.""" + done = 0 + for rollout_id in range(int(start), self.num_rollouts): + if self.should_stop() or not self._admit(rollout_id): + self.emit("rl_sampler_stopped", rollout_id=rollout_id, rounds=done) + return done + before = len(self.buffer) + started = self.clock() + self.generate(rollout_id) + self.buffer.refresh() + self.emit("rl_sampler_round", rollout_id=rollout_id, + policy_version=self.follower.version, + groups=len(self.buffer) - before, buffer_end=len(self.buffer)) + done += 1 + self._throttle(started, last=rollout_id == self.num_rollouts - 1) + self.buffer.mark_data_end() + self.emit("rl_sampler_data_end", rounds=done, buffer_end=len(self.buffer)) + return done diff --git a/yeto/rl/engine/start_phases.py b/yeto/rl/engine/start_phases.py new file mode 100644 index 000000000..4031eeb47 --- /dev/null +++ b/yeto/rl/engine/start_phases.py @@ -0,0 +1,250 @@ +"""Trainer start phases (rl-trainer-cold-start-cost L0, design D4). + +A trainer start passes seven fixed phases, in this order: + +``submitted`` the launcher built the island task (launcher clock) +``container_ready`` the island setup script began (container clock) +``setup_done`` the island setup script ended +``ray_ready`` ``ray start --head`` returned on node 0 +``actor_import_done`` every trainer actor process finished its class import +``actor_init_done`` the trainer actors finished ``init`` (model + optimizer) +``first_step_done`` the first optimizer step returned + +Each phase is one ``trainer_start_phase`` event; after ``first_step_done`` +one ``trainer_start_summary`` event carries ``total_s`` = +``first_step_done - submitted`` (``T_s`` of the scheduler). + +Times are Unix wall-clock seconds: the phases come from three processes on +two machines (launcher, setup shell, island), so a monotonic clock cannot +order them. Clock skew between the launcher and the container goes into +``submitted -> container_ready`` only. + +The first three container phases are written by the shell (:func:`shell_mark`) +into :data:`MARK_DIR`; the island replays them (:meth:`StartPhases.replay`). +Stdlib only: the launcher, the island and the head all import this module. +""" + +from __future__ import annotations + +import json +import os +import shlex +import time +from collections.abc import Callable, Iterable, Mapping +from pathlib import Path +from typing import Any + +START_PHASES = ( + "submitted", + "container_ready", + "setup_done", + "ray_ready", + "actor_import_done", + "actor_init_done", + "first_step_done", +) +START_PHASE_EVENT = "trainer_start_phase" +START_SUMMARY_EVENT = "trainer_start_summary" +# Launcher -> island: the Unix time the island task was built. +SUBMITTED_ENV = "YETO_START_SUBMITTED_UNIX" +# Setup/run shell -> island: one file per phase, content = Unix time. +MARK_DIR = "$HOME/.yeto-start-phases" +SHELL_PHASES = ("container_ready", "setup_done", "ray_ready") +# Setup shell -> island: events the setup shell could not write to the tape +# (no tape exists before the island starts), one JSON object per line. +SETUP_EVENTS_FILE = "setup-events.jsonl" +# Run manifest key (L0.2): {learner_id: summary}. +MANIFEST_KEY = "trainer_start" + +Emit = Callable[..., None] + + +def shell_mark(phase: str) -> str: + """One shell line that records ``phase`` now; it never fails the script.""" + + if phase not in SHELL_PHASES: + raise ValueError(f"{phase!r} is not a shell-recorded start phase: {SHELL_PHASES}") + return ( + f'{{ mkdir -p "{MARK_DIR}" && date +%s.%N > "{MARK_DIR}/{phase}"; }} ' + ">/dev/null 2>&1 || true" + ) + + +def shell_event(event: str, fields: Mapping[str, Any] | None = None, + int_fields: Mapping[str, str] | None = None) -> str: + """Shell lines that append one setup event and print it to stderr. + + ``fields`` are fixed values (JSON-encoded now); ``int_fields`` map a field + to a shell expression that expands to an integer at run time (``"$rc"``). + The lines never fail the script. + """ + + static = json.dumps({"event": event, **(fields or {})}, sort_keys=True) + body, args = static[:-1].replace("%", "%%"), "" + for name, expr in (int_fields or {}).items(): + body += f", {json.dumps(name)}: %d" + args += f' "{expr}"' + line = shlex.quote(body + "}\\n") + return ( + f'{{ mkdir -p "{MARK_DIR}" && printf {line}{args} >> "{MARK_DIR}/{SETUP_EVENTS_FILE}"; }} ' + ">/dev/null 2>&1 || true\n" + f"printf {shlex.quote('[yeto-setup] event ' + body + '}' + chr(92) + 'n')}{args} >&2 || true" + ) + + +def shell_reset() -> str: + """Clear marks of an earlier start in the same home (relaunch on one machine).""" + + return f'rm -rf "{MARK_DIR}" >/dev/null 2>&1 || true' + + +class StartPhases: + """Records the phases in order and emits the events. + + ``emit(event, **fields)`` writes one event. A phase out of order, or a + phase given twice, is an error: the order is the contract the scheduler + reads. A phase that could not be measured is skipped with + :meth:`missing`; the summary names it. + """ + + def __init__(self, emit: Emit, *, clock: Callable[[], float] = time.time) -> None: + self._emit = emit + self._clock = clock + self.times: dict[str, float] = {} + self.missing_phases: list[str] = [] + self.summary: dict[str, Any] | None = None + self._setup_events_seen = 0 + + def _next_index(self) -> int: + done = [*self.times, *self.missing_phases] + return max((START_PHASES.index(p) for p in done), default=-1) + 1 + + def _check_next(self, phase: str) -> None: + if phase not in START_PHASES: + raise ValueError(f"unknown start phase {phase!r}; known: {START_PHASES}") + if phase in self.times or phase in self.missing_phases: + raise ValueError(f"start phase {phase!r} recorded twice") + if START_PHASES.index(phase) < self._next_index(): + raise ValueError(f"start phase {phase!r} out of order after {list(self.times)}") + + def mark(self, phase: str, at: float | None = None) -> None: + self._check_next(phase) + t = self._clock() if at is None else float(at) + submitted = self.times.get("submitted") + self.times[phase] = t + self._emit( + START_PHASE_EVENT, + phase=phase, + index=START_PHASES.index(phase), + t_unix=t, + since_submitted_s=None if submitted is None else round(t - submitted, 3), + ) + if phase == START_PHASES[-1]: + self._summarize() + + def missing(self, phase: str, reason: str) -> None: + self._check_next(phase) + self.missing_phases.append(phase) + self._emit(START_PHASE_EVENT, phase=phase, index=START_PHASES.index(phase), + t_unix=None, missing=reason) + + def mark_once(self, phase: str) -> None: + """``mark`` unless already recorded (first train step of a long loop).""" + + if phase not in self.times and phase not in self.missing_phases: + self.mark(phase) + + def replay(self, environ: Mapping[str, str] | None = None, + mark_dir: str | os.PathLike[str] | None = None) -> None: + """Record ``submitted`` and the shell phases from env and mark files.""" + + environ = os.environ if environ is None else environ + directory = Path(os.path.expandvars(MARK_DIR) if mark_dir is None else mark_dir) + submitted = environ.get(SUBMITTED_ENV) + if submitted: + self.mark("submitted", float(submitted)) + else: + self.missing("submitted", f"{SUBMITTED_ENV} not set") + for phase in SHELL_PHASES: + try: + self.mark(phase, float((directory / phase).read_text().strip())) + except (OSError, ValueError) as error: + self.missing(phase, f"no mark file: {error}") + self.forward_setup_events(directory) + + def forward_setup_events(self, mark_dir: str | os.PathLike[str] | None = None) -> int: + """Emit setup-shell events not yet forwarded; returns how many were new. + + Called at replay and again later: a background step of the setup (the + model prefetch) can fail after the island has started. + """ + + directory = Path(os.path.expandvars(MARK_DIR) if mark_dir is None else mark_dir) + records = read_tape_records([directory / SETUP_EVENTS_FILE]) + new = records[self._setup_events_seen:] + self._setup_events_seen = len(records) + for record in new: + fields = {k: v for k, v in record.items() if k != "event"} + self._emit(str(record.get("event", "setup_event")), source="setup", **fields) + return len(new) + + def _summarize(self) -> None: + first, last = START_PHASES[0], START_PHASES[-1] + total = (round(self.times[last] - self.times[first], 3) + if first in self.times and last in self.times else None) + self.summary = { + "phases": {p: self.times.get(p) for p in START_PHASES}, + "missing": list(self.missing_phases), + "total_s": total, + } + self._emit(START_SUMMARY_EVENT, **self.summary) + + +def summaries_from_events(records: Iterable[Mapping[str, Any]]) -> dict[str, dict[str, Any]]: + """``{learner_id: summary}`` from tape records (last summary per island wins).""" + + out: dict[str, dict[str, Any]] = {} + for record in records: + if record.get("event") != START_SUMMARY_EVENT: + continue + learner = str(record.get("island_id", record.get("learner_id", "?"))) + out[learner] = {k: record.get(k) for k in ("phases", "missing", "total_s")} + return out + + +def read_tape_records(paths: Iterable[os.PathLike[str] | str]) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for path in paths: + try: + lines = Path(path).read_text(encoding="utf-8").splitlines() + except OSError: + continue + for line in lines: + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + records.append(record) + return records + + +def write_manifest_summary(manifest_path: os.PathLike[str] | str, + summaries: Mapping[str, Mapping[str, Any]]) -> bool: + """Merge ``summaries`` into the run manifest under :data:`MANIFEST_KEY`.""" + + if not summaries: + return False + path = Path(manifest_path) + manifest = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} + manifest.setdefault(MANIFEST_KEY, {}).update({k: dict(v) for k, v in summaries.items()}) + path.write_text(json.dumps(manifest, indent=1, sort_keys=True) + "\n", encoding="utf-8") + return True + + +def t_s_from_manifest(manifest: Mapping[str, Any]) -> float | None: + """``T_s`` of a finished run: the slowest island start (``None`` if unknown).""" + + totals = [s.get("total_s") for s in (manifest.get(MANIFEST_KEY) or {}).values()] + totals = [float(t) for t in totals if t is not None] + return max(totals) if totals else None diff --git a/yeto/rl/engine/sync_timing.py b/yeto/rl/engine/sync_timing.py new file mode 100644 index 000000000..eaa4ad9b9 --- /dev/null +++ b/yeto/rl/engine/sync_timing.py @@ -0,0 +1,46 @@ +"""Island-side timing of one sync boundary (rl-trainer-start-stop D10, L5.6). + +The boundary is split into three parts so a run can tell transfer from wait: + +* ``submit_ts`` - Unix time when the island finished handing its update to + the transport. +* ``global_ready_ts`` - Unix time when the next global policy was fully + received on the island. +* ``recv_seconds`` - wall seconds to receive that global payload, from its + frame header to its last byte. ``None`` when the transport does not expose + it; the wait then includes the receive time. + +``sync_wait_seconds`` is the rest: the time the island waited for the syncer +(quorum, merge, other islands) before the global payload started to arrive. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SyncTiming: + submit_ts: float | None + global_ready_ts: float | None + recv_seconds: float | None = None + + @property + def sync_wait_seconds(self) -> float | None: + if self.submit_ts is None or self.global_ready_ts is None: + return None + # A base that was already there before the submit means no wait. + return max(0.0, self.global_ready_ts - self.submit_ts - (self.recv_seconds or 0.0)) + + def fields(self) -> dict[str, float | None]: + return { + "sync/submit_ts": self.submit_ts, + "sync/global_ready_ts": self.global_ready_ts, + "sync/recv_seconds": self.recv_seconds, + "sync/sync_wait_seconds": self.sync_wait_seconds, + } + + +def unix_now() -> float: + return time.time() diff --git a/yeto/rl/engine/trainer_lifecycle.py b/yeto/rl/engine/trainer_lifecycle.py new file mode 100644 index 000000000..7ed96ae4f --- /dev/null +++ b/yeto/rl/engine/trainer_lifecycle.py @@ -0,0 +1,562 @@ +"""Trainer start/stop: the lifecycle, the durable sample buffer, the scheduler. + +rl-trainer-start-stop design D1/D2. Mechanism and policy are separate: + +* :class:`TrainerLifecycle` does what it is told (start, run, drain and cut, stop); +* :class:`TrainerScheduler` decides ``Start`` / ``Stop`` / ``Hold`` from backlog + statistics, the lifecycle state and prices (a pure function of its inputs); +* :class:`DurableSampleBuffer` holds the samples across trainer stops. + +A trainer that starts once is the special case: a scheduler that holds until the +data ends, then stops with ``FINISHED`` (:class:`AlwaysOnScheduler`). Scheduler S1 +(start on backlog, ski-rental stop) is in :mod:`.trainer_scheduler`. + +:class:`DriverLifecycle` is the lifecycle over the existing :class:`IslandDriver` +(D2). Each start builds a fresh driver (new trainer ports, new sync session) for a +new incarnation, so ``IslandDriver.rebuild_trainer`` (a same-driver rebuild) is +never needed across starts. The driver runs on its own thread; the lifecycle +wraps its sync session to observe the states and to hold the driver after a +publication until it is told to train on or to drain: + +* JOIN as a new incarnation before restoring (``Membership.join``); +* restore the last cut through the sync session's own resume path (round cut -> + ``CuttableTrainer.restore_cut``); +* the first publication after the restore is the warm-up; the trainer is then + ``TRAINING`` and waits for :meth:`DriverLifecycle.run`; +* a drain ends the driver where it waits: the final cut takes the + trigger of the stop reason (``STOPPED`` -> ``SCHEDULED``) and LEAVE carries the + P5 reason code (``Membership.leave``). + +This module is neutral core: it knows ports and the LEAVE code, no backend. +""" + +from __future__ import annotations + +import enum +import math +import threading +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +from yeto.rl.elastic_client import LeaveReason + +from . import recovery +from .leave_reason import leave_reason_for + +if TYPE_CHECKING: + from .driver import IslandDriver + + +class LifecycleError(RuntimeError): + """A lifecycle verb was called in a state that does not allow it, or a start + did not continue from the cut it was given.""" + + +class LifecycleState(str, enum.Enum): + COLD = "cold" + ACQUIRING = "acquiring" + JOINING = "joining" + RESTORING = "restoring" + WARMING = "warming" + TRAINING = "training" + DRAINING = "draining" + STOPPED = "stopped" + + +_ORDER = tuple(LifecycleState) + + +def allowed_transition(old: LifecycleState, new: LifecycleState) -> bool: + """The states move only forward one step, and from any state to STOPPED.""" + if new is LifecycleState.STOPPED: + return old is not LifecycleState.STOPPED + return _ORDER.index(new) == _ORDER.index(old) + 1 + + +# Trigger of the final cut for each stop reason (D3): a scheduler stop is SCHEDULED, +# a spot reclaim PREEMPT, the end of the data (or a user stop) MANUAL. +CUT_TRIGGER_FOR_REASON: Mapping[LeaveReason, str] = { + LeaveReason.STOPPED: recovery.SCHEDULED, + LeaveReason.PREEMPTED: recovery.PREEMPT, + LeaveReason.FINISHED: recovery.MANUAL, +} + + +# --------------------------------------------------------------------------- +# Neutral records +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class CutRef: + """A committed cut a later start can restore. ``policy_version`` is the + published version at the cut (the next start publishes it again first); + ``watermark`` is the consumed buffer watermark in the same record (D4).""" + + cut_id: str + trigger: str + policy_version: int + watermark: int | None = None + + def __post_init__(self) -> None: + if self.trigger not in recovery.TRIGGERS: + raise ValueError(f"cut trigger {self.trigger!r} is not one of {recovery.TRIGGERS}") + + +@dataclass(frozen=True) +class StartPlan: + incarnation: int + spec: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class StopCondition: + """``rounds``: train this many more rounds, then wait at the round boundary for + the next verb. None: train until the data ends.""" + + rounds: int | None = None + + def __post_init__(self) -> None: + if self.rounds is not None and self.rounds < 1: + raise ValueError("StopCondition.rounds must be positive") + + +@dataclass(frozen=True) +class RunReceipt: + rounds: int # rounds trained by this run call + first_version: int # trained version before the first of these rounds + trained_version: int # trained version after the last of these rounds + data_ended: bool # the data ended: the driver drained on its own (FINISHED) + + +@dataclass(frozen=True) +class SampleIndexEntry: + """One sample in the durable buffer: its id, the policy version that produced it + (behavior version, P0-S schema 4) and where its payload is. One sample per + entry; the cross-island group index is ``sample_pool.SampleIndexEntry``.""" + + sample_id: str + behavior_version: int + payload_ref: str + tokens: int = 0 + appended_at: float = 0.0 + + +@dataclass(frozen=True) +class Selected: + """A selected sample and its buffer position (the cursor to commit).""" + + position: int + entry: SampleIndexEntry + + +@dataclass(frozen=True) +class BacklogStats: + """``batches``: unconsumed samples in units of one training batch; + ``version_histogram``: behavior version -> unconsumed samples; + ``arrival_rate``: samples per second over the recent window; + ``data_ended``: the samplers will append nothing more; + ``max_version_gap``: current trainer version minus ``oldest_version``.""" + + batches: float + samples: int + oldest_version: int | None = None + version_histogram: Mapping[int, int] = field(default_factory=dict) + arrival_rate: float = 0.0 + data_ended: bool = False + max_version_gap: int | None = None + + +# --------------------------------------------------------------------------- +# Scheduler decisions +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class Start: + pass + + +@dataclass(frozen=True) +class Hold: + pass + + +@dataclass(frozen=True) +class Stop: + reason: LeaveReason = LeaveReason.STOPPED + + +Decision = Start | Hold | Stop + + +# --------------------------------------------------------------------------- +# Ports +# --------------------------------------------------------------------------- +@runtime_checkable +class DurableSampleBuffer(Protocol): + def append(self, entry: SampleIndexEntry) -> None: ... + def backlog(self, current_version: int) -> BacklogStats: ... + def select(self, consumer: str, limit: int) -> list[Selected]: ... + def commit_consumed(self, cursor: int) -> None: ... # written into the cut + def on_arrival(self, callback: Callable[[], None]) -> None: ... + + +@runtime_checkable +class TrainerScheduler(Protocol): + def decide(self, backlog: BacklogStats, lifecycle_state: LifecycleState, + prices: Mapping[str, float] | None = None) -> Decision: ... + + +@runtime_checkable +class Membership(Protocol): + """JOIN / LEAVE toward the coordinator (D2: the elastic client's ``Join`` with a + new incarnation per start, ``Leave`` with the P5 reason code). ``join`` returns + the coordinator's base version, or None when there is no coordinator.""" + + def join(self, incarnation: int) -> int | None: ... + def leave(self, incarnation: int, reason: LeaveReason) -> None: ... + + +@runtime_checkable +class CutLog(Protocol): + """Where the lifecycle reads the newest committed cut back.""" + + def latest(self) -> CutRef | None: ... + + +@runtime_checkable +class TrainerLifecycle(Protocol): + def plan_start(self, spec: Mapping[str, Any]) -> StartPlan: ... + def start(self, plan: StartPlan, *, resume_from: CutRef | None) -> TrainerHandle: ... + def ready(self, handle: TrainerHandle, deadline: float) -> bool: ... + def run(self, handle: TrainerHandle, source: DurableSampleBuffer | None, + until: StopCondition) -> RunReceipt: ... + def drain_and_cut(self, handle: TrainerHandle, reason: LeaveReason) -> CutRef | None: ... + def stop(self, handle: TrainerHandle, reason: LeaveReason) -> None: ... + + +class AlwaysOnScheduler: + """"Start once" (D1): start while the trainer is cold, hold while it trains, + stop with ``FINISHED`` when the data ended and the backlog is empty.""" + + def decide(self, backlog: BacklogStats, lifecycle_state: LifecycleState, + prices: Mapping[str, float] | None = None) -> Decision: + del prices # S1 and this scheduler ignore prices + if lifecycle_state is LifecycleState.COLD: + return Start() + if backlog.data_ended and backlog.samples == 0: + return Stop(LeaveReason.FINISHED) + return Hold() + + +# --------------------------------------------------------------------------- +# Handle +# --------------------------------------------------------------------------- +class TrainerHandle: + """One incarnation of the trainer. Thread-safe: the driver thread moves the + state; the caller's thread gives commands and waits on ``_cond``.""" + + def __init__(self, plan: StartPlan, *, resume_from: CutRef | None, + clock: Callable[[], float]) -> None: + self.plan = plan + self.incarnation = plan.incarnation + self.resume_from = resume_from + self.state = LifecycleState.COLD + self.history: list[tuple[str, float]] = [(self.state.value, clock())] + self.driver: IslandDriver | None = None + self.source: DurableSampleBuffer | None = None + self.start_version: int | None = None # version restored (or 0) and published first + self.published_version: int | None = None + self.rounds = 0 # rounds trained by this incarnation + self.data_ended = False + self.joined = False + self.requested_reason: LeaveReason | None = None + self.leave_reason: LeaveReason | None = None + self.cut: CutRef | None = None + self.error: BaseException | None = None + self._clock = clock + self._cond = threading.Condition() + self._target_rounds = 0 + self._drain_trigger: str | None = None + self._parks = 0 # times the driver waited for a command + self._parked = False + self._done = False # the driver thread ended + self._thread: threading.Thread | None = None + + # -- state --------------------------------------------------------------- + def advance(self, new: LifecycleState) -> None: + with self._cond: + if not allowed_transition(self.state, new): + raise LifecycleError(f"incarnation {self.incarnation}: {self.state.value} -> {new.value}") + self.state = new + self.history.append((new.value, self._clock())) + self._cond.notify_all() + + @property + def states(self) -> list[str]: + return [s for s, _ in self.history] + + @property + def trained_version(self) -> int | None: + return None if self.start_version is None else self.start_version + self.rounds + + # -- driver side --------------------------------------------------------- + def park(self) -> str | None: + """Wait at a round boundary for a command. Returns the cut trigger of a + drain, or None to train on.""" + with self._cond: + self._parked = True + self._parks += 1 + self._cond.notify_all() + self._cond.wait_for(lambda: self._drain_trigger is not None + or self.rounds < self._target_rounds) + self._parked = False + return self._drain_trigger + + def wants_park(self) -> bool: + with self._cond: + return self.rounds >= self._target_rounds or self._drain_trigger is not None + + +# --------------------------------------------------------------------------- +# Sync wrapper: the lifecycle's view into the driver +# --------------------------------------------------------------------------- +class _Drained(Exception): + """Ends the driver after a drain: the cut is written, nothing more to train.""" + + +class _LifecycleSync: + """Wraps the driver's sync session; everything else is delegated. + + The driver waits for the next verb right after each publication (the restored + policy, then every trained version), so the samplers already hold the newest + version while the trainer waits (D5). A drain there finishes the inner session + (final cut, with the trigger of the stop reason) and ends the driver.""" + + def __init__(self, inner: Any, handle: TrainerHandle, membership: Membership | None) -> None: + self.inner = inner + self.handle = handle + self.membership = membership + self.closed_reason: LeaveReason | None = None + self.drained = False + + def __getattr__(self, name: str) -> Any: + return getattr(self.inner, name) + + def start(self, driver: IslandDriver) -> Any: + h = self.handle + h.advance(LifecycleState.JOINING) + if self.membership is not None: + self.membership.join(h.incarnation) + h.joined = True + h.advance(LifecycleState.RESTORING) + start = self.inner.start(driver) + if h.resume_from is not None and int(start.rollout_id) != h.resume_from.policy_version: + raise LifecycleError( + f"incarnation {h.incarnation} restored version {start.rollout_id}, " + f"the cut {h.resume_from.cut_id} is at version {h.resume_from.policy_version}") + h.start_version = int(start.rollout_id) + h.advance(LifecycleState.WARMING) + return start + + def published(self, driver: IslandDriver, *, rollout_id: int, policy_hash: str) -> None: + self.inner.published(driver, rollout_id=rollout_id, policy_hash=policy_hash) + h = self.handle + h.published_version = int(rollout_id) + if h.state is LifecycleState.WARMING: + # the restored policy is on the samplers: warm-up done + h.advance(LifecycleState.TRAINING) + if h.state is not LifecycleState.TRAINING or not h.wants_park(): + return + trigger = h.park() + if trigger is None: + return + h.advance(LifecycleState.DRAINING) + if hasattr(self.inner, "final_trigger"): + self.inner.final_trigger = trigger + self.inner.finish(driver) # the final cut at the published version + self.drained = True + raise _Drained() + + def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Any: + out = self.inner.boundary(driver, rollout_id=rollout_id, stats=stats) + h = self.handle + commit = getattr(driver.rollout, "commit_round", None) + if callable(commit): # the round is trained: its samples are consumed (D4) + commit(rollout_id) + with h._cond: + h.rounds += 1 + if out.stop: # the data ended: drain on our own + h.data_ended = True + h.advance(LifecycleState.DRAINING) + return out + + def close(self, reason: int = LeaveReason.FINISHED) -> None: + """The driver's departure code (P5) unless the driver ended cleanly or by + our own stop: then the reason the lifecycle was given.""" + code = LeaveReason(reason) + if (code is LeaveReason.FINISHED or self.drained) and self.handle.requested_reason is not None: + code = self.handle.requested_reason + self.closed_reason = code + self.inner.close(code) + + +# --------------------------------------------------------------------------- +# Lifecycle over IslandDriver +# --------------------------------------------------------------------------- +DriverFactory = Callable[[TrainerHandle], "IslandDriver"] + + +class DriverLifecycle: + """:class:`TrainerLifecycle` over :class:`IslandDriver` (see the module doc). + + ``build_driver(handle)`` acquires the trainer for one incarnation and returns + a fresh driver whose sync session restores ``handle.resume_from`` in + ``start`` (e.g. ``LocalOnlySync`` with round cuts). ``release(handle)`` gives + the cards back after LEAVE. ``cuts`` reads the cut a drain committed.""" + + def __init__(self, build_driver: DriverFactory, *, membership: Membership | None = None, + cuts: CutLog | None = None, release: Callable[[TrainerHandle], None] | None = None, + first_incarnation: int = 0, clock: Callable[[], float] = time.monotonic) -> None: + self.build_driver = build_driver + self.membership = membership + self.cuts = cuts + self.release = release + self.clock = clock + self._next_incarnation = int(first_incarnation) + self._lock = threading.Lock() + + def plan_start(self, spec: Mapping[str, Any]) -> StartPlan: + with self._lock: + incarnation = self._next_incarnation + self._next_incarnation += 1 + return StartPlan(incarnation=incarnation, spec=dict(spec)) + + def start(self, plan: StartPlan, *, resume_from: CutRef | None) -> TrainerHandle: + handle = TrainerHandle(plan, resume_from=resume_from, clock=self.clock) + handle.advance(LifecycleState.ACQUIRING) + try: + driver = self.build_driver(handle) + except BaseException as error: + handle.error = error + self._end(handle, leave_reason_for(error)) + raise + sync = _LifecycleSync(driver.sync, handle, self.membership) + driver.sync = sync + handle.driver = driver + thread = threading.Thread(target=self._drive, args=(handle, sync), + name=f"trainer-incarnation-{plan.incarnation}", daemon=True) + handle._thread = thread + thread.start() + return handle + + def ready(self, handle: TrainerHandle, deadline: float) -> bool: + """Wait until the trainer can train (True) or failed (False), at most until + ``deadline`` on the lifecycle clock.""" + with handle._cond: + while not (handle._parked and handle.state is LifecycleState.TRAINING) \ + and handle.state is not LifecycleState.STOPPED: + left = deadline - self.clock() + if left <= 0: + return False + handle._cond.wait(timeout=min(left, 1.0)) + return handle.state is LifecycleState.TRAINING + + def run(self, handle: TrainerHandle, source: DurableSampleBuffer | None, + until: StopCondition) -> RunReceipt: + """Train ``until.rounds`` more rounds (or to the end of the data); return at + the round boundary with the trainer still in ``TRAINING``, or after it + drained on its own because the data ended.""" + handle.source = source + with handle._cond: + self._wait_parked_or_done(handle) + self._raise_if_done(handle) + first = handle.rounds + parks = handle._parks + handle._target_rounds = math.inf if until.rounds is None else first + until.rounds + handle._cond.notify_all() + handle._cond.wait_for(lambda: handle._done or (handle._parked and handle._parks > parks)) + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + assert handle.start_version is not None + return RunReceipt(rounds=handle.rounds - first, first_version=handle.start_version + first, + trained_version=handle.start_version + handle.rounds, + data_ended=handle.data_ended) + + def drain_and_cut(self, handle: TrainerHandle, reason: LeaveReason) -> CutRef | None: + """End the driver where it waits (after a publication): final cut with the + trigger of ``reason``, close the sync session. Returns the newest cut (the + restored one when nothing was trained since; None without a cut log).""" + reason = LeaveReason(reason) + trigger = CUT_TRIGGER_FOR_REASON.get(reason) + if trigger is None: + raise LifecycleError(f"no cut trigger for LEAVE reason {reason.name}") + with handle._cond: + if not handle._done: + self._wait_parked_or_done(handle) + if not handle._done: + handle.requested_reason = reason + handle._drain_trigger = trigger + handle._cond.notify_all() + self._join(handle) + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + return handle.cut + + def stop(self, handle: TrainerHandle, reason: LeaveReason) -> None: + """LEAVE with ``reason`` and release the trainer. Drains first when the + trainer still waits at a round boundary.""" + reason = LeaveReason(reason) + if handle.state is LifecycleState.STOPPED: + return + if not handle._done: + self.drain_and_cut(handle, reason) + if handle.state is LifecycleState.STOPPED: + return + self._end(handle, handle.requested_reason or reason) + + # -- internals ------------------------------------------------------------- + def _drive(self, handle: TrainerHandle, sync: _LifecycleSync) -> None: + assert handle.driver is not None + try: + handle.driver.run() + except _Drained: + pass + except BaseException as error: # noqa: BLE001 - reported on the handle + handle.error = error + if handle.error is None: + if self.cuts is not None: + handle.cut = self.cuts.latest() + with handle._cond: + handle._done = True + handle._cond.notify_all() + if handle.error is not None: + # from any state to STOPPED on failure; LEAVE with the driver's code + self._end(handle, sync.closed_reason or leave_reason_for(handle.error)) + + def _end(self, handle: TrainerHandle, reason: LeaveReason) -> None: + handle.leave_reason = LeaveReason(reason) + try: + if handle.joined and self.membership is not None: + self.membership.leave(handle.incarnation, handle.leave_reason) + finally: + try: + if self.release is not None: + self.release(handle) + finally: + handle.advance(LifecycleState.STOPPED) + + @staticmethod + def _wait_parked_or_done(handle: TrainerHandle) -> None: + handle._cond.wait_for(lambda: handle._parked or handle._done) + + @staticmethod + def _raise_if_done(handle: TrainerHandle) -> None: + if handle._done: + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + raise LifecycleError(f"incarnation {handle.incarnation} already drained") + + @staticmethod + def _join(handle: TrainerHandle) -> None: + if handle._thread is not None: + handle._thread.join() diff --git a/yeto/rl/engine/trainer_loop.py b/yeto/rl/engine/trainer_loop.py new file mode 100644 index 000000000..e4e200c3a --- /dev/null +++ b/yeto/rl/engine/trainer_loop.py @@ -0,0 +1,338 @@ +"""The controller loop of the trainer scheduler (rl-trainer-start-stop L5.3b, D1/D8). + +:class:`TrainerLoop` connects the policy to the mechanism. Each :meth:`TrainerLoop.tick`: + +1. reads the inputs: the backlog of the :class:`~.trainer_lifecycle.DurableSampleBuffer` + at the current trained version, the data-end signal, and the lifecycle state. + ``T_s`` (the start-phase events) is read by the scheduler itself; +2. asks the :class:`~.trainer_lifecycle.TrainerScheduler` for a decision; +3. executes it on the :class:`~.trainer_lifecycle.TrainerLifecycle`: + + * ``Start``: plan a new incarnation, start it from the newest cut + (``CutLog.latest``; the lifecycle JOINs and restores it), wait until ready; + * ``Stop``: ``lifecycle.stop`` (drain, cut with the trigger of the reason -- + ``STOPPED`` -> ``SCHEDULED`` --, LEAVE with the P5 code, release the cards); + * ``Hold`` while ``TRAINING`` with at least one batch: train ``rounds_per_tick`` + rounds and return at the round boundary. + +A ``Stop(FINISHED)`` with no trainer running (the data ended with less than one +batch left) ends the loop too; it is journaled as ``finished`` (no incarnation). + +After ``max_consecutive_failures`` failed incarnations in a row (no trained round +between them) the loop releases the trainer, journals ``gave_up`` and raises +:class:`TrainerLoopGaveUp`; a trained round resets the count. + +The loop journals each start and stop (kind ``trainer_loop``) in the controller +journal next to the scheduler state (kind ``trainer_scheduler_state``). After a head +restart, :meth:`TrainerLoop.restore` reads both back: the scheduler continues with +its idle time, the next incarnation number follows the last one, and a trainer the +journal shows as started and not stopped is an orphan of the dead head. The +in-process handle of an orphan is gone, so the loop releases it through +``release_orphan`` (when given) and journals it as orphaned. The next start restores +the newest committed cut. The caller resets the head's buffer to the watermark of +that cut after :meth:`TrainerLoop.restore` (the buffer is not this module's). + +:class:`TrainerLoopThread` runs ticks at a fixed period on its own thread, next to +the fleet controller's poll loop (``yeto.launch``). + +Neutral core: no backend, no cloud. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import replace +from typing import Any + +from yeto.rl.elastic_client import LeaveReason + +from .trainer_lifecycle import ( + BacklogStats, + CutLog, + Decision, + DurableSampleBuffer, + Hold, + LifecycleError, + LifecycleState, + Start, + Stop, + StopCondition, + TrainerHandle, + TrainerLifecycle, + TrainerScheduler, +) + +LOOP_KIND = "trainer_loop" + +# trainer_loop records: ``event`` is one of these. +STARTED = "started" +READY = "ready" +STOPPED = "stopped" +ORPHANED = "orphaned" +FINISHED = "finished" # the data ended with no trainer running (incarnation None) +GAVE_UP = "gave_up" # too many incarnations failed in a row: the loop stopped + +DEFAULT_MAX_CONSECUTIVE_FAILURES = 3 + + +class TrainerLoopGaveUp(RuntimeError): + """``max_consecutive_failures`` incarnations failed in a row with no trained + round between them; the loop stopped starting new ones.""" + +Record = Callable[..., Any] # journal.append(kind, **fields) + + +def next_incarnation(records: Iterable[Mapping[str, Any]]) -> int: + """The incarnation after the last one the journal names (0 for a new run). + Build the lifecycle with it (``DriverLifecycle(first_incarnation=...)``).""" + last = -1 + for r in records: + if r.get("kind") == LOOP_KIND and r.get("incarnation") is not None: + last = max(last, int(r["incarnation"])) + return last + 1 + + +class TrainerLoop: + """One trainer, one scheduler, one buffer. Not thread-safe: one caller ticks. + + ``record`` / ``records``: the controller journal (``Journal.append`` and + ``Journal.records``). ``data_ended`` tells the samplers will append nothing more. + ``release_orphan(incarnation)`` gives back the cards of a trainer the dead head + started (cloud layer).""" + + def __init__(self, *, scheduler: TrainerScheduler, lifecycle: TrainerLifecycle, + buffer: DurableSampleBuffer, cuts: CutLog, + record: Record | None = None, + data_ended: Callable[[], bool] = lambda: False, + release_orphan: Callable[[int], None] | None = None, + spec: Mapping[str, Any] | None = None, + rounds_per_tick: int = 1, ready_timeout_s: float = 1800.0, + max_consecutive_failures: int = DEFAULT_MAX_CONSECUTIVE_FAILURES, + clock: Callable[[], float] = time.monotonic) -> None: + if rounds_per_tick < 1: + raise ValueError("rounds_per_tick must be positive") + if max_consecutive_failures < 1: + raise ValueError("max_consecutive_failures must be positive") + self.scheduler = scheduler + self.lifecycle = lifecycle + self.buffer = buffer + self.cuts = cuts + self._record = record + self._data_ended = data_ended + self._release_orphan = release_orphan + self.spec = dict(spec or {}) + self.rounds_per_tick = int(rounds_per_tick) + self.ready_timeout_s = float(ready_timeout_s) + self._clock = clock + self.handle: TrainerHandle | None = None + self.started = 0 # incarnations started by this loop or a head before it + self.finished = False # stopped with FINISHED: nothing more to do + self._journaled_stops: set[int] = set() # incarnations with a stopped record + # failed incarnations since the last trained round; at the limit the loop gives up + self.max_consecutive_failures = int(max_consecutive_failures) + self.consecutive_failures = 0 + + # ------------------------------------------------------------ state + @property + def state(self) -> LifecycleState: + if self.handle is not None: + return self.handle.state + return LifecycleState.STOPPED if self.started else LifecycleState.COLD + + def current_version(self) -> int: + """The trained version: the running trainer's, else the newest cut's, else 0.""" + if self.handle is not None and self.handle.trained_version is not None: + return self.handle.trained_version + cut = self.cuts.latest() + return 0 if cut is None else cut.policy_version + + def backlog(self) -> BacklogStats: + stats = self.buffer.backlog(self.current_version()) + return replace(stats, data_ended=stats.data_ended or bool(self._data_ended())) + + # ------------------------------------------------------------ tick + def tick(self) -> Decision: + """Read the inputs, decide, execute. Returns the decision.""" + if self.finished: + return Hold() + self._note_failure() + self._give_up_at_limit() + backlog = self.backlog() + decision = self.scheduler.decide(backlog, self.state) + if isinstance(decision, Start): + self._start() + elif isinstance(decision, Stop): + self._stop(decision.reason) + elif self.state is LifecycleState.TRAINING and backlog.batches >= 1: + self._train() + self._give_up_at_limit() + return decision + + def _give_up_at_limit(self) -> None: + """Stop the loop after ``max_consecutive_failures`` failed incarnations in a + row: release a trainer still up, journal it and raise (no new Start).""" + if self.consecutive_failures < self.max_consecutive_failures: + return + self.close(LeaveReason.INFRA_FAILED) + self.finished = True + self._journal(GAVE_UP, None if self.handle is None else self.handle.incarnation, + failures=self.consecutive_failures, limit=self.max_consecutive_failures) + raise TrainerLoopGaveUp( + f"{self.consecutive_failures} trainer incarnations failed in a row " + f"(limit {self.max_consecutive_failures}); the trainer loop stopped") + + def close(self, reason: LeaveReason = LeaveReason.STOPPED) -> None: + """Stop a running trainer (run end or head shutdown).""" + if self.handle is not None and self.handle.state is not LifecycleState.STOPPED: + self._stop(reason) + + # ------------------------------------------------------------ verbs + def _start(self) -> None: + resume = self.cuts.latest() + plan = self.lifecycle.plan_start(self.spec) + # journal first: a head that dies during the start sees the orphan + self._journal(STARTED, plan.incarnation, + resume_cut=None if resume is None else resume.cut_id, + resume_version=None if resume is None else resume.policy_version) + self.started += 1 + try: + self.handle = self.lifecycle.start(plan, resume_from=resume) + except Exception as error: # noqa: BLE001 - the lifecycle stopped it with its P5 code + self.handle = None + self.consecutive_failures += 1 + self._journal(STOPPED, plan.incarnation, reason=None, error=repr(error)) + return + if self.lifecycle.ready(self.handle, self._clock() + self.ready_timeout_s): + self._journal(READY, plan.incarnation, version=self.handle.start_version) + return + if self.handle.state is not LifecycleState.STOPPED: + try: + self.lifecycle.stop(self.handle, LeaveReason.INFRA_FAILED) # not ready in time + except LifecycleError: + pass # no cut for a failure: the lifecycle released it; journaled below + self._note_failure() + + def _stop(self, reason: LeaveReason) -> None: + handle = self.handle + if handle is None or handle.state is LifecycleState.STOPPED: + if reason is LeaveReason.FINISHED: # the data ended, nothing left to train + self.finished = True + self._journal(FINISHED, None, remainder=self.buffer.backlog(self.current_version()).samples) + return + try: + self.lifecycle.stop(handle, reason) + except LifecycleError: + pass # the handle carries the error; journaled below + self._journal_stop(handle) + if handle.error is None and handle.leave_reason is LeaveReason.FINISHED: + self.finished = True + + def _train(self) -> None: + assert self.handle is not None + try: + receipt = self.lifecycle.run(self.handle, self.buffer, StopCondition(self.rounds_per_tick)) + except LifecycleError: + self._note_failure() + return + if receipt.rounds > 0: + self.consecutive_failures = 0 # a trained round: the trainer works + if receipt.data_ended: + self._stop(LeaveReason.FINISHED) + + def _note_failure(self) -> None: + """Journal a trainer that stopped on its own (failure) once.""" + handle = self.handle + if handle is not None and handle.state is LifecycleState.STOPPED \ + and handle.incarnation not in self._journaled_stops: + self._journal_stop(handle) + + def _journal_stop(self, handle: TrainerHandle) -> None: + self._journaled_stops.add(handle.incarnation) + if handle.error is not None or handle.leave_reason is LeaveReason.INFRA_FAILED: + self.consecutive_failures += 1 + cut = handle.cut + reason = handle.leave_reason + self._journal(STOPPED, handle.incarnation, + reason=None if reason is None else reason.name, + cut=None if cut is None else cut.cut_id, + trigger=None if cut is None else cut.trigger, + version=handle.trained_version, + released=getattr(handle, "released", None), + error=None if handle.error is None else repr(handle.error)) + + # ------------------------------------------------------------ journal (D8) + def _journal(self, event: str, incarnation: int | None, **fields: Any) -> None: + if self._record is not None: + self._record(LOOP_KIND, event=event, + incarnation=None if incarnation is None else int(incarnation), **fields) + + def restore(self, records: Iterable[Mapping[str, Any]]) -> bool: + """Continue after a head restart. Returns False for a new run (no records).""" + records = list(records) + restored = self.scheduler.restore(records) if hasattr(self.scheduler, "restore") else False + open_incarnation = None + for r in records: + if r.get("kind") != LOOP_KIND: + continue + if r["event"] == FINISHED: + self.finished = True + elif r["event"] == STARTED: + self.started += 1 + open_incarnation = int(r["incarnation"]) + elif r["event"] in (STOPPED, ORPHANED): + if r["event"] == STOPPED and r.get("reason") == LeaveReason.FINISHED.name: + self.finished = True + open_incarnation = None + if open_incarnation is not None: + if self._release_orphan is not None: + self._release_orphan(open_incarnation) + self._journal(ORPHANED, open_incarnation) + return bool(restored or self.started) + + +class TrainerLoopThread: + """Ticks a :class:`TrainerLoop` every ``interval_s`` on a daemon thread until + :meth:`stop` or the loop finished. A tick error ends the thread; ``error`` + holds it.""" + + def __init__(self, loop: TrainerLoop, interval_s: float, + on_error: Callable[[BaseException], None] | None = None) -> None: + if interval_s <= 0: + raise ValueError("interval_s must be positive") + self.loop = loop + self.interval_s = float(interval_s) + self._on_error = on_error # told once, on this thread, after a tick error + self.error: BaseException | None = None + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name="trainer-scheduler-loop", daemon=True) + + def start(self) -> TrainerLoopThread: + self._thread.start() + return self + + def stop(self, *, close_reason: LeaveReason | None = LeaveReason.STOPPED, + timeout: float | None = None) -> None: + """End the ticks; then stop a running trainer with ``close_reason`` + (None keeps it, e.g. for a head that hands over).""" + self._stop.set() + self._thread.join(timeout) + if close_reason is not None and not self._thread.is_alive(): + self.loop.close(close_reason) + + def join(self, timeout: float | None = None) -> bool: + """Wait up to ``timeout`` seconds for the loop to end on its own (finished + or a tick error); True when it ended.""" + self._thread.join(timeout) + return not self._thread.is_alive() + + def _run(self) -> None: + try: + while not self.loop.finished and not self._stop.is_set(): + self.loop.tick() + self._stop.wait(self.interval_s) + except BaseException as error: # noqa: BLE001 - reported to the launcher + self.error = error + if self._on_error is not None: + self._on_error(error) diff --git a/yeto/rl/engine/trainer_remote.py b/yeto/rl/engine/trainer_remote.py new file mode 100644 index 000000000..8226e994d --- /dev/null +++ b/yeto/rl/engine/trainer_remote.py @@ -0,0 +1,568 @@ +"""The trainer lifecycle across a process boundary (rl-trainer-start-stop L5.3c). + +The head runs the scheduler loop; the trainer runs on its own island (a Modal +function, a sky cluster). Each incarnation is a fresh island. Two channels connect +them, both existing paths: + +* commands, head -> island: a JSON-lines file per incarnation in the run's + checkpoint store (``trainer-control/inc-.jsonl``). The island already mounts + that store for its cuts and the sample buffer; the head writes it through the + store's own client (Modal Volume upload, bucket, local directory); +* events, island -> head: ``rl_trainer_lifecycle`` records in the island's event + tape, which the island already echoes on its log stream (``yeto.rl.event_echo``). + +Why not an elastic control frame: the elastic protocol has no head -> island +command frame (the syncer sends JoinAck, Base, Finished, Verdict), the trainer of a +start/stop run restores round cuts without a syncer, and a new frame changes the +Rust syncer and the protocol version (the v3 boundary of L9.7). + +Island side: :class:`IslandTrainerAgent` runs the island's driver inside a +:class:`~.trainer_lifecycle.DriverLifecycle` (JOIN through the sync session, +restore of the newest cut, warm-up publication) and executes the commands at round +boundaries: train N rounds, or drain with the cut trigger of the stop reason and +LEAVE with the P5 code. It reports each step as an event. + +Head side: :class:`RemoteTrainerLifecycle` is the +:class:`~.trainer_lifecycle.TrainerLifecycle` of the scheduler loop. It acquires an +island per incarnation through a :class:`TrainerIslands` port (cloud layer), sends +the commands, follows the events, and releases the island after its LEAVE. + +:class:`StoreCutLog` reads the newest round-cut pointer from the store (snapshot +or flat layout); its watermark is the consumed position the trainer committed. + +Neutral core: no backend, no cloud SDK. +""" + +from __future__ import annotations + +import json +import math +import os +import time +from collections.abc import Callable, Mapping +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +from yeto.rl.elastic_client import LeaveReason + +from . import recovery +from .resume import FLAT_MANIFEST, LATEST, STATE_DIR, CheckpointStore +from .trainer_lifecycle import ( + CUT_TRIGGER_FOR_REASON, + CutLog, + CutRef, + DriverLifecycle, + LifecycleError, + LifecycleState, + RunReceipt, + StartPlan, + StopCondition, + TrainerHandle, +) + +CONTROL_DIR = "trainer-control" +EVENT = "rl_trainer_lifecycle" +POINTER = "round-cut.json" # adapters.miles.round_cut.POINTER (state dir, synced to the store) + +# Island environment, set per incarnation by the cloud layer. +CONTROL_ENV = "YETO_TRAINER_CONTROL" # "1": the island waits for head commands +INCARNATION_ENV = "YETO_TRAINER_INCARNATION" +RESUME_ENV = "YETO_TRAINER_RESUME" # JSON of the CutRef to restore; empty for a first start +POLL_ENV = "YETO_TRAINER_COMMAND_POLL_S" + +DEFAULT_POLL_S = 5.0 +# A log stream may deliver the island's last lines after the cloud reports the call +# ended: wait this long for the stopped event before calling the exit a failure. +DEFAULT_EXIT_GRACE_S = 60.0 + +RUN, STOP = "run", "stop" + +# Event phases (field ``phase``). +STARTED, READY, TRAINED, STOPPED = "started", "ready", "trained", "stopped" + + +class RemoteTrainerError(RuntimeError): + """The trainer island failed or ended without a stopped event.""" + + +# --------------------------------------------------------------------------- +# Store files +# --------------------------------------------------------------------------- +@runtime_checkable +class StoreFiles(Protocol): + """Small whole-file reads and writes in the run's checkpoint store, by path + relative to the store root. ``read`` returns None for a missing file and sees + what other hosts committed; ``write`` is visible to them when it returns.""" + + def read(self, rel: str) -> bytes | None: ... + def write(self, rel: str, data: bytes) -> None: ... + + +class LocalStoreFiles: + """:class:`StoreFiles` over a mounted :class:`CheckpointStore` (the island's + mount, a bucket mount, a local directory).""" + + def __init__(self, store: CheckpointStore) -> None: + self.store = store + + def read(self, rel: str) -> bytes | None: + self.store.reload() + path = Path(self.store.root) / rel + return path.read_bytes() if path.is_file() else None + + def write(self, rel: str, data: bytes) -> None: + path = Path(self.store.root) / rel + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "wb") as fh: + fh.write(data) + fh.flush() + os.fsync(fh.fileno()) + os.replace(tmp, path) + self.store.commit() + + +# --------------------------------------------------------------------------- +# Cuts in the store +# --------------------------------------------------------------------------- +def cut_ref_to_json(cut: CutRef | None) -> str: + return "" if cut is None else json.dumps(asdict(cut), sort_keys=True) + + +def cut_ref_from_json(text: str | None) -> CutRef | None: + return CutRef(**json.loads(text)) if text else None + + +def cut_ref_of_pointer(pointer: Mapping[str, Any]) -> CutRef: + """The cut a round-cut pointer names. ``policy_version`` is the version the next + start publishes first: the pointer's next rollout id.""" + cut = recovery.ConsistentCut.from_record(pointer[recovery.CUT_KEY]) + version = pointer.get("next_rollout_id", pointer.get("policy_version")) + return CutRef(cut_id=cut.checkpoint_id, trigger=cut.trigger, policy_version=int(version), + watermark=cut.buffer_watermark) + + +class StoreCutLog: + """:class:`~.trainer_lifecycle.CutLog` over the store: the round-cut pointer of + the newest snapshot (``LATEST`` -> ``state//round-cut.json``) or, in the + flat layout, the pointer in the store root.""" + + def __init__(self, files: StoreFiles) -> None: + self.files = files + + def pointer(self) -> dict[str, Any] | None: + latest = self.files.read(LATEST) + if latest is not None: + seq = int(json.loads(latest)["seq"]) + raw = self.files.read(f"{STATE_DIR}/{seq:08d}/{POINTER}") + elif self.files.read(FLAT_MANIFEST) is not None: + raw = self.files.read(POINTER) + else: + raw = None + return None if raw is None else dict(json.loads(raw)) + + def latest(self) -> CutRef | None: + pointer = self.pointer() + if pointer is None or recovery.CUT_KEY not in pointer: + return None + return cut_ref_of_pointer(pointer) + + def watermark(self) -> int: + """The consumed buffer position of the newest cut (0 without a cut).""" + cut = self.latest() + return 0 if cut is None or cut.watermark is None else int(cut.watermark) + + +# --------------------------------------------------------------------------- +# Command channel +# --------------------------------------------------------------------------- +def command_path(incarnation: int) -> str: + return f"{CONTROL_DIR}/inc-{int(incarnation):06d}.jsonl" + + +@dataclass(frozen=True) +class Command: + """``run``: train ``rounds`` more rounds (None: to the end of the data); + ``stop``: drain, cut with the trigger of ``reason`` and LEAVE with it.""" + + seq: int + verb: str + rounds: int | None = None + reason: str | None = None + + def __post_init__(self) -> None: + if self.verb not in (RUN, STOP): + raise ValueError(f"unknown trainer command {self.verb!r}") + if self.verb == STOP and self.reason not in LeaveReason.__members__: + raise ValueError(f"stop command needs a LEAVE reason name, got {self.reason!r}") + + +class CommandWriter: + """Head side: appends commands to one incarnation's file (rewritten whole: the + store has no append).""" + + def __init__(self, files: StoreFiles, incarnation: int) -> None: + self.files = files + self.path = command_path(incarnation) + self.sent: list[Command] = [] + + def send(self, verb: str, *, rounds: int | None = None, reason: LeaveReason | None = None) -> int: + command = Command(seq=len(self.sent) + 1, verb=verb, rounds=rounds, + reason=None if reason is None else LeaveReason(reason).name) + lines = [json.dumps(asdict(c), sort_keys=True) for c in (*self.sent, command)] + self.files.write(self.path, ("\n".join(lines) + "\n").encode()) + self.sent.append(command) + return command.seq + + +class CommandReader: + """Island side: the commands of one incarnation not seen yet, in order.""" + + def __init__(self, files: StoreFiles, incarnation: int) -> None: + self.files = files + self.path = command_path(incarnation) + self.seen = 0 + + def poll(self) -> list[Command]: + raw = self.files.read(self.path) + if raw is None: + return [] + out = [] + for line in raw.decode().splitlines(): + if not line.strip(): + continue + command = Command(**json.loads(line)) + if command.seq > self.seen: + out.append(command) + self.seen = command.seq + return out + + +# --------------------------------------------------------------------------- +# Island side +# --------------------------------------------------------------------------- +def commanded(environ: Mapping[str, str] | None = None) -> bool: + return (os.environ if environ is None else environ).get(CONTROL_ENV) == "1" + + +def lifecycle_event(handle: TrainerHandle, phase: str, *, seq: int, + watermark: int | None = None) -> dict[str, Any]: + """The fields of one ``rl_trainer_lifecycle`` event. ``watermark``: the + consumed buffer position after the trained rounds (committed on the island).""" + cut = handle.cut + return { + "phase": phase, "incarnation": handle.incarnation, "seq": int(seq), "watermark": watermark, + "state": handle.state.value, "start_version": handle.start_version, + "rounds": handle.rounds, "trained_version": handle.trained_version, + "data_ended": handle.data_ended, + "leave_reason": None if handle.leave_reason is None else handle.leave_reason.name, + "cut": None if cut is None else asdict(cut), + "error": None if handle.error is None else repr(handle.error)[:2000], + } + + +class IslandTrainerAgent: + """Runs a built island driver as one commanded incarnation (see the module doc). + + ``emit(event, **fields)`` writes the island's event tape (``driver.emit``); + ``cuts`` reads back the cut a drain committed (``StoreCutLog`` over the island's + store mount); ``watermark(version)`` commits and returns the consumed buffer + position below ``version`` (``RoundCutCheckpoint.buffer_watermark``), so the head + counts the backlog between two cuts.""" + + def __init__(self, driver: Any, *, commands: CommandReader, incarnation: int, + resume_from: CutRef | None, emit: Callable[..., Any], cuts: CutLog | None = None, + watermark: Callable[[int], int | None] | None = None, + poll_s: float = DEFAULT_POLL_S, sleep: Callable[[float], None] = time.sleep) -> None: + self.lifecycle = DriverLifecycle(lambda _handle: driver, cuts=cuts, + first_incarnation=int(incarnation)) + self.commands = commands + self.resume_from = resume_from + self.emit = emit + self.watermark = watermark + self.poll_s = float(poll_s) + self.sleep = sleep + self.seq = 0 # last command applied + self.handle: TrainerHandle | None = None + + def run(self) -> TrainerHandle: + lc = self.lifecycle + handle = self.handle = lc.start(lc.plan_start({}), resume_from=self.resume_from) + self._event(STARTED) + if lc.ready(handle, math.inf): + self._event(READY) + while handle.state is not LifecycleState.STOPPED: + commands = self.commands.poll() + if not commands: + self.sleep(self.poll_s) + for command in commands: + self._apply(command) + if handle.state is LifecycleState.STOPPED: + break + self._event(STOPPED) + if handle.error is not None: + raise handle.error + return handle + + def _apply(self, command: Command) -> None: + lc, handle = self.lifecycle, self.handle + assert handle is not None + self.seq = command.seq + try: + if command.verb == RUN: + receipt = lc.run(handle, None, StopCondition(command.rounds)) + self._event(TRAINED) + if receipt.data_ended: + lc.stop(handle, LeaveReason.FINISHED) + else: + lc.stop(handle, LeaveReason[command.reason]) + except LifecycleError: + pass # the handle is STOPPED and carries the error: reported by run() + + def _event(self, phase: str) -> None: + handle = self.handle + assert handle is not None + watermark = None + if phase in (READY, TRAINED) and self.watermark is not None and handle.trained_version is not None: + watermark = self.watermark(handle.trained_version) + self.emit(EVENT, **lifecycle_event(handle, phase, seq=self.seq, + watermark=None if watermark is None else int(watermark))) + + +def run_commanded_island(driver: Any, *, store: CheckpointStore, + environ: Mapping[str, str] | None = None) -> TrainerHandle: + """Island entry: the incarnation and the cut to restore come from the + environment the cloud layer set for this incarnation.""" + environ = os.environ if environ is None else environ + files = LocalStoreFiles(store) + incarnation = int(environ[INCARNATION_ENV]) + watermark = getattr(getattr(driver.sync, "round_cuts", None), "buffer_watermark", None) + return IslandTrainerAgent( + driver, commands=CommandReader(files, incarnation), incarnation=incarnation, + resume_from=cut_ref_from_json(environ.get(RESUME_ENV)), emit=driver.emit, + cuts=StoreCutLog(files), watermark=watermark if callable(watermark) else None, + poll_s=float(environ.get(POLL_ENV) or DEFAULT_POLL_S)).run() + + +# --------------------------------------------------------------------------- +# Head side +# --------------------------------------------------------------------------- +@runtime_checkable +class TrainerIslands(Protocol): + """Cloud layer port: one trainer island per incarnation. + + ``launch`` starts the island with the incarnation environment (``island_env``) + and returns its job; ``events`` returns the island events that arrived since + the last call (non-blocking); ``alive`` is False once the island ended; + ``release`` gives the cards back and returns True when the cloud confirmed it; + ``release_incarnation`` does the same for an incarnation a dead head started.""" + + def launch(self, incarnation: int, resume_from: CutRef | None) -> Any: ... + def events(self, job: Any) -> list[Mapping[str, Any]]: ... + def alive(self, job: Any) -> bool: ... + def release(self, job: Any) -> bool: ... + def release_incarnation(self, incarnation: int) -> bool: ... + + +def island_env(incarnation: int, resume_from: CutRef | None, *, + poll_s: float | None = None) -> dict[str, str]: + env = {CONTROL_ENV: "1", INCARNATION_ENV: str(int(incarnation)), + RESUME_ENV: cut_ref_to_json(resume_from)} + if poll_s is not None: + env[POLL_ENV] = str(float(poll_s)) + return env + + +class RemoteHandle(TrainerHandle): + """A :class:`TrainerHandle` the head moves from the island's events.""" + + def __init__(self, plan: StartPlan, *, resume_from: CutRef | None, clock: Callable[[], float], + commands: CommandWriter) -> None: + super().__init__(plan, resume_from=resume_from, clock=clock) + self.commands = commands + self.job: Any = None + self.acked = 0 # last command the island reported done + self.ready_seen = False + self.watermark: int | None = None # consumed position the island reported + self.island_ended = False # stopped event seen, or the island exited + self.released: bool | None = None # cloud confirmed the cards are back + self.exited_at: float | None = None # first poll that saw the island not alive + + +class RemoteTrainerLifecycle: + """:class:`~.trainer_lifecycle.TrainerLifecycle` over a trainer island per + incarnation (see the module doc). Waits poll the events every ``poll_s``; an + island that ended without a stopped event is a failure after ``exit_grace_s``.""" + + def __init__(self, islands: TrainerIslands, files: StoreFiles, *, first_incarnation: int = 0, + poll_s: float = DEFAULT_POLL_S, exit_grace_s: float = DEFAULT_EXIT_GRACE_S, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep) -> None: + self.islands = islands + self.files = files + self.poll_s = float(poll_s) + self.exit_grace_s = float(exit_grace_s) + self.clock = clock + self.sleep = sleep + self._next_incarnation = int(first_incarnation) + self.current: RemoteHandle | None = None + + def running_watermark(self) -> int | None: + """The consumed position the running incarnation reported (None when no + trainer runs: then the newest cut holds it).""" + handle = self.current + if handle is None or handle.state is LifecycleState.STOPPED: + return None + return handle.watermark + + def plan_start(self, spec: Mapping[str, Any]) -> StartPlan: + plan = StartPlan(incarnation=self._next_incarnation, spec=dict(spec)) + self._next_incarnation += 1 + return plan + + def start(self, plan: StartPlan, *, resume_from: CutRef | None) -> RemoteHandle: + handle = RemoteHandle(plan, resume_from=resume_from, clock=self.clock, + commands=CommandWriter(self.files, plan.incarnation)) + self.current = handle + handle.advance(LifecycleState.ACQUIRING) + try: + handle.job = self.islands.launch(plan.incarnation, resume_from) + except BaseException as error: + handle.error = error + handle.leave_reason = LeaveReason.INFRA_FAILED + handle.advance(LifecycleState.STOPPED) # nothing acquired, nothing to release + raise + return handle + + def ready(self, handle: RemoteHandle, deadline: float) -> bool: + self._wait(handle, lambda: handle.ready_seen, deadline) + return handle.ready_seen and handle.state is LifecycleState.TRAINING + + def run(self, handle: RemoteHandle, source: Any, until: StopCondition) -> RunReceipt: + """Train ``until.rounds`` more rounds on the island (``source`` is the island's + own buffer reader; the head's view is not sent).""" + del source + self._raise_if_ended(handle) + first = handle.rounds + seq = handle.commands.send(RUN, rounds=until.rounds) + self._wait(handle, lambda: handle.acked >= seq, math.inf) + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + assert handle.start_version is not None + return RunReceipt(rounds=handle.rounds - first, first_version=handle.start_version + first, + trained_version=handle.start_version + handle.rounds, + data_ended=handle.data_ended) + + def drain_and_cut(self, handle: RemoteHandle, reason: LeaveReason) -> CutRef | None: + reason = LeaveReason(reason) + if reason not in CUT_TRIGGER_FOR_REASON: + raise LifecycleError(f"no cut trigger for LEAVE reason {reason.name}") + if not handle.island_ended: + handle.requested_reason = reason + handle.commands.send(STOP, reason=reason) + self._wait(handle, lambda: False, math.inf) + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + return handle.cut + + def stop(self, handle: RemoteHandle, reason: LeaveReason) -> None: + """Drain (when the island still runs), then release the island.""" + reason = LeaveReason(reason) + if handle.state is LifecycleState.STOPPED: + return + try: + if not handle.island_ended: + self.drain_and_cut(handle, reason) + finally: + if handle.state is not LifecycleState.STOPPED: + self._end(handle, reason) + + @staticmethod + def _raise_if_ended(handle: RemoteHandle) -> None: + if handle.island_ended or handle.state is LifecycleState.STOPPED: + if handle.error is not None: + raise LifecycleError(f"incarnation {handle.incarnation} failed") from handle.error + raise LifecycleError(f"incarnation {handle.incarnation} already drained") + + # -- events ---------------------------------------------------------------- + def _wait(self, handle: RemoteHandle, done: Callable[[], bool], deadline: float) -> None: + """Follow the events until ``done()``, the island ended, or ``deadline``.""" + while True: + self._pump(handle) + if done() or handle.island_ended or self.clock() >= deadline: + return + self.sleep(self.poll_s) + + def _pump(self, handle: RemoteHandle) -> None: + events = [e for e in self.islands.events(handle.job) + if e.get("event") == EVENT and int(e.get("incarnation", -1)) == handle.incarnation] + for event in events: + self._apply(handle, event) + if not handle.island_ended and not events and not self.islands.alive(handle.job): + if handle.exited_at is None: + handle.exited_at = self.clock() + if self.clock() - handle.exited_at >= self.exit_grace_s: + handle.island_ended = True + handle.error = handle.error or RemoteTrainerError( + f"trainer island of incarnation {handle.incarnation} ended without a " + "stopped event") + if handle.island_ended and handle.state is not LifecycleState.STOPPED: + self._end(handle, _first_reason(handle.leave_reason, handle.requested_reason, + LeaveReason.INFRA_FAILED)) + + @staticmethod + def _apply(handle: RemoteHandle, event: Mapping[str, Any]) -> None: + for key in ("start_version", "rounds"): + if event.get(key) is not None: + setattr(handle, key, int(event[key])) + handle.data_ended = bool(event.get("data_ended")) + if event.get("watermark") is not None: + handle.watermark = int(event["watermark"]) + handle.acked = max(handle.acked, int(event.get("seq") or 0)) + _catch_up(handle, LifecycleState(event["state"])) + phase = event.get("phase") + if phase == READY: + handle.ready_seen = True + elif phase == STOPPED: + handle.island_ended = True + if event.get("cut"): + handle.cut = CutRef(**event["cut"]) + if event.get("leave_reason"): + handle.leave_reason = LeaveReason[event["leave_reason"]] + if event.get("error"): + handle.error = RemoteTrainerError(event["error"]) + + def _end(self, handle: RemoteHandle, reason: LeaveReason) -> None: + """The island left (its own LEAVE): give the cards back.""" + handle.leave_reason = _first_reason(handle.leave_reason, LeaveReason(reason)) + try: + handle.released = bool(self.islands.release(handle.job)) + finally: + handle.advance(LifecycleState.STOPPED) + if not handle.released: + print(f"[trainer] incarnation {handle.incarnation}: card release not confirmed by the " + "cloud layer", flush=True) + + +def _first_reason(*reasons: LeaveReason | None) -> LeaveReason: + """The first given reason (FINISHED is 0: never test a reason for truth).""" + return next(LeaveReason(r) for r in reasons if r is not None) + + +def consumed_watermark(cuts: StoreCutLog, lifecycle: RemoteTrainerLifecycle | None) -> int: + """Where the head counts the backlog from: the position the running trainer + reported, else the newest cut's (D4: a failed trainer's rounds after its last + cut are trained again).""" + running = None if lifecycle is None else lifecycle.running_watermark() + return max(cuts.watermark(), running or 0) + + +def _catch_up(handle: TrainerHandle, target: LifecycleState) -> None: + """Advance the head's copy of the state step by step up to the island's + (events can skip states). STOPPED is the head's own step, after the release.""" + if target is LifecycleState.STOPPED: + return + order = list(LifecycleState) + while order.index(handle.state) < order.index(target): + handle.advance(order[order.index(handle.state) + 1]) diff --git a/yeto/rl/engine/trainer_scheduler.py b/yeto/rl/engine/trainer_scheduler.py new file mode 100644 index 000000000..78791aa0e --- /dev/null +++ b/yeto/rl/engine/trainer_scheduler.py @@ -0,0 +1,181 @@ +"""Scheduler S1: start on backlog, stop by ski rental (rl-trainer-start-stop D6, D8). + +* Start when the backlog ``Q`` (in batches) reaches ``Q_on``. Default + ``Q_on = A_hard + 1``. With the optional ``alpha`` (allowed start-cost share) + ``Q_on = min(Q_s1, A_hard + 1)``, ``Q_s1 = ceil(T_s * (1 - alpha) / (alpha * T_b))``. + ``alpha`` has no default. +* Stop (``STOPPED``) when the backlog is below one batch and the trainer has been idle for + ``T_idle = T_s``; stop (``FINISHED``) when the data ended and less than one batch + is left (a remainder below one batch is never trained). At the data end a cold + trainer starts for the full batches left, also below ``Q_on``. +* ``T_s`` is the last measured start (``trainer_start_summary.total_s``, the event of + rl-trainer-cold-start-cost L0.1) or, without a measurement, the configured value. + With ``fixed_start_seconds`` the configured value is an override (a check run uses + a small one to get several stops); the measured value is still journaled. +* Prices are accepted and ignored (S3 is not in this change). + +The idle start time is the only state. It is a wall-clock time, so a scheduler +restored from the controller journal after a head restart continues with the idle +time already spent (D8). The journal record has kind ``trainer_scheduler_state`` +and holds the inputs and the decision, in the same way as ``auto_state`` +(``elastic.py``). +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from typing import Any + +from yeto.rl.elastic_client import LeaveReason + +from .trainer_lifecycle import BacklogStats, Decision, Hold, LifecycleState, Start, Stop + +START_SUMMARY_EVENT = "trainer_start_summary" # rl-trainer-cold-start-cost L0.1 +SCHEDULER_STATE_KIND = "trainer_scheduler_state" + +# A cold trainer starts from these states; a running trainer is judged for stop +# only while TRAINING. The states in between are transitions: hold. +_STARTABLE = (LifecycleState.COLD, LifecycleState.STOPPED) + +Record = Callable[..., None] # controller._record(kind, tx_id=None, **fields) + + +def last_start_seconds(events: Iterable[Mapping[str, Any]]) -> float | None: + """``total_s`` of the last complete ``trainer_start_summary``, or None.""" + last = None + for e in events: + if e.get("event") == START_SUMMARY_EVENT and e.get("total_s") is not None: + last = float(e["total_s"]) + return last + + +@dataclass(frozen=True) +class S1Config: + """``start_seconds``: configured ``T_s``, used when no start was measured. + ``batch_seconds`` (``T_b``) is needed only with ``alpha``. ``fixed_start_seconds``: + use ``start_seconds`` even after a start was measured.""" + + a_hard: int + start_seconds: float + alpha: float | None = None + batch_seconds: float | None = None + fixed_start_seconds: bool = False + + def __post_init__(self) -> None: + if self.a_hard < 0: + raise ValueError(f"a_hard must be >= 0, got {self.a_hard}") + if self.start_seconds < 0: + raise ValueError(f"start_seconds must be >= 0, got {self.start_seconds}") + if self.alpha is not None: + if not 0.0 < self.alpha <= 1.0: + raise ValueError(f"alpha must be in (0, 1], got {self.alpha}") + if not self.batch_seconds or self.batch_seconds <= 0: + raise ValueError("alpha needs batch_seconds > 0 (T_b)") + + +class S1Scheduler: + """:class:`~.trainer_lifecycle.TrainerScheduler` S1. + + ``start_events`` returns the start-phase events read so far (any iterable of + event dicts); the newest summary gives ``T_s``. ``record`` is the controller + journal writer; without it the state is not journaled.""" + + def __init__(self, config: S1Config, *, + start_events: Callable[[], Iterable[Mapping[str, Any]]] | None = None, + record: Record | None = None, + clock: Callable[[], float] = time.time) -> None: + self.config = config + self._start_events = start_events + self._record = record + self._clock = clock + self.idle_since: float | None = None + + # ------------------------------------------------------------ inputs + def measured_start_seconds(self) -> float | None: + return last_start_seconds(self._start_events()) if self._start_events else None + + def start_seconds(self) -> float: + measured = None if self.config.fixed_start_seconds else self.measured_start_seconds() + return self.config.start_seconds if measured is None else measured + + def q_on(self, start_seconds: float) -> float: + cap = self.config.a_hard + 1 + if self.config.alpha is None: + return cap + a, t_b = self.config.alpha, self.config.batch_seconds + q_s1 = math.ceil(start_seconds * (1.0 - a) / (a * t_b)) + return min(q_s1, cap) + + # ------------------------------------------------------------ decision + def decide(self, backlog: BacklogStats, lifecycle_state: LifecycleState, + prices: Mapping[str, float] | None = None) -> Decision: + del prices # S1 ignores prices (D6) + now = self._clock() + t_s = self.start_seconds() + q_on = self.q_on(t_s) + before = self.idle_since + decision = self._decide(backlog, lifecycle_state, now, t_s, q_on) + self._journal(backlog, lifecycle_state, now, t_s, q_on, decision, before) + return decision + + def _decide(self, backlog: BacklogStats, state: LifecycleState, now: float, + t_s: float, q_on: float) -> Decision: + if backlog.data_ended and backlog.batches < 1: + # less than one batch is left and none will come: nothing more to train + if state is not LifecycleState.TRAINING and state not in _STARTABLE: + return Hold() # a start or a drain is under way + self.idle_since = None + return Stop(LeaveReason.FINISHED) + if state in _STARTABLE: + self.idle_since = None + if backlog.batches >= q_on or backlog.data_ended: + return Start() # at the data end, train the full batches below Q_on + return Hold() + if state is not LifecycleState.TRAINING: + return Hold() + if backlog.batches >= 1: + self.idle_since = None + return Hold() + # below one batch the trainer cannot train: it is idle + if self.idle_since is None: + self.idle_since = now + if now - self.idle_since >= t_s: + self.idle_since = None + return Stop(LeaveReason.STOPPED) + return Hold() + + # ------------------------------------------------------------ journal (D8) + def _journal(self, backlog: BacklogStats, state: LifecycleState, now: float, + t_s: float, q_on: float, decision: Decision, before: float | None) -> None: + # Write when the state changes or the decision acts; a steady Hold is not news. + if self._record is None or (self.idle_since == before and isinstance(decision, Hold)): + return + self._record(SCHEDULER_STATE_KIND, tx_id=None, + state={"idle_since": self.idle_since}, + inputs={"at": now, "lifecycle_state": state.value, + "batches": backlog.batches, "samples": backlog.samples, + "data_ended": backlog.data_ended, + "start_seconds": t_s, "q_on": q_on, + **({"measured_start_seconds": self.measured_start_seconds()} + if self.config.fixed_start_seconds else {})}, + decision=_decision_dict(decision)) + + def restore(self, records: Iterable[Mapping[str, Any]]) -> bool: + """Apply the last journaled state; False if there is none.""" + last = None + for r in records: + if r.get("kind") == SCHEDULER_STATE_KIND: + last = r + if last is None: + return False + self.idle_since = last["state"].get("idle_since") + return True + + +def _decision_dict(decision: Decision) -> dict[str, Any]: + if isinstance(decision, Stop): + return {"action": "stop", "reason": decision.reason.name} + return {"action": "start" if isinstance(decision, Start) else "hold"} diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 104da01b0..d68bfca88 100644 --- a/yeto/wandb_tape.py +++ b/yeto/wandb_tape.py @@ -31,6 +31,7 @@ from collections.abc import Callable, Iterator from pathlib import Path +from .rl.engine.merge_metrics import MergeMetricsFolder from .wandb_logger import NullRun, WandbRun log = logging.getLogger("wandb-tape") @@ -44,6 +45,19 @@ def _is_merge_record(record: dict) -> bool: return "step" in record and record.get("event") != POLICY_SWEEP_LEDGER_EVENT +def _is_elastic_record(record: dict) -> bool: + """An elastic syncer tape line (``kind`` key, see syncer/src/elastic.rs).""" + return "kind" in record and "step" not in record + + +def record_metrics(rec: dict, folder: MergeMetricsFolder) -> dict | None: + """Metrics to log for one tape record, or None. Elastic records go + through ``folder`` and log only once per outer step (L5.5).""" + if _is_elastic_record(rec): + return folder.feed(rec) + return tape_metrics(rec) + + def tape_metrics(rec: dict) -> dict: """One merge record -> the metrics logged for it. @@ -268,6 +282,7 @@ def _forward(self) -> None: start = self.offsets.read() if self.offsets is not None else 0 if start: log.info("resuming the event tape at byte %d", start) + folder = MergeMetricsFolder() try: for rec in follow_jsonl( self.tape_path, @@ -275,7 +290,9 @@ def _forward(self) -> None: on_offset=None if self.offsets is None else self.offsets.record, stop=self._stop, ): - self.run.log(tape_metrics(rec)) + metrics = record_metrics(rec, folder) + if metrics is not None: + self.run.log(metrics) if _is_merge_record(rec): self.records += 1 except Exception as e: # noqa: BLE001 - telemetry never breaks the run @@ -295,6 +312,7 @@ def replay( ) -> int: """Log every record of a finished tape; returns the merge-record count.""" count = 0 + folder = MergeMetricsFolder() with open(os.path.expanduser(str(tape_path)), encoding="utf-8", errors="replace") as f: for line in f: line = line.strip() @@ -305,7 +323,9 @@ def replay( except json.JSONDecodeError as e: log.warning("skipping unparsable event-tape line: %s", e) continue - run.log(tape_metrics(rec)) + metrics = record_metrics(rec, folder) + if metrics is not None: + run.log(metrics) if _is_merge_record(rec): count += 1 run.summary({"sync/tape_records": count})