From 8d56b2d14e0d76120fdcb003eb9f1893c90b5ab3 Mon Sep 17 00:00:00 2001 From: ShubyM Date: Thu, 29 Oct 2026 16:13:06 -0700 Subject: [PATCH 1/3] feat(telemetry): record worker processes and the request ledger through one buffer Each process that records keeps one buffer in memory. A hook on a request path only appends to it, swallows its own errors and returns nothing. One thread drains the buffer into Redis in one pipeline every 5 s. A worker's record of itself goes in the same pipeline, so the thread is also its heartbeat. A failing or hanging Redis costs a flush. The sync Redis client now has socket timeouts, so a hung call cannot hold the thread forever. Trainers, in-process samplers and pool dispatchers keep a worker record: phase, GPU turns and the waits for them, outstanding requests by op and run, and the settings env without secrets. The time slicer's acquire and the wake and sleep inside the turn show on it. The record says exited after a last flush, which an exiting process waits up to 5 s for. The ledger counts each run's trainer requests enqueued (API server), popped and answered (trainer) in open_rl:ledger:. The API server starts its own recorder, which also writes for its in-process trainer. Worker pods get POD_NAME, POD_UID and NODE_NAME for the records. --- Makefile | 2 +- src/server/api_server.py | 5 + src/server/sampling_dispatcher.py | 4 +- src/server/scheduler_worker_manager.py | 2 + src/server/store.py | 51 ++++- src/server/telemetry.py | 230 ++++++++++++++++++++++ src/server/training_requests_processor.py | 16 +- src/server/vllm_sampler.py | 11 +- tests/test_telemetry.py | 169 ++++++++++++++++ 9 files changed, 482 insertions(+), 8 deletions(-) create mode 100644 src/server/telemetry.py create mode 100644 tests/test_telemetry.py 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/src/server/api_server.py b/src/server/api_server.py index a677e856..a1caff94 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, @@ -568,6 +569,7 @@ async def enqueue(command: Command) -> str: active_set_id = await _resolve_active_set_id(command.model_id) 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 +746,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 +782,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/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index a3eebe3d..70d9b293 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -19,6 +19,7 @@ import redis.asyncio as redis from server.sampling_store import SamplingStore, job_of +from server.telemetry import serving, set_phase from server.vllm_options import sample_params, sample_result logger = logging.getLogger(__name__) @@ -368,7 +369,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/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..146c3175 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,11 @@ 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 + class InMemoryStateStore(StateStore): def __init__(self): @@ -292,6 +310,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 +351,19 @@ 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 + class RedisStateStore(StateStore): def __init__(self, client: redis.Redis, sync_client: sync_redis.Redis): @@ -370,6 +402,21 @@ 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() + @cache def redis_client(redis_url: str) -> redis.Redis: @@ -388,5 +435,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..8252e3b7 --- /dev/null +++ b/src/server/telemetry.py @@ -0,0 +1,230 @@ +"""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:ledger: a hash counting the run's trainer requests enqueued, popped and answered + open_rl:worker: a worker process's phase, GPU turns and the waits for them, outstanding requests and settings env +""" + +import asyncio +import json +import os +import re +import socket +import threading +import time +import uuid +from collections import Counter, deque +from collections.abc import Iterable +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass +from functools import wraps + +from server.store import Writes, get_state_store + +TTL_SECONDS = 7 * 86400 +MAX_TURNS = 200 +FLUSH_SECONDS = 5 +FINISH_SECONDS = 5 # how long an exiting process waits for its last flush +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: +# ("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 + + +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 plan(entries: Iterable[tuple]) -> Writes: + """The writes for buffered entries.""" + writes = Writes() + for kind, run, *rest in entries: + if kind == "count": + writes.fields.setdefault(key("ledger", run), Counter())[rest[0]] += 1 + 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, worker: Worker | None) -> None: + self.worker = worker + 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): + writes = plan(entries) + if self.worker is not None: + writes.values[key("worker", self.worker.record["id"])] = self.worker.snapshot() + get_state_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(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) -> None: + if (w := worker()) is not None and (entry := w.outstanding.pop(request_id, None)) is not None: + add("count", entry.run_id, "answered") diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index 53c586f3..ee76803e 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, 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) 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)) @@ -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/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 00000000..2feb382e --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,169 @@ +"""Run telemetry: what the recorder writes, and that recording never changes or slows what a worker answers.""" + +import asyncio +import json +import os +import threading +import time +import unittest +from unittest.mock import Mock, patch + +from server import api_server, telemetry +from server.store import InMemoryStateStore, InMemoryStore, StateStore +from server.training_requests_processor import TrainingRequestsProcessor +from tests.api_client import asgi_client +from tests.test_fft_batch_failure import BatchStore, SlicerStub + +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)] + + +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") + 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}) + + +if __name__ == "__main__": + unittest.main() From 09bf16d5a8ef33b973cc8ff8c4746a03e7e93456 Mon Sep 17 00:00:00 2001 From: ShubyM Date: Thu, 29 Oct 2026 16:19:59 -0700 Subject: [PATCH 2/3] feat(telemetry): record trainer ops by step and serve a run's steps The trainer records one op record per op in open_rl:operations:, capped at the newest 2,000: operation, status, error type, timing, finite metrics and the step it belongs to. A forward-only call records as forward. Its own sample ops keep a separate list. A failure answered before its op ran, such as a failed wake or parse, gets a record too. Each successful optim_step adds to open_rl:steps: and open_rl:step_times:, and a failed op to the ledger's failed count. A step stamp counts this process's steps on top of the run's count, which the recorder reads once per run before it writes the run's first op. A restarted trainer continues the run's numbering. The API server stamps the enqueue time on each command it queues, as an openrl-enqueued-at entry in trace_context. The op record carries queue_seconds, the time from enqueue to the op's start. The stamp rides in trace_context because workers from before it forbid unknown command fields but already accept any trace_context entry, so the API server can roll ahead of the workers. Workers still replace NaN and infinity before answering. sanitize_metrics names the replaced metrics to telemetry first, so a NaN grad norm is recorded in nonfinite instead of as 0.0. GET /api/v1/inspect/runs/{id}/steps?last=&fields= groups the op records by step, with each op's count, failures, seconds and metric spread. Ops after the last optim_step that do no training, such as the final save, join the last step instead of forming an empty one. --- docs/configuration.md | 6 + src/server/api_server.py | 1 + src/server/observability/router.py | 21 +++ src/server/observability/steps.py | 32 +++++ src/server/store.py | 13 ++ src/server/telemetry.py | 155 ++++++++++++++++++++-- src/server/training_requests_processor.py | 6 +- src/training/automodel_worker.py | 2 +- src/training/commands.py | 8 ++ src/training/fft_trainer_worker.py | 14 +- src/training/lora_trainer_worker.py | 6 +- src/training/trainer_worker.py | 8 +- tests/test_commands.py | 7 + tests/test_telemetry.py | 117 +++++++++++++++- 14 files changed, 365 insertions(+), 31 deletions(-) create mode 100644 src/server/observability/steps.py diff --git a/docs/configuration.md b/docs/configuration.md index bf3d570e..6348a72a 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. +Trainers record each op, the run's step count and its request counts in the +API server's Redis, and each worker records its phase, GPU turns and settings. +Records expire after 7 days. `/api/v1/inspect/runs//steps?last=&fields=` +returns the run's 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 a1caff94..d878aed8 100644 --- a/src/server/api_server.py +++ b/src/server/api_server.py @@ -568,6 +568,7 @@ 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 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/store.py b/src/server/store.py index 146c3175..cece00fd 100644 --- a/src/server/store.py +++ b/src/server/store.py @@ -303,6 +303,11 @@ 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): @@ -364,6 +369,9 @@ def write_sync(self, writes: Writes, ttl_seconds: float) -> None: 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): @@ -417,6 +425,11 @@ def write_sync(self, writes: Writes, ttl_seconds: float) -> None: 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: diff --git a/src/server/telemetry.py b/src/server/telemetry.py index 8252e3b7..2dded28d 100644 --- a/src/server/telemetry.py +++ b/src/server/telemetry.py @@ -4,12 +4,18 @@ 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:ledger: a hash counting the run's trainer requests enqueued, popped and answered - open_rl:worker: a worker process's phase, GPU turns and the waits for them, outstanding requests and settings env + open_rl:operations: trainer op records, each stamped with the run's step, the newest SAMPLE_LIMIT + open_rl:operations::samples the trainer's sample op records, kept apart so sampling never evicts training ops + 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 @@ -17,22 +23,31 @@ import time import uuid from collections import Counter, deque -from collections.abc import Iterable +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: -# ("count", run, event) one more of the run's trainer requests reached event +# ("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 +# ("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) @@ -62,6 +77,7 @@ class Outstanding: op: str run_id: str since: float + observed: bool = False # observe_operation records the op, so answered need not class Worker: @@ -95,11 +111,25 @@ def snapshot(self) -> str: return json.dumps({**self.record, "heartbeat_at": time.time(), "outstanding": list(groups.values())}) -def plan(entries: Iterable[tuple]) -> Writes: - """The writes for buffered entries.""" +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]) -> 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 is this process's pod, which every record names.""" writes = Writes() for kind, run, *rest in entries: - if kind == "count": + 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 == "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 return writes @@ -109,7 +139,11 @@ class Recorder: record, and once more when stopped. It is the only thread that writes, so a flush never overlaps another.""" def __init__(self, worker: Worker | None) -> None: - self.worker = worker + self.worker, self.pod = 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) @@ -126,10 +160,14 @@ def flush(self) -> None: while True: entries.append(buffer.popleft()) with suppress(Exception): - writes = plan(entries) + 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) if self.worker is not None: writes.values[key("worker", self.worker.record["id"])] = self.worker.snapshot() - get_state_store().write_sync(writes, TTL_SECONDS) + store.write_sync(writes, TTL_SECONDS) recorder: Recorder | None = None @@ -225,6 +263,97 @@ def popped(requests: list[dict]) -> None: @best_effort -def answered(request_id: str) -> None: - if (w := worker()) is not None and (entry := w.outstanding.pop(request_id, None)) is not None: - add("count", entry.run_id, "answered") +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"], 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, 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": "trainer", + "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, 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 diff --git a/src/server/training_requests_processor.py b/src/server/training_requests_processor.py index ee76803e..243ec45c 100644 --- a/src/server/training_requests_processor.py +++ b/src/server/training_requests_processor.py @@ -22,7 +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, popped, serving, set_phase +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 @@ -208,7 +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) + answered(request_id, result) async def run_once(self) -> None: batch = await self.next_batch() @@ -306,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() 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 index 2feb382e..606c1a88 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -8,11 +8,14 @@ import unittest from unittest.mock import Mock, patch +from fastapi.testclient import TestClient + from server import api_server, telemetry from server.store import InMemoryStateStore, InMemoryStore, StateStore from server.training_requests_processor import TrainingRequestsProcessor from tests.api_client import asgi_client -from tests.test_fft_batch_failure import BatchStore, SlicerStub +from tests.test_fft_batch_failure import BatchStore, BrokenWorker, SlicerStub +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 @@ -64,6 +67,31 @@ 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() @@ -135,7 +163,7 @@ 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") + telemetry.answered("f0", {"type": "forward_backward_completed"}) return json.loads(telemetry.worker().snapshot())["outstanding"] outstanding, _, _ = self.serve(pop) @@ -165,5 +193,90 @@ async def 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}) + + if __name__ == "__main__": unittest.main() From 47a7f81d14e14697a9bd53de30eeb00446a41249 Mon Sep 17 00:00:00 2001 From: ShubyM Date: Thu, 29 Oct 2026 16:24:47 -0700 Subject: [PATCH 3/3] feat(telemetry): record the samples each run's clients got The hook is in SamplingStore.set_result, where the first saved result wins, so it sees every sampling answer a client gets exactly once, from a dispatcher or from the API server's deadline sweep and job failures. A result that loses to an earlier one is not recorded. A failed answer keeps its own record in open_rl:operations::samples. Successful answers are only counted in memory. Each flush writes one group record per run with their count and the window they span, so there is no record per sample. A run's sends are the sum of its group counts plus its failed records. There is no separate sample counter. A send that raises is sent again and is counted once, when it is answered. A failed weight update answers nothing, and the requests wait for the restarted sampler. A revision no replica can load fails each of its requests once. Samples answered before the engine dies are written by the exit's last flush. accepted_at moves to sampling_store, which now needs it for a failure saved without a send. --- docs/configuration.md | 10 +-- src/server/sampling_dispatcher.py | 14 ++-- src/server/sampling_store.py | 17 +++- src/server/telemetry.py | 49 ++++++++--- tests/test_telemetry.py | 130 ++++++++++++++++++++++++++++++ 5 files changed, 191 insertions(+), 29 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 6348a72a..19e258ed 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -138,11 +138,11 @@ 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. -Trainers record each op, the run's step count and its request counts in the -API server's Redis, and each worker records its phase, GPU turns and settings. -Records expire after 7 days. `/api/v1/inspect/runs//steps?last=&fields=` -returns the run's ops grouped by step, with their counts, failures, time and -metrics. +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 | | --- | --- | --- | diff --git a/src/server/sampling_dispatcher.py b/src/server/sampling_dispatcher.py index 70d9b293..08d338b1 100644 --- a/src/server/sampling_dispatcher.py +++ b/src/server/sampling_dispatcher.py @@ -18,7 +18,7 @@ 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 @@ -45,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} @@ -171,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()) @@ -186,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: 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/telemetry.py b/src/server/telemetry.py index 2dded28d..c5fdb422 100644 --- a/src/server/telemetry.py +++ b/src/server/telemetry.py @@ -5,7 +5,8 @@ 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 the trainer's sample op records, kept apart so sampling never evicts training ops + 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 @@ -46,6 +47,7 @@ # ("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. @@ -115,10 +117,11 @@ 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]) -> Writes: +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 is this process's pod, which every record names.""" - writes = Writes() + 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 @@ -126,11 +129,18 @@ def plan(entries: Iterable[tuple], step_starts: dict[str, int], pod: dict[str, s 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 @@ -138,8 +148,8 @@ 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, worker: Worker | None) -> None: - self.worker, self.pod = worker, identity() + 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. @@ -164,7 +174,7 @@ def flush(self) -> None: # 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) + 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) @@ -189,7 +199,7 @@ def start(role: str, runtime_id: str | None = None, *, worker_record: bool = Tru if recorder is not None or os.getenv("RANK", "0") != "0": return with suppress(Exception): - started = Recorder(Worker(role, runtime_id) if worker_record else None) + started = Recorder(role, Worker(role, runtime_id) if worker_record else None) started.thread.start() recorder = started @@ -270,7 +280,7 @@ def answered(request_id: str, result: dict) -> 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"], entry.since, FAILED)) + keep_op(op_record(entry.op, entry.run_id, request_id, w.record["runtime_id"], "trainer", entry.since, FAILED)) # Trainer op records @@ -285,7 +295,7 @@ def nonfinite(names: Iterable[str]) -> None: found.update(names) -def op_record(op: str, run_id: str, request_id: str | None, runtime_id: str | None, started_at: float, error_type: str | None) -> dict: +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, @@ -294,7 +304,7 @@ def op_record(op: str, run_id: str, request_id: str | None, runtime_id: str | No "request_id": request_id, "run_id": run_id, "runtime_id": runtime_id, - "role": "trainer", + "role": role, "status": "failed" if error_type else "succeeded", "error_type": error_type, "elapsed_seconds": now - started_at, @@ -326,7 +336,7 @@ def op_finished(command: Command, runtime_id: str | None, started_at: float, fou 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, started_at, error_type) + 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} @@ -357,3 +367,18 @@ async def observe_operation(command: Command, runtime_id: str | None, call: Call 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/tests/test_telemetry.py b/tests/test_telemetry.py index 606c1a88..eb1f1e4c 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -3,18 +3,29 @@ 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 @@ -278,5 +289,124 @@ def test_metrics_are_spread_per_op_and_kept_by_prefix(self) -> None: 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()