feat(telemetry): record trainer and sampler activity for the inspection API

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),)

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:<runtime>`. 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/<id>/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:<resource>` reads that Cloud Logging resource instead, such as `projects/<project>/locations/global/buckets/<bucket>/views/<view>`. |

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

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,

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)]

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):

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."""

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
 
 

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,13 +298,24 @@ 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):
     self.kv_store: dict[str, str] = {}
     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()