feat(observability): show GPU activity from placements and DCGM

docs/configuration.md

@@ -127,7 +127,9 @@ for. On one machine every pool has one replica, in process. See
 
 The API server serves a read-only inventory of runs at
 `/api/v1/inspect/inventory` and a run's logs at
-`/api/v1/inspect/logs?run_id=<id>`. See `/docs` for their fields.
+`/api/v1/inspect/logs?run_id=<id>`. `/api/v1/inspect/activity` shows which
+`Workload`s held which GPUs in a time window, and `/api/v1/inspect/query` reads
+those GPUs' DCGM metrics. See `/docs` for their fields.
 
 `placements` lists the `Workload`s that hold GPUs now, from their status in
 Kubernetes. Each entry has the node, pod, GPU UUIDs, `placed_at` in Unix
@@ -139,6 +141,7 @@ days.
 | 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>`. |
+| `OPEN_RL_PROMETHEUS_URL` | unset | A Prometheus-compatible query endpoint with the DCGM exporter's metrics, such as an in-cluster Managed Service for Prometheus frontend at `http://gmp-frontend:9090`. Unset turns metric reads off. |
 
 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
@@ -153,17 +156,20 @@ The `lora`, `fft` and `distributed-fft-timeslice` bundles ship an NVIDIA DCGM
 exporter (`k8s/deploy/dcgm`). GKE's managed DCGM exporter skips DRA node pools,
 so without it those nodes report no GPU metrics.
 
-The steps below are for the inspection API's metrics and logs reads. It reads
-Cloud Monitoring and Cloud Logging as the API server's Kubernetes
-ServiceAccount, `open-rl-sa`, through Workload Identity. Enable Workload
-Identity when you create the cluster (`--workload-pool=<project>.svc.id.goog`).
-Enabling it on an existing cluster replaces its nodes.
+The inspection API reads GPU metrics from the Prometheus-compatible endpoint in
+`OPEN_RL_PROMETHEUS_URL`, such as the in-cluster Managed Service for Prometheus
+frontend. It reads Cloud Logging as the API server's Kubernetes ServiceAccount,
+`open-rl-sa`, through Workload Identity. Enable Workload Identity when you
+create the cluster (`--workload-pool=<project>.svc.id.goog`). Enabling it on an
+existing cluster replaces its nodes.
 
 ```bash
 PRINCIPAL=principal://iam.googleapis.com/projects/<number>/locations/global/workloadIdentityPools/<project>.svc.id.goog/subject/ns/openrl-system/sa/open-rl-sa
-gcloud projects add-iam-policy-binding <project> --role=roles/monitoring.viewer --member="$PRINCIPAL"
 gcloud projects add-iam-policy-binding <project> --role=roles/logging.viewer --member="$PRINCIPAL"
 ```
 
+Grant `roles/monitoring.viewer` the same way only if `OPEN_RL_PROMETHEUS_URL`
+is Cloud Monitoring's own Prometheus API.
+
 `<number>` is the project number. If you read logs through a log bucket view,
 also grant `roles/logging.viewAccessor` on that view.

src/server/observability/fleet.py

@@ -12,6 +12,7 @@
 from server.worker_manager import FAILED, STOPPED, WAITING_FOR_PLACEMENT, owner_id
 
 CACHE_SECONDS, STALE_SECONDS = 5, 60
+PLACEMENT_FIELDS = ("name", "role", "node", "pod", "gpu_uuids", "placed_at", "released_at", "training_kind")
 Lifecycle = Literal["ended", "failed", "waiting", "placed", "active"]
 
 
@@ -80,13 +81,17 @@ def join(run_id: str, raw: str, live: set[str], cluster: Cluster) -> dict:
   return run | {"pods": [{"name": p["name"], "problem": p["problem"]} for p in pods]}
 
 
+def redis_failed(exc: Exception) -> str:
+  return f"Redis read failed: {f'no reply in {READ_TIMEOUT_SECONDS} s' if isinstance(exc, TimeoutError) else exc}"
+
+
 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:")]
   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}")
+    errors.append(redis_failed(failed))
     metas, sessions = {}, {}
   live = {name.removeprefix("open_rl:session:") for name in sessions}
   for name, raw in metas.items():
@@ -120,3 +125,79 @@ def placements(workloads: list[dict] | None, metas: dict[str, str]) -> list[dict
     for w in workloads or []
     if w["holds_gpus"]
   ]
