638d7c733f
GPU metrics polling, per-epoch RL metrics, OTLP tracing in training binaries, and three Grafana dashboards (training cockpit, traces, infra cockpit trace row). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
185 lines
9.1 KiB
Markdown
185 lines
9.1 KiB
Markdown
# Training Binary Observability + Grafana Dashboards Design
|
||
|
||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||
|
||
**Goal:** Instrument training binaries with GPU metrics polling, per-epoch Prometheus metrics, and OTLP tracing so we can validate whether the GPU is maxed out during real DQN/PPO training runs. Deploy Grafana dashboards to visualize everything.
|
||
|
||
**Architecture:** Training binaries poll `nvidia-smi` every 5s for GPU metrics, emit per-fold training metrics via existing Prometheus helpers, and export OTLP traces to Tempo. Grafana consumes Prometheus (metrics), Loki (logs), and Tempo (traces) with cross-linking.
|
||
|
||
**Tech Stack:** Rust tracing + tracing-opentelemetry (already in workspace), Prometheus client (already used), nvidia-smi CLI, Grafana 8.x dashboards (JSON ConfigMaps).
|
||
|
||
---
|
||
|
||
## Problem Statement
|
||
|
||
The training binaries (`train_baseline_rl`, `train_baseline_supervised`, `hyperopt_baseline_rl`, `hyperopt_baseline_supervised`) have a Prometheus metrics server running on port 9094 with 18 training metrics registered, but:
|
||
|
||
1. **Zero GPU metrics emitted** — `nvidia-smi` is called once at startup for batch sizing, never polled during training
|
||
2. **`train_baseline_rl.rs` is a metrics black box** — DQN/PPO trainers are delegated to but never call back into `common::metrics::training_metrics`. Only `record_data_load()` and `set_active_workers()` are called
|
||
3. **No OTLP tracing** — training binaries use plain `tracing_subscriber::fmt()`, not `init_observability()`
|
||
4. **No `OTEL_EXPORTER_OTLP_ENDPOINT`** in K8s training job template
|
||
5. **No Grafana dashboards** for training or traces (design exists, implementation doesn't)
|
||
|
||
## Data Flow
|
||
|
||
```
|
||
Training Binary → nvidia-smi poll (5s) → Prometheus gauges → /metrics :9094 → Prometheus scrape → Grafana
|
||
Training Binary → set_epoch_loss() etc. → Prometheus gauges → /metrics :9094 → Prometheus scrape → Grafana
|
||
Training Binary → init_observability() → OTLP gRPC :4317 → Tempo → Grafana (Traces)
|
||
Training Binary → stdout → containerd → Promtail → Loki → Grafana (linked via trace_id)
|
||
```
|
||
|
||
## Component 1: GPU Metrics Polling
|
||
|
||
New file: `crates/common/src/metrics/gpu_metrics.rs`
|
||
|
||
### Metrics Registered
|
||
|
||
| Metric Name | Type | Description |
|
||
|---|---|---|
|
||
| `foxhunt_gpu_utilization_percent` | Gauge | SM utilization 0-100 |
|
||
| `foxhunt_gpu_memory_used_bytes` | Gauge | VRAM used in bytes |
|
||
| `foxhunt_gpu_memory_total_bytes` | Gauge | VRAM total in bytes |
|
||
| `foxhunt_gpu_temperature_celsius` | Gauge | GPU temperature |
|
||
| `foxhunt_gpu_power_draw_watts` | Gauge | Power consumption |
|
||
|
||
### Implementation
|
||
|
||
- `poll_gpu_metrics()` — shells out to `nvidia-smi --query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw --format=csv,noheader,nounits`, parses CSV, sets gauges. MiB values from nvidia-smi are converted to bytes. Silently logs warning and skips update on parse failure (nvidia-smi not available on CPU-only machines).
|
||
- `start_gpu_poller(interval_secs: u64)` — spawns a daemon thread named `"gpu-metrics-poller"` that calls `poll_gpu_metrics()` in a loop with `std::thread::sleep(Duration::from_secs(interval_secs))`. Returns immediately. Thread is detached (no join handle needed).
|
||
|
||
### Integration
|
||
|
||
All 4 training binaries call `gpu_metrics::start_gpu_poller(5)` in `main()` right after `metrics_server::start_metrics_server(9094)`.
|
||
|
||
## Component 2: Wire Per-Epoch Metrics into train_baseline_rl.rs
|
||
|
||
Since DQN/PPO trainers return only a single `f64` (best validation loss), we instrument at the **fold level** in the binary, not inside the trainers.
|
||
|
||
### Per-Fold Metrics Calls
|
||
|
||
After each DQN/PPO fold completes:
|
||
```rust
|
||
metrics::set_epoch(model_name, fold_idx, fold_idx as f64); // fold as "epoch"
|
||
metrics::set_epoch_loss(model_name, fold_idx, best_val_loss);
|
||
metrics::set_validation_loss(model_name, fold_idx, best_val_loss);
|
||
metrics::set_iteration_seconds(model_name, fold_idx, fold_elapsed.as_secs_f64());
|
||
```
|
||
|
||
After data loading:
|
||
```rust
|
||
metrics::record_data_load(model_name, load_elapsed.as_secs_f64()); // already done
|
||
```
|
||
|
||
On checkpoint save (in the checkpoint callback):
|
||
```rust
|
||
metrics::record_checkpoint_save(model_name, fold_idx, ckpt_bytes.len(), save_elapsed);
|
||
```
|
||
|
||
### Hyperopt Binaries
|
||
|
||
`hyperopt_baseline_rl.rs` and `hyperopt_baseline_supervised.rs` — add `set_epoch()` and `set_epoch_loss()` per trial iteration. Lower priority since hyperopt is exploratory, but the metrics are cheap.
|
||
|
||
## Component 3: OTLP Tracing in Training Binaries
|
||
|
||
### Change
|
||
|
||
In all 4 training binary `main()` functions, replace:
|
||
```rust
|
||
tracing_subscriber::fmt()
|
||
.with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")))
|
||
.init();
|
||
```
|
||
|
||
With:
|
||
```rust
|
||
common::observability::init_observability(
|
||
"train_baseline_rl", // service name varies per binary
|
||
std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok().as_deref(),
|
||
);
|
||
```
|
||
|
||
When `OTEL_EXPORTER_OTLP_ENDPOINT` is unset, `init_observability` falls back to console-only logging (same as today). When set, spans flow to Tempo.
|
||
|
||
### K8s Training Job Template
|
||
|
||
Add to `infra/k8s/training/job-template.yaml` pod env:
|
||
```yaml
|
||
- name: OTEL_EXPORTER_OTLP_ENDPOINT
|
||
value: "http://tempo.foxhunt.svc.cluster.local:4317"
|
||
```
|
||
|
||
## Component 4: Grafana Training Cockpit Dashboard
|
||
|
||
New file: `infra/k8s/monitoring/dashboards/foxhunt-training-cockpit.json`
|
||
|
||
Provisioned as a ConfigMap, auto-loaded by Grafana's dashboard provisioner (same pattern as existing `foxhunt-cockpit.json`).
|
||
|
||
### Dashboard Layout (8 rows)
|
||
|
||
Per the existing design doc's Component 4 spec, using the actual metric names from `common/src/metrics/training_metrics.rs` and the new GPU metrics:
|
||
|
||
| Row | Focus | Panels |
|
||
|---|---|---|
|
||
| 1 | Active Jobs | active_workers gauge, fold progress |
|
||
| 2 | Training Progress | epoch_loss + validation_loss timeseries, batches_per_second, iteration_seconds |
|
||
| 3 | GPU Resources | gpu_utilization_percent, gpu_memory_used/total, gpu_temperature, gpu_power_draw |
|
||
| 4 | Model Quality | eval_accuracy, eval_f1, eval_precision, eval_recall |
|
||
| 5 | Checkpoints | checkpoint_saves_total, checkpoint_failures_total, checkpoint_duration_seconds, checkpoint_size_bytes |
|
||
| 6 | Data Pipeline | batches_processed, data_load_seconds |
|
||
| 7 | Errors | nan_detected_total, gradient_explosion_total, feature_errors_total |
|
||
| 8 | Timing | iteration_seconds histogram, data_load_seconds histogram |
|
||
|
||
### Template Variables
|
||
|
||
- `$model` — dropdown: `dqn`, `ppo`, `tft`, etc. Filters all panels by `model` label
|
||
- `$fold` — dropdown: `0`, `1`, `2`, etc. Filters by `fold` label
|
||
- `$interval` — auto interval for rate() queries
|
||
|
||
## Component 5: Grafana Traces Dashboard
|
||
|
||
New file: `infra/k8s/monitoring/dashboards/foxhunt-traces.json`
|
||
|
||
- Service map panel (Tempo node graph)
|
||
- Trace search table (service, operation, duration, status filters)
|
||
- Trace duration histogram
|
||
- Error rate by service timeseries
|
||
|
||
## Component 6: Existing Dashboard Updates
|
||
|
||
- `foxhunt-cockpit.json` — add one "Trace Health" row: spans/sec and error rate from Tempo metrics generator
|
||
- Service annotations — add Prometheus scrape annotations to 8 service deployment YAMLs (per existing design doc Component 1)
|
||
|
||
## Component 7: Service OTLP Wiring
|
||
|
||
Per existing design doc Component 2: Add `init_tracing()` call to 8 service `main()` functions. Add `#[instrument(skip_all)]` to ~20-30 key async functions (gRPC handlers, DB queries, ML inference paths). No instrumentation on hot-path tick processing.
|
||
|
||
## File Changes Summary
|
||
|
||
| Category | Files | Change |
|
||
|----------|-------|--------|
|
||
| New GPU metrics module | `common/src/metrics/gpu_metrics.rs` | nvidia-smi polling + 5 gauges |
|
||
| Training binary metrics | `ml/examples/train_baseline_rl.rs` | Wire per-fold metrics calls |
|
||
| Training binary OTLP | 4 training binary `main()` functions | Replace fmt with init_observability |
|
||
| K8s training job | `infra/k8s/training/job-template.yaml` | Add OTEL endpoint env var |
|
||
| New dashboards | 2 JSON files | training-cockpit, traces |
|
||
| Updated dashboards | 1 JSON file | cockpit (trace row) |
|
||
| K8s service annotations | 8 service YAMLs | Prometheus scrape annotations |
|
||
| Service OTLP | 8 service main.rs | Add init_tracing() |
|
||
| Instrumentation | ~20-30 key functions | `#[instrument(skip_all)]` |
|
||
|
||
## What We Don't Change
|
||
|
||
- No changes to DQN/PPO trainer internals (instrument at fold level in binaries)
|
||
- No changes to hot-path tick processing (no `#[instrument]` on SIMD/RDTSC paths)
|
||
- No new Rust dependencies (tracing-opentelemetry already in workspace)
|
||
- No changes to Tempo, Loki, or Prometheus deployment configs
|
||
- Existing infra cockpit layout stays as-is (just one row added)
|
||
|
||
## Risks and Mitigations
|
||
|
||
- **nvidia-smi subprocess overhead**: 5s poll interval, `nvidia-smi` takes ~50ms. 1% overhead on a 5s window. Acceptable.
|
||
- **nvidia-smi unavailable**: CPU-only dev machines don't have it. Poller logs a warning once and exits the thread gracefully.
|
||
- **Trace volume from training**: Training runs are ~30-60min with ~1000 folds. At one span per fold, volume is negligible (~1000 spans per job).
|
||
- **Metric cardinality**: `model` × `fold` labels. Max ~10 models × ~50 folds = 500 series. Low cardinality.
|