diff --git a/Makefile b/Makefile index f8ff69b0..41bec14d 100644 --- a/Makefile +++ b/Makefile @@ -19,7 +19,7 @@ HOST ?= 127.0.0.1 PORT ?= 9003 # The fully qualified base URL used by local CLI tools and clients BASE_URL ?= http://$(HOST):$(PORT) -UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_perf_compare_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_sampler_batch_limits tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_prefetch tests.test_pool_accept tests.test_sampling_dispatcher tests.test_vllm_sampler tests.test_vllm_options tests.test_observability +UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_trainer_logprob_chunks tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_perf_compare_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_sampler_batch_limits tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_prefetch tests.test_pool_accept tests.test_sampling_dispatcher tests.test_vllm_sampler tests.test_vllm_options tests.test_observability tests.test_telemetry # Only forward BASE_URL to e2e when the user supplied it. The Makefile default # is for local CLI usage; e2e should start its own backend by default. TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),) diff --git a/docs/configuration.md b/docs/configuration.md index bf3d570e..19e258ed 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -138,6 +138,12 @@ run's workers, it also saves each `Workload`'s last placement to the Redis list `open_rl:placements:`. The list keeps the newest 200 records for seven days. +Workers record each run's trainer ops, step count, request counts, samples and +failed samples in the API server's Redis for 7 days. Each worker also records +its own phase, GPU turns and settings. +`/api/v1/inspect/runs//steps?last=&fields=` returns the run's trainer ops +grouped by step, with their counts, failures, time and metrics. + | Env var | Default | What it does | | --- | --- | --- | | `OPEN_RL_LOG_STORE` | unset | Where it reads logs. Unset turns log reads off. `cloud-logging` reads Cloud Logging in the project of the server's Google credentials. `cloud-logging:` reads that Cloud Logging resource instead, such as `projects//locations/global/buckets//views/`. | diff --git a/src/server/api_server.py b/src/server/api_server.py index a677e856..d878aed8 100644 --- a/src/server/api_server.py +++ b/src/server/api_server.py @@ -27,6 +27,7 @@ from starlette.exceptions import HTTPException as StarletteHTTPException from server import openrl_status, proto_codec +from server import telemetry as run_records from server.lora_snapshots import snapshot_name, snapshot_root from server.model_metadata import ( TrainingModelMetadata, @@ -567,7 +568,9 @@ async def enqueue(command: Command) -> str: carrier = await open_future(command.request_id) active_set_id = await _resolve_active_set_id(command.model_id) + carrier[commands.ENQUEUED_AT_KEY] = str(time.time()) await store.put_request(commands.wire(command.model_copy(update={"trace_context": carrier})), active_set_id=active_set_id) + run_records.enqueued(command.model_id) # One line per training request so a request that never reaches a worker can # be traced end to end (the workers log the same id when they pop it). print(f"[API_SERVER] enqueued op={command.op} request_id={command.request_id} model_id={command.model_id} active_set={active_set_id}") @@ -744,6 +747,8 @@ async def lifespan(_: FastAPI): global worker_manager task = None check_sampler_backend() + # The in-process trainer's records, and the requests this server enqueues, are written from here. + run_records.start("api_server", worker_record=False) if is_fft_enabled() or os.getenv("REDIS_URL") or os.getenv("OPEN_RL_WORKER_MANAGER"): worker_manager = create_worker_manager() if is_single_process_mode(): @@ -778,6 +783,7 @@ async def lifespan(_: FastAPI): if worker_manager is not None: worker_manager.close() worker_manager = None + await run_records.finish() # The release tag, baked in at image build (OPEN_RL_VERSION build arg); the diff --git a/src/server/observability/router.py b/src/server/observability/router.py index cb9525c9..17a23943 100644 --- a/src/server/observability/router.py +++ b/src/server/observability/router.py @@ -9,7 +9,9 @@ from server.observability import cloud_logging, fleet, metrics from server.observability.kubernetes import namespace from server.observability.log_selector import ClientLines, WorkerPods +from server.observability.steps import steps from server.store import get_state_store +from server.telemetry import key from server.worker_manager import owner_id router = APIRouter(tags=["inspection"]) @@ -78,6 +80,25 @@ async def query( return {**result, "available": False, "error": f"The metrics source is unavailable: {type(exc).__name__}"} +@router.get("/api/v1/inspect/runs/{run_id}/steps") +async def run_steps( + run_id: str, + last: int = Query(20, ge=1, le=1000), + fields: Annotated[str | None, Query(description="Comma-separated prefixes of the op metrics to keep, e.g. loss,grad_norm. Default: all")] = None, +): + """The run's last steps, oldest first, from the trainer's op records: each step's window, whether its optim_step + succeeded, the metrics that were NaN or infinite, and per op its count, failures, seconds and metric spread.""" + store = get_state_store() + if await store.get_value(key("model_meta", run_id)) is None: + raise HTTPException(404, f"No run {run_id}") + every = steps(await store.read_samples(key("operations", run_id))) + if keep := tuple(p for p in (fields or "").split(",") if p): + for step in every[-last:]: + for op in step["ops"].values(): + op["metrics"] = {name: spread for name, spread in op["metrics"].items() if name.startswith(keep)} + return {"run_id": run_id, "total": len(every), "steps": every[-last:]} + + @router.get("/api/v1/inspect/logs") async def logs( run_id: str, diff --git a/src/server/observability/steps.py b/src/server/observability/steps.py new file mode 100644 index 00000000..b277157a --- /dev/null +++ b/src/server/observability/steps.py @@ -0,0 +1,32 @@ +TRAINING = {"forward_backward", "optim_step"} + + +def steps(ops: list[dict]) -> list[dict]: + """Trainer op records grouped by the step they belong to: the step's window, whether its optim_step succeeded, and + per op its count, failures, seconds and metric spread. Ops after the last optim_step that do no training, such as + the final save or an eval, belong to the last step rather than to one of their own.""" + numbered, completed = [], 0 + for s in ops: + # A record kept before ops carried their step is placed by the optim_steps before it. + numbered.append((s["step"] if isinstance(s.get("step"), int) else completed, s)) + completed += s["operation"] == "optim_step" and s["status"] == "succeeded" + numbers = sorted({number for number, _ in numbered}) + if len(numbers) > 1 and not any(s["operation"] in TRAINING for number, s in numbered if number == numbers[-1]): + numbered = [(numbers[-2] if number == numbers[-1] else number, s) for number, s in numbered] + found: dict[int, dict] = {} + for number, s in numbered: + step = found.setdefault(number, {"step": number, "started_at": s["started_at"], "ended_at": s["at"], "ops": {}, "nonfinite": set()}) + step["started_at"], step["ended_at"] = min(step["started_at"], s["started_at"]), max(step["ended_at"], s["at"]) + op = step["ops"].setdefault(s["operation"], {"count": 0, "failed": 0, "seconds": 0.0, "metrics": {}}) + op["count"] += 1 + op["failed"] += s["status"] != "succeeded" + op["seconds"] += s["elapsed_seconds"] + for name, value in s["metrics"].items(): + op["metrics"].setdefault(name, []).append(value) + step["nonfinite"].update(s.get("nonfinite", ())) + for step in found.values(): + optim = step["ops"].get("optim_step", {"count": 0, "failed": 0}) + step.update(nonfinite=sorted(step["nonfinite"]), complete=optim["count"] > optim["failed"]) + for op in step["ops"].values(): + op["metrics"] = {name: {"min": min(v), "max": max(v), "mean": sum(v) / len(v)} for name, v in op["metrics"].items()} + return [found[k] for k in sorted(found)] diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index a3eebe3d..08d338b1 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -18,7 +18,8 @@ import httpx import redis.asyncio as redis -from server.sampling_store import SamplingStore, job_of +from server.sampling_store import SamplingStore, accepted_at, job_of +from server.telemetry import serving, set_phase from server.vllm_options import sample_params, sample_result logger = logging.getLogger(__name__) @@ -44,11 +45,6 @@ def revision_file(pool: str) -> Path: return Path(tmp_root(), "sampling", pool, "revision") -def accepted_at(r: dict[str, Any]) -> float: - """When the record's request was accepted. A stream id starts with that time in milliseconds.""" - return int(r["entry_id"].split("-")[0]) / 1000 - - def failed(message: str) -> dict[str, str]: return {"type": "RequestFailedResponse", "error_message": message} @@ -170,13 +166,14 @@ def to_send(self) -> list[dict[str, Any]]: turns = sorted(jobs.items(), key=lambda job: self.served.get(job[0], -1)) return [r for turn in itertools.zip_longest(*(records for _, records in turns)) for r in turn if r is not None] - async def set_result(self, r: dict[str, Any], result: dict[str, Any]) -> None: - await self.store.set_result(r, result) + async def set_result(self, r: dict[str, Any], result: dict[str, Any], started_at: float | None = None) -> None: + await self.store.set_result(r, result, started_at) self.records.pop(r["entry_id"], None) async def send(self, r: dict[str, Any]) -> None: """Saves the sampler's answer. None, or an error, leaves the record pending to be sent again.""" self.served[job_of(r)] = next(self.sends) + started_at = time.time() try: # A sampler that never answers must not hold the send past its record's deadline. timeout = max(0.0, accepted_at(r) + SAMPLE_DEADLINE_SECONDS - time.time()) @@ -185,7 +182,7 @@ async def send(self, r: dict[str, Any]) -> None: logger.info("Retrying %s after %r", r["request_id"], exc) return if result is not None: - await self.set_result(r, result) + await self.set_result(r, result, started_at) class HttpSampler: @@ -368,7 +365,8 @@ async def main() -> None: logging.getLogger("httpx").setLevel(logging.WARNING) router = os.getenv("OPEN_RL_SAMPLER_ROUTER", "http://127.0.0.1:8081") store, pool = SamplingStore(redis_client(os.environ["REDIS_URL"])), os.environ["OPEN_RL_SAMPLER_POOL"] - async with http_client() as http: + async with serving("dispatcher", pool), http_client() as http: + set_phase("ready") sampler = HttpSampler(http, pool, router, []) task = asyncio.create_task(Dispatcher(store, pool, sampler, int(os.getenv("OPEN_RL_DISPATCH_CONCURRENCY", "256"))).run()) for sig in (signal.SIGTERM, signal.SIGINT): diff --git a/src/server/sampling_store.py b/src/server/sampling_store.py index 7335eb6e..9ff98b01 100644 --- a/src/server/sampling_store.py +++ b/src/server/sampling_store.py @@ -8,6 +8,8 @@ import redis.asyncio as redis +from server.telemetry import sample_answered + RESULT_TTL_SECONDS = 300 # An accepted request may never get a result, so its idempotency key expires instead of staying until Redis is full. PENDING_TTL_SECONDS = 3 * 24 * 3600 @@ -52,6 +54,11 @@ def job_and_revision(fields: dict[str, str]) -> tuple[str, str]: return job_of(record), record["revision"] +def accepted_at(r: dict[str, Any]) -> float: + """When the record's request was accepted. A stream id starts with that time in milliseconds.""" + return int(r["entry_id"].split("-")[0]) / 1000 + + def entry_record(entry_id: str, fields: dict[str, str]) -> dict[str, Any]: return {**json.loads(fields["record"], object_hook=compact_prompt), "entry_id": entry_id} @@ -88,8 +95,9 @@ async def pending(self, pool: str, limit: int | None = 1000, after: str = "0-0") async def wait_for_new(self, pool: str, after: str, timeout_ms: int) -> None: await self.redis.xread({stream_key(pool): after}, count=1, block=timeout_ms) - async def set_result(self, record: dict[str, Any], result: dict[str, Any]) -> bool: - """Save the first final result and drop the pending entry. False if it already has one.""" + async def set_result(self, record: dict[str, Any], result: dict[str, Any], started_at: float | None = None) -> bool: + """Save the first final result and drop the pending entry. False if it already has one. A saved result is recorded + as a sample its client got, started at started_at, or when the request was accepted.""" stream, future, entry_id = stream_key(record["pool"]), future_key(record["request_id"]), record["entry_id"] async def save(pipe) -> bool: @@ -104,7 +112,10 @@ async def save(pipe) -> bool: return True # Every set_result writes the result key, so watching it lets only one through. - return await self.redis.transaction(save, future, value_from_callable=True) + saved = await self.redis.transaction(save, future, value_from_callable=True) + if saved: + sample_answered(job_of(record), record["pool"], record["request_id"], result, started_at or accepted_at(record)) + return saved async def entries(self, stream: str, until: str = "+"): """The stream's entries up to `until`, a page at a time.""" diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 8949c0e9..a4b7158a 100644 --- a/src/server/scheduler_worker_manager.py +++ b/src/server/scheduler_worker_manager.py @@ -170,6 +170,8 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]: env.append(downward("OPEN_RL_ACCEL_TIMESLICER_HOST", "status.hostIP")) if worker.pool: env.append(downward("POD_IP", "status.podIP")) + # The pod a worker names on the records it keeps. + env += [downward("POD_NAME", "metadata.name"), downward("POD_UID", "metadata.uid"), downward("NODE_NAME", "spec.nodeName")] return env diff --git a/src/server/store.py b/src/server/store.py index e1b840be..cece00fd 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -5,6 +5,8 @@ import os import time from abc import ABC, abstractmethod +from collections import Counter +from dataclasses import dataclass, field from functools import cache from typing import Any @@ -234,6 +236,17 @@ async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None await asyncio.sleep(0.1) +@dataclass +class Writes: + """Writes applied together, each key expiring a TTL after: JSON items appended to lists that keep their newest + `limit`, amounts added to counters and to counts in hashes, and values set.""" + + appends: dict[str, tuple[list[str], int]] = field(default_factory=dict) # key -> (items, limit) + counts: Counter[str] = field(default_factory=Counter) + fields: dict[str, Counter[str]] = field(default_factory=dict) + values: dict[str, str] = field(default_factory=dict) + + class StateStore(ABC): """Application state stored as values with expiry and sets of members.""" @@ -285,6 +298,16 @@ async def recent(self, key: str) -> list[str]: """The list push_recent_sync keeps at key, newest first.""" pass + @abstractmethod + def write_sync(self, writes: Writes, ttl_seconds: float) -> None: + """Apply writes in one round trip, without the event loop, for a thread.""" + pass + + @abstractmethod + async def read_samples(self, key: str) -> list[dict[str, Any]]: + """The JSON items of a list write_sync appends to, oldest first.""" + pass + class InMemoryStateStore(StateStore): def __init__(self): @@ -292,6 +315,7 @@ def __init__(self): self.expiries: dict[str, float] = {} self.sets: dict[str, set[str]] = {} self.lists: dict[str, list[str]] = {} + self.hashes: dict[str, Counter[str]] = {} async def set_value(self, key: str, value: str, ttl_seconds: float | None = None) -> None: self.kv_store[key] = value @@ -332,6 +356,22 @@ def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: async def recent(self, key: str) -> list[str]: return list(self.lists.get(key, [])) + def write_sync(self, writes: Writes, ttl_seconds: float) -> None: + # Lists and hashes here never expire. + for key, (items, limit) in writes.appends.items(): + self.lists[key] = [*self.lists.get(key, []), *items][-limit:] + for key, amount in writes.counts.items(): + self.kv_store[key] = str(int(self.get_value_sync(key) or 0) + amount) + self.expiries[key] = time.monotonic() + ttl_seconds + for key, counts in writes.fields.items(): + self.hashes.setdefault(key, Counter()).update(counts) + for key, value in writes.values.items(): + self.kv_store[key] = value + self.expiries[key] = time.monotonic() + ttl_seconds + + async def read_samples(self, key: str) -> list[dict[str, Any]]: + return [json.loads(item) for item in self.lists.get(key, [])] + class RedisStateStore(StateStore): def __init__(self, client: redis.Redis, sync_client: sync_redis.Redis): @@ -370,6 +410,26 @@ def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: async def recent(self, key: str) -> list[str]: return await self.redis.lrange(key, 0, -1) + def write_sync(self, writes: Writes, ttl_seconds: float) -> None: + ttl = int(ttl_seconds) + pipe = self.sync_redis.pipeline(transaction=False) + for key, (items, limit) in writes.appends.items(): + pipe.rpush(key, *items).ltrim(key, -limit, -1).expire(key, ttl) + for key, amount in writes.counts.items(): + pipe.incrby(key, amount).expire(key, ttl) + for key, counts in writes.fields.items(): + for name, amount in counts.items(): + pipe.hincrby(key, name, amount) + pipe.expire(key, ttl) + for key, value in writes.values.items(): + pipe.set(key, value, ex=ttl) + pipe.execute() + + async def read_samples(self, key: str) -> list[dict[str, Any]]: + raw = await self.redis.lrange(key, 0, -1) + # Decoded off the event loop, since a list holds up to thousands of records. + return await asyncio.to_thread(lambda: [json.loads(item) for item in raw]) + @cache def redis_client(redis_url: str) -> redis.Redis: @@ -388,5 +448,7 @@ def get_store() -> RequestStore: def get_state_store() -> StateStore: redis_url = os.environ.get("REDIS_URL") if redis_url: - return RedisStateStore(redis_client(redis_url), sync_redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5)) + # Timeouts, so a Redis that stops answering costs a thread its call rather than the thread. + sync_client = sync_redis.Redis.from_url(redis_url, decode_responses=True, socket_timeout=5, socket_connect_timeout=5) + return RedisStateStore(redis_client(redis_url), sync_client) return InMemoryStateStore() diff --git a/src/server/telemetry.py b/src/server/telemetry.py new file mode 100644 index 00000000..c5fdb422 --- /dev/null +++ b/src/server/telemetry.py @@ -0,0 +1,384 @@ +"""The records OpenRL keeps about its runs, for the inspection API to read. + +A hook on a request path only appends to this process's buffer in memory. It swallows its own errors and returns +nothing, so telemetry never blocks or changes what it observes. One thread writes the buffer to Redis in one pipeline +every FLUSH_SECONDS, with a worker's record of itself, so a failing or hanging Redis costs a flush and nothing more. + + open_rl:operations: trainer op records, each stamped with the run's step, the newest SAMPLE_LIMIT + open_rl:operations::samples sample records, kept apart so sampling never evicts training ops: per flush, one for + the samples of the run its clients got, with their count, and one per failed sample + open_rl:steps: the run's optimizer steps, and open_rl:step_times: when its last STEP_TIMES ended + open_rl:ledger: a hash counting the run's trainer requests enqueued, popped, answered and failed + open_rl:worker: a worker process's phase, GPU turns and the waits for them, outstanding requests and + settings env +""" + +import asyncio +import contextvars +import json +import math +import os +import re +import socket +import threading +import time +import uuid +from collections import Counter, deque +from collections.abc import Awaitable, Callable, Iterable +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass +from functools import wraps +from numbers import Real + +from server.store import Writes, get_state_store +from training.commands import Command + +TTL_SECONDS = 7 * 86400 +SAMPLE_LIMIT = 2000 # records kept per list +STEP_TIMES = 64 +MAX_TURNS = 200 +FLUSH_SECONDS = 5 +FINISH_SECONDS = 5 # how long an exiting process waits for its last flush +FAILED = "RequestFailedResponse" +SECRET = re.compile(r"authorization|cookie|api[-_]?key|access[-_]?token|secret|password", re.I) +CONFIG_ENV = ("OPEN_RL_", "OPENRL_", "VLLM_", "NCCL_", "CUDA_", "PYTORCH_", "TORCH_") + +# Entries waiting for the next flush, oldest first, each a tuple whose first item names its kind: +# ("op", run, record) a trainer op record, whose step counts this process's steps until a flush adds the run's +# count before them +# ("sample", run, record) a sample record +# ("sampled", run, runtime, started_at, at) one sample a client got +# ("step", run, at) an optimizer step that ended at `at` +# ("count", run, event) one more of the run's trainer requests reached event +# While flushes fail the oldest go. +buffer: deque[tuple] = deque(maxlen=100_000) + + +def key(*parts: str) -> str: + return ":".join(("open_rl", *parts)) + + +def best_effort(hook): + @wraps(hook) + def quiet(*args, **kwargs) -> None: + with suppress(Exception): + hook(*args, **kwargs) + + return quiet + + +def identity() -> dict[str, str]: + """The pod this process runs in, as the scheduler passes it.""" + return {field: value for field, env in (("pod", "POD_NAME"), ("pod_uid", "POD_UID"), ("node", "NODE_NAME")) if (value := os.getenv(env))} + + +@dataclass +class Outstanding: + """A request this worker popped and has not answered.""" + + op: str + run_id: str + since: float + observed: bool = False # observe_operation records the op, so answered need not + + +class Worker: + """A worker process's record of itself. The event loop changes it and the recorder's thread writes it.""" + + def __init__(self, role: str, runtime_id: str | None) -> None: + now = time.time() + self.record = { + # One record per process, so a container restarted in place does not overwrite its predecessor's. + "id": uuid.uuid4().hex[:16], + "node": socket.gethostname(), + **identity(), + "role": role, + "runtime_id": runtime_id, + "started_at": now, + "phase": "starting", + "phase_since": now, + "turns": [], + "waits": [], + # The settings the scheduler passed as env, which is what this process runs with. + "config": {"env": {k: v for k, v in sorted(os.environ.items()) if k.startswith(CONFIG_ENV) and not SECRET.search(k)}}, + } + self.outstanding: dict[str, Outstanding] = {} + + def snapshot(self) -> str: + """The record as JSON, with outstanding requests counted by op and run, each group showing its oldest.""" + groups: dict[tuple, dict] = {} + for entry in list(self.outstanding.values()): # a copy, since the event loop changes it meanwhile + group = groups.setdefault((entry.op, entry.run_id), {"op": entry.op, "run_id": entry.run_id, "since": entry.since, "count": 0}) + group["count"] += 1 + return json.dumps({**self.record, "heartbeat_at": time.time(), "outstanding": list(groups.values())}) + + +def append(writes: Writes, list_key: str, item, limit: int) -> None: + writes.appends.setdefault(list_key, ([], limit))[0].append(json.dumps(item)) + + +def plan(entries: Iterable[tuple], step_starts: dict[str, int], pod: dict[str, str], role: str) -> Writes: + """The writes for buffered entries. step_starts holds each trainer op's run's step count before this process's + first step, and pod and role are this process's, which every record names. The samples a run's clients got make + one record, so its sampling is visible without a record per sample.""" + writes, groups = Writes(), {} + for kind, run, *rest in entries: + if kind == "op": + record = rest[0] | pod + record["step"] += step_starts[run] + append(writes, key("operations", run), record, SAMPLE_LIMIT) + elif kind == "sample": + append(writes, key("operations", run, "samples"), rest[0] | pod, SAMPLE_LIMIT) + elif kind == "sampled": + runtime, started_at, at = rest + if (group := groups.get(run)) is None: + group = groups[run] = op_record("sample", run, None, runtime, role, started_at, None) | {"at": at, "count": 0} + group.update(started_at=min(group["started_at"], started_at), at=max(group["at"], at), count=group["count"] + 1) + elif kind == "step": + writes.counts[key("steps", run)] += 1 + append(writes, key("step_times", run), rest[0], STEP_TIMES) + elif kind == "count": + writes.fields.setdefault(key("ledger", run), Counter())[rest[0]] += 1 + for run, group in groups.items(): + append(writes, key("operations", run, "samples"), group | {"elapsed_seconds": group["at"] - group["started_at"]} | pod, SAMPLE_LIMIT) + return writes + + +class Recorder: + """This process's writer: a thread that drains the buffer into one pipeline every FLUSH_SECONDS, with the worker's + record, and once more when stopped. It is the only thread that writes, so a flush never overlaps another.""" + + def __init__(self, role: str, worker: Worker | None) -> None: + self.role, self.worker, self.pod = role, worker, identity() + # Each run's steps done in this process, counted on the event loop as its ops finish. + self.steps_done: Counter[str] = Counter() + # Each run's step count before this process's first step, read on this thread before its first op is written. + self.step_starts: dict[str, int] = {} + self.stopped = threading.Event() + self.thread = threading.Thread(target=self.beat, name="telemetry", daemon=True) + + def beat(self) -> None: + while not self.stopped.wait(FLUSH_SECONDS): + self.flush() + if self.worker is not None: + self.worker.record.update(phase="exited", phase_since=time.time()) + self.flush() + + def flush(self) -> None: + entries = [] + with suppress(IndexError): + while True: + entries.append(buffer.popleft()) + with suppress(Exception): + store = get_state_store() + # This process's steps for a run land with or after its first op, so the count read here precedes them. + for run in {entry[1] for entry in entries if entry[0] == "op"} - self.step_starts.keys(): + self.step_starts[run] = int(store.get_value_sync(key("steps", run)) or 0) + writes = plan(entries, self.step_starts, self.pod, self.role) + if self.worker is not None: + writes.values[key("worker", self.worker.record["id"])] = self.worker.snapshot() + store.write_sync(writes, TTL_SECONDS) + + +recorder: Recorder | None = None + + +def add(*entry) -> None: + # Nothing is buffered without a recorder to write it, so a process that does not record holds nothing. + if recorder is not None: + buffer.append(entry) + + +def worker() -> Worker | None: + return recorder.worker if recorder is not None else None + + +def start(role: str, runtime_id: str | None = None, *, worker_record: bool = True) -> None: + """Record from now on. Only rank 0 records, since every torchrun rank runs the same ops.""" + global recorder + if recorder is not None or os.getenv("RANK", "0") != "0": + return + with suppress(Exception): + started = Recorder(role, Worker(role, runtime_id) if worker_record else None) + started.thread.start() + recorder = started + + +async def finish() -> None: + """Stop recording: a worker's record says it exited, and the buffer is flushed a last time within FINISH_SECONDS. + Also called before an os._exit.""" + global recorder + stopping, recorder = recorder, None + if stopping is None: + return + stopping.stopped.set() + deadline = time.monotonic() + FINISH_SECONDS + while stopping.thread.is_alive() and time.monotonic() < deadline: + await asyncio.sleep(0.02) + + +@asynccontextmanager +async def serving(role: str, runtime_id: str | None): + """Keep this process's worker record while the body runs, and mark it exited after.""" + start(role, runtime_id) + try: + yield + finally: + await finish() + + +@best_effort +def set_phase(phase: str) -> None: + if (w := worker()) is not None and w.record["phase"] != phase: + w.record.update(phase=phase, phase_since=time.time()) + + +@best_effort +def note(kind: str, started: float, ended: float) -> None: + if (w := worker()) is not None: + w.record[kind] = [*w.record[kind][1 - MAX_TURNS :], [started, ended]] + + +@asynccontextmanager +async def gpu_turn(slot): + """Take the time slicer's slot for the body and show it on the worker record: the wait from asking to the grant, then + the turn, which holds the wake and the sleep around the work.""" + asked = time.time() + async with slot: + granted = time.time() + note("waits", asked, granted) + set_phase("on_gpu") + try: + yield + finally: + note("turns", granted, time.time()) + set_phase("ready") + + +@best_effort +def enqueued(run_id: str) -> None: + add("count", run_id, "enqueued") + + +@best_effort +def popped(requests: list[dict]) -> None: + """Requests this worker took from its queue, each outstanding until answered. A shared trainer's delete_model is + queued without being counted as enqueued, so it is not counted here either.""" + if (w := worker()) is None: + return + for request in requests: + if request.get("request_id") and request.get("op") != "delete_model": + w.outstanding[request["request_id"]] = Outstanding(request["op"], request["model_id"], time.time()) + add("count", request["model_id"], "popped") + + +@best_effort +def answered(request_id: str, result: dict) -> None: + """The request's answer is written. A failure answered without its op running, such as a failed wake or parse, gets + its op record here, so every failed answer has one.""" + if (w := worker()) is None or (entry := w.outstanding.pop(request_id, None)) is None: + return + add("count", entry.run_id, "answered") + if not entry.observed and result.get("type") == FAILED: + keep_op(op_record(entry.op, entry.run_id, request_id, w.record["runtime_id"], "trainer", entry.since, FAILED)) + + +# Trainer op records + +# The names of the metrics an op's worker found NaN or infinite, before it replaced them for JSON. +nonfinite_names: contextvars.ContextVar[set[str] | None] = contextvars.ContextVar("nonfinite_names", default=None) + + +@best_effort +def nonfinite(names: Iterable[str]) -> None: + if (found := nonfinite_names.get()) is not None: + found.update(names) + + +def op_record(op: str, run_id: str, request_id: str | None, runtime_id: str | None, role: str, started_at: float, error_type: str | None) -> dict: + now = time.time() + return { + "at": now, + "started_at": started_at, + "operation": op, + "request_id": request_id, + "run_id": run_id, + "runtime_id": runtime_id, + "role": role, + "status": "failed" if error_type else "succeeded", + "error_type": error_type, + "elapsed_seconds": now - started_at, + "metrics": {}, + } + + +def keep_op(record: dict) -> None: + if recorder is None: + return + run = record["run_id"] + if record["operation"] == "sample": + add("sample", run, record) + return + record["step"] = recorder.steps_done[run] # before this op's own step + add("op", run, record) + if record["status"] == "failed": + add("count", run, "failed") + elif record["operation"] == "optim_step": + recorder.steps_done[run] += 1 + add("step", run, record["at"]) + + +@best_effort +def op_finished(command: Command, runtime_id: str | None, started_at: float, found: set[str], error_type: str | None = None, result=None) -> None: + """Record how one op went: failed with error_type, or as its result says. No payload, output or error message is + kept. A metric is either a finite number or named in nonfinite, from the raw value its worker saw.""" + result = result if isinstance(result, dict) else {} + if error_type is None and result.get("type") == FAILED: + error_type = FAILED + op = "forward" if getattr(command, "forward_only", False) else command.op # an eval loss must not mix with training losses + record = op_record(op, command.model_id, command.request_id, runtime_id, "trainer", started_at, error_type) + numbers = {k: float(v) for k, v in (result.get("metrics") or {}).items() if isinstance(v, Real) and not isinstance(v, bool)} + found = found | {k for k, v in numbers.items() if not math.isfinite(v)} + record["metrics"] = {k: v for k, v in numbers.items() if k not in found} + if found: + record["nonfinite"] = sorted(found) + if command.enqueued_at is not None: + record["queue_seconds"] = started_at - command.enqueued_at + keep_op(record) + + +async def observe_operation(command: Command, runtime_id: str | None, call: Callable[[], Awaitable[dict]]) -> dict: + """Run one trainer op and record how it went. Returns or raises exactly as the op does.""" + started_at, found = time.time(), set() + mark_observed(command.request_id) + token = nonfinite_names.set(found) # the op's worker threads copy this context, so they add to found + try: + result = await call() + except Exception as error: + op_finished(command, runtime_id, started_at, found, type(error).__name__) + raise + finally: + nonfinite_names.reset(token) + op_finished(command, runtime_id, started_at, found, result=result) + return result + + +@best_effort +def mark_observed(request_id: str) -> None: + if (w := worker()) is not None and (entry := w.outstanding.get(request_id)) is not None: + entry.observed = True + + +# Sample records + + +@best_effort +def sample_answered(run_id: str, pool: str, request_id: str, result: dict, started_at: float) -> None: + """A client got its sampling request's result. A failure keeps its own record; a success counts toward the run's next + group record.""" + if recorder is None: + return + if result.get("type") == FAILED: + add("sample", run_id, op_record("sample", run_id, request_id, pool, recorder.role, started_at, FAILED)) + else: + add("sampled", run_id, pool, started_at, time.time()) diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index 53c586f3..243ec45c 100644 --- a/src/server/training_requests_processor.py +++ b/src/server/training_requests_processor.py @@ -22,6 +22,7 @@ from server.model_metadata import get_model_metadata from server.sampling_store import SamplingStore from server.store import RedisStore, RequestStore, StateStore, get_state_store, get_store +from server.telemetry import answered, finish, gpu_turn, observe_operation, popped, serving, set_phase from training import commands from training.commands import parse_command from training.distributed import broadcast_object, is_distributed, is_primary, local_rank @@ -126,6 +127,7 @@ async def run(self) -> None: if self.time_slicer is not None: await self.time_slicer.register(self.workload) self.snapshot_registered = True + set_phase("ready") while True: try: await self.announce_readiness() @@ -183,6 +185,7 @@ async def exit_gracefully(self, unregister: bool = True) -> None: await self.time_slicer.close() except Exception: pass + await finish() os._exit(0) async def next_batch(self) -> list[dict[str, Any]]: @@ -205,6 +208,7 @@ async def publish_result(self, request_id: str | None, result: dict[str, Any]) - # Only rank 0 writes to the store; the other ranks computed the same result. if request_id is not None and is_primary(): await self.store.set_result(request_id, result) + answered(request_id, result) async def run_once(self) -> None: batch = await self.next_batch() @@ -214,6 +218,7 @@ async def run_once(self) -> None: shutdown = any(req.get("op") == "shutdown_workers" for req in batch) work = [req for req in batch if req.get("op") != "shutdown_workers"] + popped(work) model_id = batch[0].get("model_id", "default") results: list[tuple[str | None, dict[str, Any]]] = [] @@ -268,7 +273,7 @@ async def handle_batch(self, requests: list[dict[str, Any]], results: list[tuple save_reqs = [r for r in requests if r.get("op") in save_ops] if gpu_reqs: - async with self.time_slicer.acquire(self.workload): + async with gpu_turn(self.time_slicer.acquire(self.workload)): await asyncio.to_thread(self.worker.wake_up) try: for request in gpu_reqs: @@ -282,7 +287,7 @@ async def handle_batch(self, requests: list[dict[str, Any]], results: list[tuple for request in save_reqs: results.append(await self.handle_request(request)) else: - async with self.time_slicer.acquire(self.workload): + async with gpu_turn(self.time_slicer.acquire(self.workload)): for request in save_reqs: results.append(await self.handle_request(request)) @@ -301,7 +306,7 @@ async def handle_request(self, raw_request: dict[str, Any]) -> tuple[str | None, ctx = propagate.extract(command.trace_context) if command.trace_context else None token = otel_context.attach(ctx) if ctx else None - result = await self.dispatch_operation(command) + result = await observe_operation(command, self.model_id, lambda: self.dispatch_operation(command)) return request_id, result except Exception as exc: traceback.print_exc() @@ -472,6 +477,11 @@ def run_probe_server(): ) +async def serve_trainer(args: argparse.Namespace) -> None: + async with serving("trainer", args.model_id): + await main_async(args) + + def start_request_processing_loop() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model-id", help="Model id whose per-model request queue this dedicated trainer worker drains.") @@ -485,7 +495,7 @@ def start_request_processing_loop() -> None: cuda_devs = os.getenv("CUDA_VISIBLE_DEVICES", "ALL") print(f"-> Hardware : CUDA_VISIBLE_DEVICES={cuda_devs}") - asyncio.run(main_async(args)) + asyncio.run(serve_trainer(args)) if __name__ == "__main__": diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index 791f1c3f..fc9d906e 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -19,6 +19,7 @@ from server.lora_snapshots import snapshot_root from server.sampling_dispatcher import Dispatcher, RevisionUnloadable, delta_applies, failed, redis_client from server.sampling_store import SamplingStore +from server.telemetry import gpu_turn, serving, set_phase from server.vllm_options import engine_args, sample_params, sample_result logger = logging.getLogger(__name__) @@ -133,7 +134,7 @@ async def holding_gpu(time_slicer: TimeSlicerClient | None, workload: WorkloadRe if time_slicer is None: yield return - async with time_slicer.acquire(workload): + async with gpu_turn(time_slicer.acquire(workload)): if sampler is not None: await sampler.wake() try: @@ -189,6 +190,7 @@ async def run_sampling_worker(pool: str) -> None: if time_slicer is not None: await sampler.sleep() # give the memory back before releasing the turn Path(READY_FILE).touch() + set_phase("ready") concurrency = int(os.getenv("OPEN_RL_DISPATCH_CONCURRENCY", "256")) if time_slicer is not None: dispatcher = TimeSlicedDispatcher(store, pool, sampler, concurrency, time_slicer, workload) @@ -216,12 +218,17 @@ async def run_sampling_worker(pool: str) -> None: await store.redis.aclose() +async def serve_pool(pool: str) -> None: + async with serving("sampler", pool): + await run_sampling_worker(pool) + + def main() -> None: parser = argparse.ArgumentParser(description="Open-RL sampling pool replica with vLLM in this process") parser.add_argument("--pool", required=True, help="The sampling pool whose stream this replica serves") args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") - asyncio.run(run_sampling_worker(args.pool)) + asyncio.run(serve_pool(args.pool)) if __name__ == "__main__": diff --git a/src/training/automodel_worker.py b/src/training/automodel_worker.py index a2cb5f69..bec73995 100644 --- a/src/training/automodel_worker.py +++ b/src/training/automodel_worker.py @@ -486,7 +486,7 @@ def optim_step(self, adam_params: dict[str, Any], model_id: str | None = None) - total_norm = self.clip_gradients(max_grad_norm if max_grad_norm > 0 else math.inf) state.optimizer.step() state.optimizer.zero_grad(set_to_none=True) - return {"metrics": {"grad_norm:mean": self.sanitize_float(total_norm)}} + return {"metrics": self.sanitize_metrics({"grad_norm:mean": total_norm})} def new_optimizer(self, adam_params: dict[str, Any]) -> torch.optim.Optimizer: return torch.optim.AdamW( diff --git a/src/training/commands.py b/src/training/commands.py index b2def1f8..e2590126 100644 --- a/src/training/commands.py +++ b/src/training/commands.py @@ -12,6 +12,8 @@ from training.types import Datum, FFTConfig, FineTuningType, LoraConfig SHUTDOWN_REQUEST_ID = "SHUTDOWN_SENTINEL" +# Rides in trace_context, which workers from before the stamp already accept; they forbid a new field. +ENQUEUED_AT_KEY = "openrl-enqueued-at" class Command(BaseModel): @@ -22,6 +24,12 @@ class Command(BaseModel): model_id: str trace_context: dict[str, str] | None = None + @property + def enqueued_at(self) -> float | None: + """When the API server queued it, in Unix seconds, for the trainer's record of how long it waited.""" + stamp = (self.trace_context or {}).get(ENQUEUED_AT_KEY) + return None if stamp is None else float(stamp) + class CreateModel(Command): op: Literal["create_model"] = "create_model" diff --git a/src/training/fft_trainer_worker.py b/src/training/fft_trainer_worker.py index c8562026..19f40e79 100644 --- a/src/training/fft_trainer_worker.py +++ b/src/training/fft_trainer_worker.py @@ -447,12 +447,14 @@ def optim_step(self, adam_params: dict[str, Any], model_id: str | None = None) - ) return { - "metrics": { - "grad_norm:mean": self.sanitize_float(total_norm.item()), - "time/compute_delta_diff": self.sanitize_float(delta_compute_time), - "time/optimizer_step": self.sanitize_float(step_time), - "time/clip_grad_norm": self.sanitize_float(clip_time), - }, + "metrics": self.sanitize_metrics( + { + "grad_norm:mean": total_norm.item(), + "time/compute_delta_diff": delta_compute_time, + "time/optimizer_step": step_time, + "time/clip_grad_norm": clip_time, + } + ), } def generate( diff --git a/src/training/lora_trainer_worker.py b/src/training/lora_trainer_worker.py index a51b0905..6b822f4e 100644 --- a/src/training/lora_trainer_worker.py +++ b/src/training/lora_trainer_worker.py @@ -332,11 +332,7 @@ def optim_step(self, adam_params: dict[str, Any], model_id: str) -> dict[str, An self.save_adapter(model_id) - return { - "metrics": { - "grad_norm:mean": self.sanitize_float(total_norm.item()), - }, - } + return {"metrics": self.sanitize_metrics({"grad_norm:mean": total_norm.item()})} def generate( self, diff --git a/src/training/trainer_worker.py b/src/training/trainer_worker.py index 0141ab7e..4eba178c 100644 --- a/src/training/trainer_worker.py +++ b/src/training/trainer_worker.py @@ -8,6 +8,7 @@ from torch.utils.checkpoint import checkpoint from transformers import PreTrainedModel, PreTrainedTokenizerBase +from server import telemetry from training import losses from training.distributed import local_rank from training.types import Datum @@ -199,7 +200,7 @@ def _finish(self, data: list[Datum], loss_fn_outputs: list[dict[str, Any] | None completed_loss_fn_outputs.append(output) return { - "metrics": {"loss:mean": self.sanitize_float(mean_loss), "loss:sum": self.sanitize_float(total_loss)}, + "metrics": self.sanitize_metrics({"loss:mean": mean_loss, "loss:sum": total_loss}), "loss_fn_outputs": completed_loss_fn_outputs, "loss_fn_output_type": "ArrayRecord", } @@ -419,6 +420,11 @@ def prompt_logprobs(self, model: PreTrainedModel, input_tensor: torch.Tensor) -> return prompt_logprobs + def sanitize_metrics(self, metrics: dict[str, float]) -> dict[str, float]: + """The metrics with NaN and infinity replaced, for JSON. Telemetry gets the names of the replaced ones first.""" + telemetry.nonfinite(name for name, value in metrics.items() if not math.isfinite(value)) + return {name: self.sanitize_float(value) for name, value in metrics.items()} + def sanitize_float(self, val: float) -> float: if math.isinf(val): return -9999.0 if val < 0 else 9999.0 diff --git a/tests/test_commands.py b/tests/test_commands.py index 609814a8..4e0fd48e 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -61,6 +61,13 @@ def test_wire_data_is_readable_by_existing_workers(self) -> None: datum = commands.wire(command)["payload"]["data"][0] self.assertEqual([token for chunk in datum["model_input"]["chunks"] for token in chunk["tokens"]], [1, 2]) + def test_enqueue_stamp_keeps_the_envelope_existing_workers_accept(self) -> None: + trace_context = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", commands.ENQUEUED_AT_KEY: "1700000000.5"} + raw = commands.wire(commands.OptimStep(request_id="r", model_id="m", trace_context=trace_context)) + # A worker from before the stamp forbids any envelope key it does not know. + self.assertEqual(set(raw), {"op", "request_id", "model_id", "trace_context", "payload"}) + self.assertEqual(commands.parse_command(raw).enqueued_at, 1700000000.5) + def test_ambiguous_envelope_is_rejected(self) -> None: with self.assertRaises(ValueError): commands.parse_command({"request_id": "r", "model_id": "m", "op": "optim_step", "payload": {"model_id": "other"}}) diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 00000000..eb1f1e4c --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,412 @@ +"""Run telemetry: what the recorder writes, and that recording never changes or slows what a worker answers.""" + +import asyncio +import json +import os +import sys +import tempfile +import threading +import time +import types +import unittest +import uuid +from pathlib import Path +from unittest.mock import Mock, patch + +from fastapi.testclient import TestClient + +from server import api_server, telemetry +from server.sampling_dispatcher import Dispatcher, RevisionUnloadable, failed +from server.sampling_store import SamplingStore, future_key +from server.store import InMemoryStateStore, InMemoryStore, StateStore +from server.training_requests_processor import TrainingRequestsProcessor +from server.vllm_sampler import Sampler +from tests.api_client import asgi_client +from tests.redis_server import RedisTestCase +from tests.test_fft_batch_failure import BatchStore, BrokenWorker, SlicerStub +from tests.test_sampling_dispatcher import SAMPLE, until +from tests.test_vllm_sampler import FakeEngine, fake_vllm +from training.trainer_worker import BaseTrainerWorker + +BOUND = 1.0 # an exiting process waits this long for its last flush; a serve must finish well inside it + + +class FaultyStore: + """Every call raises, as with Redis down, or waits forever, as with a Redis that never answers.""" + + def __init__(self, hangs: bool): + self.hangs = hangs + + def __getattr__(self, name): + if not self.hangs: + raise ConnectionError(f"telemetry store {name} failed") + return lambda *args, **kwargs: threading.Event().wait() + + +def no_store(): + raise RuntimeError("no telemetry store configured") + + +class RecorderOnlyStore(InMemoryStateStore): + """Notes, and raises on, any use off the recorder's thread, since a Redis call on a request path is the bug.""" + + def __init__(self): + super().__init__() + self.intrusions: list[str] = [] + + def __getattribute__(self, name): + if name in StateStore.__abstractmethods__ and threading.current_thread().name != "telemetry": + super().__getattribute__("intrusions").append(name) + raise AssertionError(f"{name} called on a request path") + return super().__getattribute__(name) + + +def trainer() -> Mock: + return Mock(cpu_offload=True, optim_step=lambda adam_params, model_id: {"metrics": {"grad_norm:mean": 0.5}}) + + +async def train(worker: Mock, requests: list[dict]) -> dict: + """One time-sliced trainer batch. Returns what each request was answered.""" + store = BatchStore(requests) + with patch.dict(os.environ, {"REDIS_URL": "redis://test"}): + processor = TrainingRequestsProcessor(store, worker, "run-a", time_slicer=SlicerStub()) + await processor.run_once() + return store.futures + + +def optim_steps(count: int) -> list[dict]: + return [{"request_id": f"os-{i}", "op": "optim_step", "model_id": "run-a", "payload": {"adam_params": {}}} for i in range(count)] + + +def request(request_id: str, op: str, **payload) -> dict: + return {"request_id": request_id, "op": op, "model_id": "run-a", "payload": payload} + + +class ScriptedTrainer(BaseTrainerWorker): + """Answers forward_backward with a loss and optim_step with a NaN grad norm, through the workers' own sanitizing, + and fails an optim_step whose learning rate is negative.""" + + cpu_offload = True + + def wake_up(self) -> None: + pass + + def sleep(self) -> None: + pass + + def forward_backward(self, data, loss_fn, loss_config, model_id, forward_only=False) -> dict: + return {"metrics": self.sanitize_metrics({"loss:mean": 1.5})} + + def optim_step(self, adam_params, model_id) -> dict: + if adam_params.get("learning_rate", 0) < 0: + raise ValueError("negative learning rate") + return {"metrics": self.sanitize_metrics({"grad_norm:mean": float("nan"), "time/optimizer_step": 0.25})} + + +class TelemetryTest(unittest.TestCase): + def setUp(self) -> None: + self.store = InMemoryStateStore() + self.enterContext(patch.object(telemetry, "get_state_store", lambda: self.store)) + self.enterContext(patch.object(telemetry, "FINISH_SECONDS", BOUND)) + + def serve(self, work, role: str = "trainer") -> tuple: + """Run work while recording. Returns its result, how long it took, and how long the exit after took.""" + + async def run(): + async with telemetry.serving(role, "run-a"): + started = time.monotonic() + result = await asyncio.wait_for(work(), 2 * BOUND) + exiting = time.monotonic() + return result, exiting - started, time.monotonic() - exiting + + return asyncio.run(run()) + + def worker_records(self) -> list[dict]: + return [json.loads(v) for k, v in self.store.kv_store.items() if k.startswith(telemetry.key("worker", ""))] + + +class FaultTest(TelemetryTest): + """The same work, with a healthy telemetry store and with one that fails, gets the same answers in about the same time.""" + + def test_a_failing_hanging_or_missing_store_changes_no_answer(self) -> None: + expected, _, _ = self.serve(lambda: train(trainer(), optim_steps(2))) + self.assertEqual(self.store.hashes[telemetry.key("ledger", "run-a")], {"popped": 2, "answered": 2}, "the healthy serve recorded") + for fault, store in {"raises": lambda: FaultyStore(False), "hangs": lambda: FaultyStore(True), "is missing": no_store}.items(): + with self.subTest(store=fault), patch.object(telemetry, "get_state_store", store): + result, took, exiting = self.serve(lambda: train(trainer(), optim_steps(2))) + self.assertEqual(result, expected) + self.assertLess(took, BOUND) + self.assertLess(exiting, 2 * BOUND) + + def test_no_request_path_touches_the_store_and_the_recorder_writes_what_it_saw(self) -> None: + store = RecorderOnlyStore() + with patch.object(telemetry, "get_state_store", lambda: store): + result, _, _ = self.serve(lambda: train(trainer(), optim_steps(2))) + self.assertEqual([r["type"] for r in result.values()], ["optim_step_completed"] * 2) + self.assertEqual(store.intrusions, []) + self.assertEqual(store.hashes[telemetry.key("ledger", "run-a")], {"popped": 2, "answered": 2}) + + +class WorkerRecordTest(TelemetryTest): + def test_the_record_shows_settings_turns_and_a_clean_exit(self) -> None: + env = {"OPEN_RL_FINE_TUNING_TYPE": "full", "OPEN_RL_API_KEY": "sk-private", "HOME": "/root"} + with patch.dict(os.environ, env, clear=True): + self.serve(lambda: train(trainer(), optim_steps(1))) + (record,) = self.worker_records() + self.assertEqual((record["role"], record["runtime_id"], record["phase"]), ("trainer", "run-a", "exited")) + self.assertEqual(record["config"]["env"], {"OPEN_RL_FINE_TUNING_TYPE": "full"}) + self.assertNotIn("sk-private", json.dumps(record)) + self.assertEqual((len(record["waits"]), len(record["turns"])), (1, 1)) + + def test_the_record_stays_fresh_while_the_event_loop_is_blocked(self) -> None: + async def block(): + blocked_at = time.time() + time.sleep(0.3) # loading weights holds the loop, and the recorder must not need it + return blocked_at + + with patch.object(telemetry, "FLUSH_SECONDS", 0.02), patch.object(self.store, "write_sync", wraps=self.store.write_sync) as write: + blocked_at, _, _ = self.serve(block) + beats = [json.loads(v)["heartbeat_at"] for call in write.call_args_list for v in call.args[0].values.values()] + self.assertTrue(any(beat - blocked_at > 0.2 for beat in beats[:-1]), "a beat landed while the loop was blocked, before the exit's") + + def test_popped_requests_stay_outstanding_until_answered(self) -> None: + async def pop(): + telemetry.popped([{"request_id": f"f{i}", "op": "forward_backward", "model_id": "run-a"} for i in range(3)]) + # A shared trainer's delete_model was never counted as enqueued, so it is not outstanding either. + telemetry.popped([{"request_id": "d1", "op": "delete_model", "model_id": "run-a"}]) + telemetry.answered("f0", {"type": "forward_backward_completed"}) + return json.loads(telemetry.worker().snapshot())["outstanding"] + + outstanding, _, _ = self.serve(pop) + self.assertEqual([(g["op"], g["run_id"], g["count"]) for g in outstanding], [("forward_backward", "run-a", 2)]) + self.assertEqual(self.store.hashes[telemetry.key("ledger", "run-a")], {"popped": 3, "answered": 1}) + + +class LedgerTest(TelemetryTest): + def setUp(self) -> None: + super().setUp() + self.enterContext(patch.object(api_server, "store", InMemoryStore())) + self.enterContext(patch.object(api_server, "state", InMemoryStateStore())) + self.enterContext(patch("server.worker_manager.get_state_store", return_value=api_server.state)) + self.enterContext(patch.dict(os.environ, {"SAMPLING_BACKEND": "torch"})) + + def test_a_shared_trainers_delete_is_counted_neither_enqueued_nor_popped(self) -> None: + async def run(): + async with asgi_client() as client: + created = (await client.post("/api/v1/create_model", json={"base_model": "m"})).json() + (await client.post("/api/v1/delete_model", json={"model_id": created["request_id"]})).raise_for_status() + requests = await api_server.store.get_requests(active_set_id="m-1") + self.assertEqual([r["op"] for r in requests], ["create_model", "delete_model"]) + telemetry.popped(requests) + return created["request_id"] + + run_id, _, _ = self.serve(run) + self.assertEqual(self.store.hashes[telemetry.key("ledger", run_id)], {"enqueued": 1, "popped": 1}) + + +class OpRecordTest(TelemetryTest): + def records(self, *parts: str) -> list[dict]: + return asyncio.run(self.store.read_samples(telemetry.key("operations", *parts))) + + def test_ops_carry_the_runs_step_and_a_nan_metric_is_named_not_zeroed(self) -> None: + asyncio.run(self.store.set_value(telemetry.key("steps", "run-a"), "5")) # a restarted trainer resumes the run's count + requests = [ + request("fb-1", "forward_backward", data=[]), + request("os-1", "optim_step"), + request("os-2", "optim_step", adam_params={"learning_rate": -1}), + request("fb-2", "forward_backward", data=[]), + request("ev-1", "forward_backward", data=[], forward_only=True), + ] + answers, _, _ = self.serve(lambda: train(ScriptedTrainer(), requests)) + self.assertEqual(answers["os-1"]["metrics"]["grad_norm:mean"], 0.0, "the client still gets the sanitized value") + ops = self.records("run-a") + self.assertEqual( + [(r["request_id"], r["operation"], r["status"], r["step"]) for r in ops], + [("fb-1", "forward_backward", "succeeded", 5), ("os-1", "optim_step", "succeeded", 5), ("os-2", "optim_step", "failed", 6)] + + [("fb-2", "forward_backward", "succeeded", 6), ("ev-1", "forward", "succeeded", 6)], + ) + self.assertEqual((ops[1]["metrics"], ops[1]["nonfinite"]), ({"time/optimizer_step": 0.25}, ["grad_norm:mean"])) + self.assertEqual(self.store.get_value_sync(telemetry.key("steps", "run-a")), "6") + self.assertEqual(len(self.store.lists[telemetry.key("step_times", "run-a")]), 1) + self.assertEqual(self.store.hashes[telemetry.key("ledger", "run-a")], {"popped": 5, "answered": 5, "failed": 1}) + + def test_a_failure_answered_before_the_op_runs_is_recorded_once(self) -> None: + async def fail(): + store = BatchStore([request("fb-1", "forward_backward", data=[]), request("bad-1", "optim_step", adam_params="not a dict")]) + with patch.dict(os.environ, {"REDIS_URL": "redis://test"}): + await TrainingRequestsProcessor(store, BrokenWorker(), "run-a", time_slicer=SlicerStub()).run_once() + + with self.assertRaisesRegex(RuntimeError, "CUDA out of memory"): + self.serve(fail) + self.assertEqual([(r["request_id"], r["status"]) for r in self.records("run-a")], [("fb-1", "failed"), ("bad-1", "failed")]) + + def test_an_op_records_how_long_it_waited_in_the_queue(self) -> None: + async def run(): + async with asgi_client() as client: + created = (await client.post("/api/v1/create_model", json={"base_model": "m"})).json() + await asyncio.sleep(0.2) + await TrainingRequestsProcessor(api_server.store, Mock(), active_tenant_set_id="m-1").run_once() + return created["request_id"] + + self.enterContext(patch.object(api_server, "store", InMemoryStore())) + self.enterContext(patch.object(api_server, "state", InMemoryStateStore())) + self.enterContext(patch("server.worker_manager.get_state_store", return_value=api_server.state)) + with patch.dict(os.environ, {"SAMPLING_BACKEND": "torch"}): + run_id, _, _ = self.serve(run) + (created,) = self.records(run_id) + self.assertGreaterEqual(created["queue_seconds"], 0.2) + self.assertLess(created["queue_seconds"], 2) + + +class StepsRouteTest(TelemetryTest): + def op(self, step: int, operation: str, metrics: dict | None = None) -> dict: + timing = {"started_at": step, "at": step + 0.5, "elapsed_seconds": 0.5} + return {"step": step, "operation": operation, "status": "succeeded", "metrics": metrics or {}, **timing} + + def get(self, ops: list[dict], query: str = "") -> dict: + asyncio.run(self.store.set_value(telemetry.key("model_meta", "run-a"), "{}")) + self.store.lists[telemetry.key("operations", "run-a")] = [json.dumps(op) for op in ops] + with patch("server.observability.router.get_state_store", return_value=self.store): + return TestClient(api_server.app).get(f"/api/v1/inspect/runs/run-a/steps{query}").json() + + def test_ops_after_the_last_step_join_it_and_a_step_in_progress_stands_alone(self) -> None: + train = [self.op(0, "forward_backward", {"loss": 2.0}), self.op(0, "optim_step"), self.op(1, "forward_backward", {"loss": 1.0})] + finished = self.get([*train, self.op(1, "optim_step"), self.op(2, "save_state"), self.op(2, "forward", {"loss": 0.5})]) + self.assertEqual( + [(s["step"], sorted(s["ops"]), s["complete"]) for s in finished["steps"]], + [ + (0, ["forward_backward", "optim_step"], True), + (1, ["forward", "forward_backward", "optim_step", "save_state"], True), + ], + ) + in_progress = self.get([*train, self.op(1, "optim_step"), self.op(2, "save_weights_for_sampler"), self.op(2, "forward_backward")]) + self.assertEqual([(s["step"], s["complete"]) for s in in_progress["steps"]], [(0, True), (1, True), (2, False)]) + + def test_metrics_are_spread_per_op_and_kept_by_prefix(self) -> None: + fb = [self.op(0, "forward_backward", {"loss": 2.0, "kl": 0.1, "entropy": 3.0}), self.op(0, "forward_backward", {"loss": 1.0})] + (step,) = self.get([*fb, self.op(0, "optim_step")], "?fields=loss,kl")["steps"] + spread = {"loss": {"min": 1.0, "max": 2.0, "mean": 1.5}, "kl": {"min": 0.1, "max": 0.1, "mean": 0.1}} + self.assertEqual(step["ops"]["forward_backward"], {"count": 2, "failed": 0, "seconds": 1.0, "metrics": spread}) + + +class ScriptedSampler: + """Answers prompt [1], fails the first send of [2] so the dispatcher sends it again, refuses [3], and can find no + weights a revision applies on.""" + + def __init__(self, unloadable: bool = False) -> None: + self.unloadable, self.raised = unloadable, False + + async def ensure_weights(self, revision: str) -> bool: + if revision and self.unloadable: + raise RevisionUnloadable(revision, ["default"]) + return True + + async def generate(self, r: dict) -> dict: + prompt = list(r["prompt_token_ids"]) + if prompt == [2] and not self.raised: + self.raised = True + raise ConnectionError("the replica went away") + return failed("Sampler rejected the request") if prompt == [3] else SAMPLE + + +class BrokenUpdateEngine(FakeEngine): + async def pause_generation(self, **kwargs) -> None: + raise RuntimeError("CUDA out of memory") + + +class SampleRecordTest(RedisTestCase): + """A run's sends are read back from its sample records: each group record's count, plus one per failed record.""" + + async def asyncSetUp(self) -> None: + await super().asyncSetUp() + self.store, self.pool, self.telemetry = SamplingStore(self.redis), f"sampler-{uuid.uuid4()}", InMemoryStateStore() + self.enterContext(patch.object(telemetry, "get_state_store", lambda: self.telemetry)) + vllm = fake_vllm() | {name: types.ModuleType(name) for name in ("vllm.distributed", "vllm.distributed.weight_transfer.base")} + vllm["vllm.distributed.weight_transfer.base"].WeightTransferUpdateRequest = dict + self.enterContext(patch.dict(sys.modules, vllm)) + self.tmp = self.enterContext(tempfile.TemporaryDirectory()) + self.enterContext(patch.dict(os.environ, {"OPEN_RL_TMP_DIR": self.tmp})) + + async def accepted(self, prompt: list[int], revision: str = "") -> dict: + r = {"request_id": str(uuid.uuid4()), "pool": self.pool, "job": "run", "revision": revision, "model": "base", "idempotency_key": None} + r |= {"prompt_token_ids": prompt, "sampling_params": {"n": 1, "max_tokens": 1}, "include_prompt_logprobs": False} + await self.store.accept(r) + return r + + async def answer(self, r: dict) -> str | None: + raw = await self.redis.lindex(future_key(r["request_id"]), 0) + return json.loads(raw)["type"] if raw else None + + def serve(self, sampler) -> None: + """Serve the pool as a sampler process does, until it exits.""" + telemetry.start("sampler", self.pool) + self.running = asyncio.create_task(Dispatcher(self.store, self.pool, sampler, 4).run()) + self.addAsyncCleanup(self.exit) + + async def exit(self) -> None: + self.running.cancel() + await asyncio.gather(self.running, return_exceptions=True) + await telemetry.finish() + + async def sent(self) -> dict: + """The run's sends and failures, read back from what the sampler process recorded by the time it exited.""" + await self.exit() + records = await self.telemetry.read_samples(telemetry.key("operations", "run", "samples")) + return {"sends": sum(r.get("count", 1) for r in records), "failed": sum(r["status"] == "failed" for r in records)} + + async def test_a_request_sent_again_after_an_error_and_one_refused_are_each_sent_once(self) -> None: + records = [await self.accepted([prompt]) for prompt in (1, 2, 3)] + self.serve(ScriptedSampler()) + await until(lambda: all_answered(self, records)) + self.assertEqual([await self.answer(r) for r in records], ["sample", "sample", "RequestFailedResponse"]) + self.assertEqual(await self.sent(), {"sends": 3, "failed": 1}) + + async def test_an_answer_saved_after_another_one_is_not_counted(self) -> None: + r = await self.accepted([1]) + release = asyncio.Event() + + class SlowSampler(ScriptedSampler): + async def generate(self, r: dict) -> dict: + await release.wait() + return SAMPLE + + self.serve(SlowSampler()) + await asyncio.sleep(0.2) + # The API server's deadline sweep fails the request while its send is in flight. + await self.store.fail_older_than(0, failed("No sampler answered")) + release.set() + await asyncio.sleep(0.2) + self.assertEqual(await self.answer(r), "RequestFailedResponse") + self.assertEqual(await self.sent(), {"sends": 1, "failed": 1}) + + async def test_a_failed_weight_update_answers_nothing_and_unloadable_weights_fail_each_request_once(self) -> None: + revision = Path(self.tmp, "w1") + revision.mkdir() + (revision / "metadata.json").write_text("{}") + records = [await self.accepted([1], str(revision)) for _ in range(2)] + self.serve(Sampler(BrokenUpdateEngine(), "base")) + await asyncio.sleep(0.5) # the worker exits now, and the records wait for the restarted one + self.assertEqual([await self.answer(r) for r in records], [None, None]) + self.assertEqual(await self.sent(), {"sends": 0, "failed": 0}) + self.serve(ScriptedSampler(unloadable=True)) + await until(lambda: all_answered(self, records)) + self.assertEqual(await self.sent(), {"sends": 2, "failed": 2}) + + async def test_samples_answered_before_the_engine_died_are_sent_once_as_is_the_one_it_failed(self) -> None: + engine = FakeEngine(poison=[6, 6, 6]) + self.serve(Sampler(engine, "base")) + answered = [await self.accepted([1]), await self.accepted([2])] + await until(lambda: all_answered(self, answered)) + poison = await self.accepted([6, 6, 6]) + await until(lambda: self.answer(poison)) + self.assertEqual([await self.answer(r) for r in [*answered, poison]], ["sample", "sample", "RequestFailedResponse"]) + self.assertTrue(engine.errored) + self.assertEqual(await self.sent(), {"sends": 3, "failed": 1}) + + +async def all_answered(test: SampleRecordTest, records: list[dict]) -> bool: + return all([await test.answer(r) for r in records]) + + +if __name__ == "__main__": + unittest.main()