+
+
+async def holdings() -> dict:
+  return await cached("holdings", CACHE_SECONDS, read_holdings)
+
+
+async def read_holdings() -> dict:
+  """Every placement known: live ones from Workload status, released ones from open_rl:placements:<runtime>; and the runs
+  that could have used them."""
+  now = time.time()
+  history = asyncio.ensure_future(asyncio.wait_for(read_history(), READ_TIMEOUT_SECONDS))
+  cluster = await kubernetes()
+  live = [{f: w[f] for f in PLACEMENT_FIELDS} | {"runtime": w["model_id"]} for w in cluster["workloads"] or [] if w["holds_gpus"]]
+  try:
+    runs, released, errors = await history
+  except Exception as exc:
+    runs, released, errors = [], [], [redis_failed(exc)]
+  return {"observed_at": now, "runs": runs, "placements": live + released, "errors": cluster["errors"] + errors}
+
+
+async def read_history() -> tuple[list[dict], list[dict], list[str]]:
+  """From Redis: the runs, the placements their runtimes released, and an error per unreadable history."""
+  store, runs, released, errors = get_state_store(), [], [], []
+  for name, raw in (await store.scan_values("open_rl:model_meta:")).items():
+    try:
+      meta, run_id = decode_model_metadata(raw), name.removeprefix("open_rl:model_meta:")
+    except Exception:
+      continue  # the inventory reports it
+    kind = "lora" if meta.fine_tuning_type == "lora" else "fft"
+    runs.append({"run_id": run_id, "runtime": meta.runtime(run_id), "kind": kind, "created_at": meta.created_at, "completed_at": meta.completed_at})
+  runtimes = sorted({r["runtime"] for r in runs})
+  for runtime, records in zip(runtimes, await asyncio.gather(*(store.recent(f"open_rl:placements:{r}") for r in runtimes))):
+    try:
+      released += [{f: record[f] for f in PLACEMENT_FIELDS} | {"runtime": runtime} for record in map(json.loads, records)]
+    except (ValueError, KeyError) as exc:
+      errors.append(f"placement history of {runtime} unreadable: {exc!r}")
+  return runs, released, errors
+
+
+def lanes(held: dict, low: float, high: float, run_id: str | None = None, node_id: str | None = None) -> list[dict]:
+  """One lane per Workload and the GPUs it held in [low, high]: the intervals it held them, clipped to the window, and the
+  runs on its runtime that lived then. Lanes that held one GPU at overlapping times are shared_gpu."""
+  now, found = held["observed_at"], {}
+  for p in held["placements"]:
+    start, end = p["placed_at"], p["released_at"] or now
+    if start >= high or end <= low:
+      continue
+    fresh = {k: p[k] for k in ("name", "role", "node", "pod", "gpu_uuids")} | {"run_ids": set(), "intervals": [], "ended": True}
+    lane = found.setdefault((p["name"], p["node"], tuple(p["gpu_uuids"])), fresh)
+    lane["run_ids"].update(
+      r["run_id"]
+      for r in held["runs"]
+      if (r["runtime"], r["kind"]) == (p["runtime"], p["training_kind"]) and r["created_at"] <= end and (r["completed_at"] or now) >= start
+    )
+    lane["intervals"].append([max(low, start), min(high, end)])
+    lane["ended"] &= p["released_at"] is not None
+  result = [lane | {"run_ids": sorted(lane["run_ids"]), "intervals": merge(lane["intervals"])} for lane in found.values()]
+  for lane in result:
+    lane["shared_gpu"] = any(
+      other is not lane and set(lane["gpu_uuids"]) & set(other["gpu_uuids"]) and overlap(lane["intervals"], other["intervals"]) for other in result
+    )
+  return [lane for lane in result if (not run_id or run_id in lane["run_ids"]) and (not node_id or lane["node"] == node_id)]
+
+
+def merge(spans: list[list[float]]) -> list[list[float]]:
+  merged: list[list[float]] = []
+  for start, end in sorted(spans):
+    if merged and start <= merged[-1][1]:
+      merged[-1][1] = max(end, merged[-1][1])
+    else:
+      merged.append([start, end])
+  return merged
+
+
+def overlap(a: list[list[float]], b: list[list[float]]) -> bool:
+  return any(a0 < b1 and b0 < a1 for a0, a1 in a for b0, b1 in b)

