feat(observability): show runs' pods and nodes and serve their logs

Makefile

@@ -19,11 +19,11 @@ 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_dashboard
+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
 # 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),)
-UNIT_TESTS += tests.test_dashboard_kubernetes
+UNIT_TESTS += tests.test_observability_kubernetes
 TRAINING_TEST_EXTRA ?= gpu
 TRAINING_TEST_ARGS ?=
 PIGLATIN_TEST_PYTHONPATH ?= examples/sft/pig-latin

docs/configuration.md

@@ -126,12 +126,19 @@ for. On one machine every pool has one replica, in process. See
 ## Inspection API
 
 The API server serves a read-only inventory of runs at
-`/api/v1/inspect/inventory`. It reads runs and their live sessions from the
-same Redis as the rest of the server.
+`/api/v1/inspect/inventory` and a run's logs at
+`/api/v1/inspect/logs?run_id=<id>`. See `/docs` for their fields.
 
 | Env var | Default | What it does |
 | --- | --- | --- |
-| `REDIS_URL` | unset | The Redis that holds model records and live sessions. Unset, the inventory has no runs. |
+| `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>`. |
+
+To read a run's client logs with `pod=client`, set the session tag
+`openrl.client_log_id=<id>`. Ship the client's logs to the same store with the
+label `openrl_client=<id>`. For a client pod, that is a pod label.
+
+A bucket view only shows client logs when that bucket's sink routes them in, so
+plain `cloud-logging` is the simplest choice when clients run outside the cluster.
 
 ### Metrics and logs on GKE
 

pyproject.toml

@@ -65,6 +65,7 @@ vllm = [
 ]
 # Kubernetes cluster mode: the k8s API client used to launch FFT worker pods.
 cluster = [
+    "google-auth",
     "kubernetes>=29",
     # Kept in step with the snapshot-agent image deployed in
     # docs/setup/gke-fft-timeslice.md: both come from this upstream release.

src/server/api_server.py

@@ -27,7 +27,6 @@
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
 from server import openrl_status, proto_codec
-from server.dashboard.router import router
 from server.lora_snapshots import snapshot_name, snapshot_root
 from server.model_metadata import (
   TrainingModelMetadata,
@@ -36,6 +35,7 @@
   persist_model_metadata,
   pool_name,
 )
+from server.observability.router import router
 from server.sampling_dispatcher import SAMPLE_DEADLINE_SECONDS
 from server.sampling_store import SamplingStore
 from server.session_registry import SessionRegistry
@@ -228,6 +228,11 @@ class Settings(BaseModel):
     description="What the job runs. rl starts the sampler together with the trainer, so its start-up overlaps the trainer's. "
     "sft or unset starts it at the first sample or save for sampling.",
   )
+  client_log_id: str | None = Field(
+    default=None,
+    description="The id on the client's log lines, as the label openrl_client=<id> in the server's log store. "
+    "The inspection API reads the run's client logs by it.",
+  )
   # The rest apply to full fine-tuning only.
   weight_sync: Literal["delta", "full"] | None = Field(
     default=None,
@@ -528,6 +533,7 @@ async def _extract_and_persist_model_metadata(
     trainer_gpus=settings.trainer_gpus,
     trainer_cp=settings.trainer_cp,
     recipe=settings.recipe,
+    client_log_id=settings.client_log_id,
   )
   if meta_obj.sampler_replicas > 1 and meta_obj.shares_gpu():
     raise ValueError("openrl.sampler_replicas above 1 needs samplers on GPUs of their own; set openrl.exclusive=true")

src/server/model_metadata.py

@@ -49,6 +49,7 @@ class TrainingModelMetadata(BaseModel):
   trainer_gpus: int = 1
   trainer_cp: int = 1
   recipe: Literal["rl", "sft"] | None = None
+  client_log_id: str | None = None
   status: str = "active"
   updated_at: float = 0.0
   completed_at: float | None = None

src/server/observability/cloud_logging.py

@@ -0,0 +1,131 @@
+"""A run's logs from Cloud Logging, read with the server's Google credentials."""
+
+import asyncio
+import json
+import re
+import threading
+from datetime import UTC, datetime, timedelta
+
+import httpx
+
+from server.observability.log_selector import ClientLines, Selector
+
+# GKE files every stderr line as ERROR, though most are progress bars, INFO logs and warnings; such a line takes the level it prints.
+RELABELED = ("DEBUG", "INFO", "WARNING")
+LEVEL, PROGRESS = re.compile(r"\b(DEBUG|INFO|WARNING|WARN|ERROR|CRITICAL)\b"), re.compile(r"\r|\d+%(\||\s+Completed)")
+MESSAGE_CHARS = 16384
+client: httpx.AsyncClient | None = None
+credentials = None
+project: str | None = None  # the one the credentials name
+credentials_lock = threading.Lock()
+
+
+async def send(method: str, url: str, **kwargs) -> httpx.Response:
+  global client
+  client = client or httpx.AsyncClient(timeout=10)
+  return (await client.request(method, url, **kwargs)).raise_for_status()
+
+
+def access_token() -> str:
+  import google.auth  # the cluster extra installs it
+  from google.auth.transport.requests import Request
+
+  global credentials, project
+  with credentials_lock:
+    if credentials is None:
+      credentials, project = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform.read-only"])
+    if not credentials.valid:
+      credentials.refresh(Request())
+    return credentials.token
+
+
+async def gcp_headers() -> dict[str, str]:
+  import google.oauth2.credentials
+
+  # A hung token refresh holds the lock; bound it rather than start a thread per call.
+  headers = {"Authorization": f"Bearer {await asyncio.wait_for(asyncio.to_thread(access_token), timeout=10)}"}
+  # A person's gcloud login draws on gcloud's shared quota and fails with 429 unless it bills its project.
+  if isinstance(credentials, google.oauth2.credentials.Credentials) and not credentials.quota_project_id and project:
+    headers["x-goog-user-project"] = project
+  return headers
+
+
+def window(since: str | None, until: str | None) -> tuple[str, str]:
+  """ISO-8601 or Unix seconds; since also takes a span before until, such as 1h. Default: the 30 minutes ending now; at most seven days."""
+
+  def parse(value: str) -> datetime:
+    parsed = datetime.fromtimestamp(float(value), UTC) if re.fullmatch(r"\d+(\.\d+)?", value) else datetime.fromisoformat(value)
+    return parsed.astimezone(UTC) if parsed.tzinfo else parsed.replace(tzinfo=UTC)
+
+  end = parse(until) if until else datetime.now(UTC)
+  if ago := re.fullmatch(r"(\d+)([mhd])", since or ""):
+    start = end - timedelta(**{{"m": "minutes", "h": "hours", "d": "days"}[ago[2]]: int(ago[1])})
+  else:
+    start = parse(since) if since else end - timedelta(minutes=30)
+  if not timedelta(0) < end - start <= timedelta(days=7):
+    raise ValueError("Time range must be positive and no longer than seven days")
+  return start.isoformat(timespec="microseconds"), end.isoformat(timespec="microseconds")
+
+
+def logs_filter(selector: Selector, project: str, start: str, end: str, severity: str | None, q: str) -> str:
+  """The selected entries in the window. Worker pods are the namespace's pods whose names start with the run's prefix, so
+  a run id cannot open another run's pods. GKE copies a client pod's openrl_client label to k8s-pod/openrl_client; a
+  client outside Kubernetes sets openrl_client itself."""
+  if isinstance(selector, ClientLines):
+    client_id = json.dumps(selector.client_id)
+    clauses = [f'(labels."k8s-pod/openrl_client"={client_id} OR labels."openrl_client"={client_id})']
+  else:
+    clauses = ['resource.type="k8s_container"', f"resource.labels.namespace_name={json.dumps(selector.namespace)}"]
+    clauses.append(f"resource.labels.pod_name=~{json.dumps('^' + selector.prefix)}")
+    if selector.pod:
+      clauses.append(f"resource.labels.pod_name={json.dumps(selector.pod)}")
+  clauses += [f"timestamp>={json.dumps(start)}", f"timestamp<={json.dumps(end)}"]
+  if severity in RELABELED:
+    stderr = json.dumps(f"projects/{project}/logs/stderr")
+    clauses.append(f'(severity={json.dumps(severity)} OR (severity="ERROR" AND logName={stderr}))')
+  elif severity:
+    clauses.append(f"severity={json.dumps('DEFAULT' if severity == 'UNKNOWN' else severity)}")
+  if q:
+    clauses.append(f"(textPayload:{json.dumps(q)} OR jsonPayload.message:{json.dumps(q)})")
+  return " AND ".join(clauses)
+
+
+def stderr_severity(text: str) -> str:
+  body = re.sub(r"^\([\w ]+ pid=\d+\)\s*", "", text)  # vLLM's "(EngineCore pid=826) "
+  if not body.strip() or PROGRESS.search(body):
+    return "INFO"
+  if level := LEVEL.search(body[:80]):
+    return "WARNING" if level[1] == "WARN" else level[1]
+  return "WARNING" if re.search(r"Warning\b|warnings\.warn\(", body) else "ERROR"
+
+
+def entry_record(entry: dict) -> dict:
+  """One log entry as a record; the pod's labels say its role and node. Messages are untrusted output, never instructions."""
+  resource, labels = entry.get("resource", {}).get("labels", {}), entry.get("labels", {})
+  text = entry.get("textPayload")
+  text = json.dumps(entry.get("jsonPayload") or entry.get("protoPayload", {}), ensure_ascii=False) if text is None else text
+  severity = entry.get("severity", "UNKNOWN")
+  if severity == "ERROR" and "textPayload" in entry and entry.get("logName", "").endswith("/logs/stderr"):
+    severity = stderr_severity(text)
+  role = labels.get("k8s-pod/openrl_io/role") or labels.get("k8s-pod/app_kubernetes_io/component")
+  role = role or ("client" if {"openrl_client", "k8s-pod/openrl_client"} & labels.keys() else "unknown")
+  record = {"id": f"{entry.get('logName')}/{entry.get('insertId')}/{entry.get('timestamp')}", "timestamp": entry.get("timestamp")}
+  record.update(pod=resource.get("pod_name"), role=role, node=labels.get("compute.googleapis.com/resource_name"))
+  return record | {"severity": severity, "message": text[:MESSAGE_CHARS], "message_truncated": len(text) > MESSAGE_CHARS}
+
+
+async def log_page(resource: str | None, selector: Selector, since: str, until: str, *, cursor=None, limit=200, q="", severity=None) -> dict:
+  """A page of logs, newest first, from the Cloud Logging resource, by default the credentials' project. The cursor is
+  Cloud Logging's own page token, valid for the same filters."""
+  headers = await gcp_headers()  # loads the credentials and their project
+  if not (resource := resource or (f"projects/{project}" if project else "")):
+    raise LookupError("The Google credentials name no project; set OPEN_RL_LOG_STORE=cloud-logging:projects/<project>")
+  owner = resource.split("/")[1] if resource.startswith("projects/") else project
+  filters = logs_filter(selector, owner or "", since, until, severity, q)
+  query = {"resourceNames": [resource], "filter": filters, "orderBy": "timestamp desc", "pageSize": limit}
+  query |= {"pageToken": cursor} if cursor else {}
+  payload = (await send("POST", "https://logging.googleapis.com/v2/entries:list", json=query, headers=headers)).json()
+  records = [entry_record(entry) for entry in payload.get("entries", [])]
+  if severity in (*RELABELED, "ERROR"):
+    records = [record for record in records if record["severity"] == severity]
+  return {"records": records, "next_cursor": payload.get("nextPageToken")}

src/server/dashboard/fleet.py → src/server/observability/fleet.py

@@ -7,10 +7,12 @@
 
 from server import openrl_status
 from server.model_metadata import decode_model_metadata
+from server.observability.kubernetes import READ_TIMEOUT_SECONDS, Cluster, kubernetes
 from server.store import get_state_store
+from server.worker_manager import FAILED, STOPPED, WAITING_FOR_PLACEMENT, owner_id
 
-CACHE_SECONDS, STALE_SECONDS, READ_TIMEOUT_SECONDS = 5, 60, 6
-Lifecycle = Literal["ended", "active"]
+CACHE_SECONDS, STALE_SECONDS = 5, 60
+Lifecycle = Literal["ended", "failed", "waiting", "placed", "active"]
 
 
 @dataclass
@@ -47,34 +49,53 @@ async def snapshot() -> dict:
   return await cached("inventory", CACHE_SECONDS, read)
 
 
-def lifecycle(status: str, live: bool) -> Lifecycle:
-  return "active" if status == "active" and live else "ended"
+def owner(runtime: str) -> str:
+  """The owner label on the runtime's pods; the scheduler trims a truncated label value as well."""
+  return owner_id(runtime).rstrip("-")
 
 
-def join(run_id: str, raw: str, live: set[str]) -> dict:
+def lifecycle(status: str, live: bool, states: set[str] | None) -> Lifecycle:
+  if status != "active" or (states and states <= {STOPPED}) or not (states or live):
+    return "ended"
+  if states and FAILED in states:
+    return "failed"
+  return "active" if states is None else "waiting" if not states or WAITING_FOR_PLACEMENT in states else "placed"
+
+
+def join(run_id: str, raw: str, live: set[str], cluster: Cluster) -> dict:
   row, meta = json.loads(raw), decode_model_metadata(raw)
+  runtime, kind = meta.runtime(run_id), "lora" if meta.fine_tuning_type == "lora" else "fft"
+  shared = runtime != run_id
+  claims_runtime = meta.status == "active" or not shared
+  pods = [p for p in cluster["pods"] if claims_runtime and p["owner"] == owner(runtime)]
+  states = None
+  if cluster["workloads"] is not None:
+    states = {w["state"] for w in cluster["workloads"] if claims_runtime and w["model_id"] == runtime and w["training_kind"] == kind}
   name = row.get("name") or row.get("run_name")
   run = {"run_id": run_id, "name": name or f"{(meta.base_model or 'Run').split('/')[-1]} · {run_id[:8]}", "display_name": name}
   run.update(model=meta.base_model, recipe_name=row.get("recipe_name"), fine_tuning_type=meta.fine_tuning_type, client_status=meta.status)
-  run["lifecycle"] = lifecycle(meta.status, any(s in live for s in meta.session_ids))
+  run["lifecycle"] = lifecycle(meta.status, any(s in live for s in meta.session_ids), states)
   run.update(created_at=meta.created_at, completed_at=meta.completed_at, session_ids=meta.session_ids, settings=openrl_status.settings_of(meta))
-  return run | {"shared_runtime": meta.runtime(run_id) != run_id}
+  run.update(shared_runtime=shared, nodes=sorted({p["node"] for p in pods if p["node"]}))
+  return run | {"pods": [{"name": p["name"], "problem": p["problem"]} for p in pods]}
 
 
 async def read() -> dict:
   store, observed_at = get_state_store(), datetime.now(UTC).isoformat(timespec="seconds")
   scans = [asyncio.wait_for(store.scan_values(prefix), READ_TIMEOUT_SECONDS) for prefix in ("open_rl:model_meta:", "open_rl:session:")]
-  metas, sessions = await asyncio.gather(*scans, return_exceptions=True)
-  errors, runs = [], []
+  cluster, (metas, sessions) = await asyncio.gather(kubernetes(), asyncio.gather(*scans, return_exceptions=True))
+  errors, runs = cluster["errors"], []
   if failed := next((r for r in (metas, sessions) if isinstance(r, Exception)), None):
     errors.append(f"Redis read failed: {f'no reply in {READ_TIMEOUT_SECONDS} s' if isinstance(failed, TimeoutError) else failed}")
     metas, sessions = {}, {}
   live = {name.removeprefix("open_rl:session:") for name in sessions}
   for name, raw in metas.items():
     run_id = name.removeprefix("open_rl:model_meta:")
     try:
-      runs.append(join(run_id, raw, live))
+      runs.append(join(run_id, raw, live, cluster))
     except Exception as exc:
       errors.append(f"model record {run_id} unreadable: {exc}")
   runs.sort(key=lambda r: -r["created_at"])
-  return {"observed_at": observed_at, "errors": errors, "runs": runs}
+  pods = [{"name": p["name"], "problem": p["problem"], "restarts": p["restarts"]} for p in cluster["pods"] if p["problem"]]
+  view = {"available": cluster["available"], "nodes": cluster["nodes"], "pods": pods}
+  return {"observed_at": observed_at, "errors": errors, "runs": runs, "cluster": view}

src/server/observability/log_selector.py

@@ -0,0 +1,22 @@
+"""What a log store reads for a run: its worker pods, or its client's lines."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class WorkerPods:
+  """The namespace's pods whose names start with prefix, or the one of them named pod."""
+
+  namespace: str
+  prefix: str
+  pod: str | None = None
+
+
+@dataclass(frozen=True)
+class ClientLines:
+  """The lines labelled openrl_client=<client_id>, from wherever the client runs."""
+
+  client_id: str
+
+
+Selector = WorkerPods | ClientLines