src/server/observability/metrics.py

@@ -0,0 +1,43 @@
+"""Named PromQL queries over DCGM's GPU metrics, read from the Prometheus-compatible endpoint at OPEN_RL_PROMETHEUS_URL."""
+
+import json
+import math
+import os
+import re
+from string import Template
+
+import httpx
+
+DCGM_SECONDS = 10  # how often k8s/deploy/dcgm collects; a finer step only repeats samples
+# Per name, the metric each query reads and the query, tried in order until one has series. $gpus matches the GPU UUIDs,
+# $step is the seconds between points.
+QUERIES = {
+  # The percent of each step the graphics engine was busy, so points=1 is the window's duty cycle. GPU_UTIL, where the
+  # profiling metric is missing, samples only about the last second of each collection interval.
+  "gpu_busy": {
+    "DCGM_FI_PROF_GR_ENGINE_ACTIVE": "max by (Hostname, UUID, gpu) (100 * avg_over_time(DCGM_FI_PROF_GR_ENGINE_ACTIVE{UUID=~$gpus}[${step}s]))",
+    "DCGM_FI_DEV_GPU_UTIL": "max by (Hostname, UUID, gpu) (avg_over_time(DCGM_FI_DEV_GPU_UTIL{UUID=~$gpus}[${step}s]))",
+  },
+  "gpu_util": {"DCGM_FI_DEV_GPU_UTIL": "max by (Hostname, UUID, gpu) (DCGM_FI_DEV_GPU_UTIL{UUID=~$gpus})"},
+  "gpu_memory_mib": {"DCGM_FI_DEV_FB_USED": "max by (Hostname, UUID, gpu) (DCGM_FI_DEV_FB_USED{UUID=~$gpus})"},
+}
+
+
+async def get(url: str, params: dict) -> dict:
+  async with httpx.AsyncClient(timeout=10) as client:
+    return (await client.get(url, params=params)).raise_for_status().json()
+
+
+async def query(name: str, gpus: list[str], low: float, high: float, points: int) -> dict:
+  """The named query over these GPUs in [low, high], from the first of its metrics that has series."""
+  url, step = os.environ["OPEN_RL_PROMETHEUS_URL"].rstrip("/"), max(DCGM_SECONDS, int((high - low) / points))
+  matcher = json.dumps("|".join(re.escape(uuid) for uuid in gpus))
+  for metric, template in QUERIES[name].items():
+    text = Template(template).substitute(gpus=matcher, step=step)
+    params = {"query": text, "start": high if points == 1 else low, "end": high, "step": step}
+    rows = (await get(f"{url}/api/v1/query_range", params))["data"]["result"]
+    series = [{"labels": r["metric"], "points": [[float(t), float(v)] for t, v in r["values"] if math.isfinite(float(v))]} for r in rows]
+    answer = {"metric": metric, "query": text, "step": step, "series": series}
+    if series:
+      return answer
+  return answer

src/server/observability/router.py

@@ -1,11 +1,12 @@
 import os
-from typing import Annotated
+from datetime import datetime
+from typing import Annotated, Literal
 
 import httpx
 from fastapi import APIRouter, HTTPException, Query
 
 from server.model_metadata import decode_model_metadata
-from server.observability import cloud_logging, fleet
+from server.observability import cloud_logging, fleet, metrics
 from server.observability.kubernetes import namespace
 from server.observability.log_selector import ClientLines, WorkerPods
 from server.store import get_state_store
@@ -28,6 +29,55 @@ async def inventory():
   return await fleet.snapshot()
 
 
+def bounds(since: str | None, until: str | None) -> tuple[str, str]:
+  try:
+    return cloud_logging.window(since, until)
+  except (ValueError, OverflowError, OSError) as exc:
+    raise HTTPException(400, "Use ISO-8601 or Unix-second since/until and a positive window of at most seven days") from exc
+
+
+def epoch(bound: str) -> float:
+  return datetime.fromisoformat(bound).timestamp()
+
+
+@router.get("/api/v1/inspect/activity")
+async def activity(since: Bound = None, until: Bound = None, run_id: str | None = None, node_id: str | None = None):
+  """One lane per Workload that held GPUs in the window, a run's or a node's: its role, node, pod, GPU UUIDs, the runs on
+  its runtime that lived then, the intervals it held them in Unix seconds clipped to the window, and whether it let go.
+  shared_gpu means another lane held one of its GPUs at an overlapping time, so its GPUs' metrics are not its own."""
+  since, until = bounds(since, until)
+  held = await fleet.holdings()
+  return {"since": since, "until": until, "errors": held["errors"], "lanes": fleet.lanes(held, epoch(since), epoch(until), run_id, node_id)}
+
+
+@router.get("/api/v1/inspect/query")
+async def query(
+  name: Literal[tuple(metrics.QUERIES)],
+  run_id: str | None = None,
+  node_id: str | None = None,
+  since: Bound = None,
+  until: Bound = None,
+  points: int = Query(300, ge=1, le=1000),
+):
+  """A named DCGM query over the GPUs a run's or a node's placements held in the window, from OPEN_RL_PROMETHEUS_URL:
+  series of [Unix seconds, value] `step` seconds apart. points=1 evaluates once over the whole window. metric is the DCGM
+  metric that answered; gpu_busy falls back to DCGM_FI_DEV_GPU_UTIL where DCGM_FI_PROF_GR_ENGINE_ACTIVE has no data."""
+  since, until = bounds(since, until)
+  result = {"name": name, "since": since, "until": until, "series": []}
+  if not os.getenv("OPEN_RL_PROMETHEUS_URL"):
+    return {**result, "available": False, "error": "No metrics source configured. Set OPEN_RL_PROMETHEUS_URL."}
+  low, high = epoch(since), epoch(until)
+  gpus = sorted({uuid for lane in fleet.lanes(await fleet.holdings(), low, high, run_id, node_id) for uuid in lane["gpu_uuids"]})
+  if not gpus:
+    return {**result, "available": True, "error": "No GPUs held by this run or node in this window"}
+  try:
+    return {**result, **await metrics.query(name, gpus, low, high, points), "available": True}
+  except httpx.HTTPStatusError as exc:
+    return {**result, "available": False, "error": f"The metrics source refused the query (HTTP {exc.response.status_code})"}
+  except Exception as exc:
+    return {**result, "available": False, "error": f"The metrics source is unavailable: {type(exc).__name__}"}
+
+
 @router.get("/api/v1/inspect/logs")
 async def logs(
   run_id: str,
@@ -40,10 +90,7 @@ async def logs(
   cursor: Annotated[str | None, Query(max_length=16384, description="next_cursor of a page with the same filters and window")] = None,
 ):
   """Newest first, from the OPEN_RL_LOG_STORE log store: the run's worker pods, or with pod=client its client's lines."""
-  try:
-    since, until = cloud_logging.window(since, until)
-  except (ValueError, OverflowError, OSError) as exc:
-    raise HTTPException(400, "Use ISO-8601 or Unix-second since/until and a positive window of at most seven days") from exc
+  since, until = bounds(since, until)
   if (raw := await get_state_store().get_value(f"open_rl:model_meta:{run_id}")) is None:
     raise HTTPException(404, f"No run {run_id}")
   meta = decode_model_metadata(raw)

src/server/store.py

@@ -280,6 +280,11 @@ def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds:
     ttl_seconds after the last push."""
     pass
 
+  @abstractmethod
+  async def recent(self, key: str) -> list[str]:
+    """The list push_recent_sync keeps at key, newest first."""
+    pass
+
 
 class InMemoryStateStore(StateStore):
   def __init__(self):
@@ -324,6 +329,9 @@ async def scan_values(self, prefix: str) -> dict[str, str]:
   def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: float) -> None:
     self.lists[key] = (values[::-1] + self.lists.get(key, []))[:keep]
 
+  async def recent(self, key: str) -> list[str]:
+    return list(self.lists.get(key, []))
+
 
 class RedisStateStore(StateStore):
   def __init__(self, client: redis.Redis, sync_client: sync_redis.Redis):
@@ -359,6 +367,9 @@ async def scan_values(self, prefix: str) -> dict[str, str]:
   def push_recent_sync(self, key: str, values: list[str], keep: int, ttl_seconds: float) -> None:
     self.sync_redis.pipeline().lpush(key, *values).ltrim(key, 0, keep - 1).expire(key, int(ttl_seconds)).execute()
 
+  async def recent(self, key: str) -> list[str]:
+    return await self.redis.lrange(key, 0, -1)
+
 
 @cache
 def redis_client(redis_url: str) -> redis.Redis: