docs: training observability implementation plan (reduced scope)
DCGM exporter + dashboards already exist. Actual work: 3 tasks — per-fold metrics in train_baseline_rl, OTLP tracing in all 6 binaries, OTEL env var in K8s job template. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
291
docs/plans/2026-03-01-training-observability-implementation.md
Normal file
291
docs/plans/2026-03-01-training-observability-implementation.md
Normal file
@@ -0,0 +1,291 @@
|
||||
# Training Binary Observability Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** Wire per-fold Prometheus metrics and OTLP tracing into training binaries so Grafana dashboards (already deployed) show real training + GPU data.
|
||||
|
||||
**Architecture:** Training binaries already have Prometheus metrics registered and a metrics server on :9094. DCGM exporter (already deployed as DaemonSet) provides GPU metrics. We wire the missing per-fold metric calls, replace `tracing_subscriber::fmt()` with `init_observability()` for OTLP traces, and add the OTEL env var to the K8s job template.
|
||||
|
||||
**Tech Stack:** Rust tracing + tracing-opentelemetry (already in workspace), common::metrics::training_metrics (18 metrics registered), common::observability (OTLP init), DCGM exporter (already deployed).
|
||||
|
||||
---
|
||||
|
||||
### Scope Reduction from Design Doc
|
||||
|
||||
The design doc (`2026-03-01-training-observability-design.md`) proposed 7 components. During research we discovered:
|
||||
|
||||
- **Component 1 (GPU Polling) — SKIP**: DCGM exporter DaemonSet (`infra/k8s/monitoring/dcgm-exporter.yaml`) already provides `dcgm_gpu_utilization`, `dcgm_fb_used`, `dcgm_gpu_temp`, `dcgm_power_usage` on GPU nodes.
|
||||
- **Components 4-6 (Dashboards) — SKIP**: All dashboards already exist in `infra/k8s/monitoring/dashboards/` — `foxhunt-training-cockpit.json`, `foxhunt-gpu-training.json`, `foxhunt-traces.json`, `foxhunt-cockpit.json`. Import script (`import.sh`) handles ConfigMap deployment.
|
||||
- **Component 7 (Service OTLP) — SKIP**: Services already have `init_tracing()`. Out of scope for training observability.
|
||||
|
||||
**Remaining work (3 tasks):**
|
||||
1. Wire per-fold metrics in `train_baseline_rl.rs` (Component 2)
|
||||
2. Replace `tracing_subscriber::fmt()` with OTLP-capable init in all 6 training binaries (Component 3)
|
||||
3. Add `OTEL_EXPORTER_OTLP_ENDPOINT` env var to K8s job template (Component 3, K8s part)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Wire Per-Fold Metrics into train_baseline_rl.rs
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/examples/train_baseline_rl.rs:636-676` (DQN + PPO fold match blocks)
|
||||
|
||||
**Context:** The fold loop at line 584 iterates `windows`. Each fold calls `train_dqn_fold()` (line 638) and/or `train_ppo_fold()` (line 660), which return `Ok(best_loss)`. Currently only `record_data_load()` is called. The training cockpit dashboard (`foxhunt-training-cockpit.json`) expects `foxhunt_training_current_epoch`, `foxhunt_training_epoch_loss`, `foxhunt_training_validation_loss`, `foxhunt_training_iteration_seconds` with `model` and `fold` labels.
|
||||
|
||||
**Step 1: Add fold timing + metric calls to DQN fold block**
|
||||
|
||||
At `train_baseline_rl.rs:636`, wrap the DQN block with timing and add metric calls after `Ok(best_loss)`:
|
||||
|
||||
```rust
|
||||
// Train DQN
|
||||
if train_dqn {
|
||||
let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
|
||||
let fold_start = std::time::Instant::now();
|
||||
let fold_str = fold_idx.to_string();
|
||||
match train_dqn_fold(
|
||||
window.fold,
|
||||
&train_norm,
|
||||
&val_norm,
|
||||
&train_bars_aligned,
|
||||
&val_bars_aligned,
|
||||
args,
|
||||
&args.output_dir,
|
||||
&hp,
|
||||
) {
|
||||
Ok(best_loss) => {
|
||||
let elapsed = fold_start.elapsed().as_secs_f64();
|
||||
metrics::set_epoch("dqn", &fold_str, fold_idx as f64);
|
||||
metrics::set_epoch_loss("dqn", &fold_str, best_loss);
|
||||
metrics::set_validation_loss("dqn", &fold_str, best_loss);
|
||||
metrics::set_iteration_seconds("dqn", &fold_str, elapsed);
|
||||
dqn_results.push((window.fold, best_loss));
|
||||
}
|
||||
Err(e) => {
|
||||
error!(" [DQN] Fold {} failed: {}", window.fold, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Step 2: Add fold timing + metric calls to PPO fold block**
|
||||
|
||||
Same pattern at `train_baseline_rl.rs:658`:
|
||||
|
||||
```rust
|
||||
// Train PPO
|
||||
if train_ppo {
|
||||
let hp = load_hyperopt_params(&args.hyperopt_params, "ppo");
|
||||
let fold_start = std::time::Instant::now();
|
||||
let fold_str = fold_idx.to_string();
|
||||
match train_ppo_fold(
|
||||
window.fold,
|
||||
&train_norm,
|
||||
&val_norm,
|
||||
&train_bars_aligned,
|
||||
&val_bars_aligned,
|
||||
args,
|
||||
&args.output_dir,
|
||||
&hp,
|
||||
) {
|
||||
Ok(best_loss) => {
|
||||
let elapsed = fold_start.elapsed().as_secs_f64();
|
||||
metrics::set_epoch("ppo", &fold_str, fold_idx as f64);
|
||||
metrics::set_epoch_loss("ppo", &fold_str, best_loss);
|
||||
metrics::set_validation_loss("ppo", &fold_str, best_loss);
|
||||
metrics::set_iteration_seconds("ppo", &fold_str, elapsed);
|
||||
ppo_results.push((window.fold, best_loss));
|
||||
}
|
||||
Err(e) => {
|
||||
error!(" [PPO] Fold {} failed: {}", window.fold, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Step 3: Verify it compiles**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --example train_baseline_rl`
|
||||
Expected: compiles with 0 errors
|
||||
|
||||
**Step 4: Run existing tests**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- --test-threads=4`
|
||||
Expected: all 2390 tests pass (metrics calls don't affect test behavior)
|
||||
|
||||
**Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/examples/train_baseline_rl.rs
|
||||
git commit -m "feat(ml): wire per-fold Prometheus metrics into train_baseline_rl
|
||||
|
||||
Emit set_epoch, set_epoch_loss, set_validation_loss, set_iteration_seconds
|
||||
after each DQN/PPO fold completes. These feed the training cockpit
|
||||
Grafana dashboard (foxhunt-training-cockpit.json)."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Replace tracing_subscriber::fmt() with OTLP-Capable Init in Training Binaries
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/examples/train_baseline_rl.rs:747-754` (main fn)
|
||||
- Modify: `crates/ml/examples/train_baseline_supervised.rs` (main fn)
|
||||
- Modify: `crates/ml/examples/evaluate_baseline.rs` (main fn)
|
||||
- Modify: `crates/ml/examples/hyperopt_baseline_rl.rs` (main fn)
|
||||
- Modify: `crates/ml/examples/hyperopt_baseline_supervised.rs` (main fn)
|
||||
- Modify: `crates/ml/examples/evaluate_supervised.rs` (main fn)
|
||||
|
||||
**Context:** `common::observability::init_observability()` is `async fn` but its body is 100% sync (uses `connect_lazy()`). Training binary `main()` is sync `fn main() -> Result<()>`. Rather than pulling in a tokio runtime just for init, we replicate the subscriber stack inline using the already-public `build_otel_tracer()` (which IS sync).
|
||||
|
||||
The pattern: check `OTEL_EXPORTER_OTLP_ENDPOINT` env var. If set, build OTLP layer via `build_otel_tracer()`. If unset, just use fmt layer (same as today). This way dev machines (no Tempo) keep working.
|
||||
|
||||
**Step 1: Update train_baseline_rl.rs main()**
|
||||
|
||||
Replace lines 747-754:
|
||||
|
||||
```rust
|
||||
fn main() -> Result<()> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||||
)
|
||||
.init();
|
||||
```
|
||||
|
||||
With:
|
||||
|
||||
```rust
|
||||
fn main() -> Result<()> {
|
||||
// Initialize tracing with optional OTLP export to Tempo
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
|
||||
let fmt_layer = tracing_subscriber::fmt::layer()
|
||||
.json()
|
||||
.with_target(true)
|
||||
.with_current_span(true)
|
||||
.with_span_list(false);
|
||||
|
||||
let otel_layer = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok().and_then(|endpoint| {
|
||||
let config = common::observability::TracingConfig {
|
||||
service_name: "train_baseline_rl".to_string(),
|
||||
otlp_endpoint: endpoint,
|
||||
enable_export: true,
|
||||
};
|
||||
common::observability::build_otel_tracer(&config)
|
||||
.map(|tracer| tracing_opentelemetry::layer().with_tracer(tracer))
|
||||
});
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
.with(fmt_layer)
|
||||
.with(otel_layer)
|
||||
.init();
|
||||
```
|
||||
|
||||
**Step 2: Update all other 5 binaries with the same pattern**
|
||||
|
||||
Each binary gets the same subscriber block, only `service_name` changes:
|
||||
- `train_baseline_supervised.rs` → `"train_baseline_supervised"`
|
||||
- `evaluate_baseline.rs` → `"evaluate_baseline"`
|
||||
- `hyperopt_baseline_rl.rs` → `"hyperopt_baseline_rl"`
|
||||
- `hyperopt_baseline_supervised.rs` → `"hyperopt_baseline_supervised"`
|
||||
- `evaluate_supervised.rs` → `"evaluate_supervised"`
|
||||
|
||||
For `hyperopt_baseline_rl.rs` and `hyperopt_baseline_supervised.rs`, which use `.with_max_level(Level::INFO).with_target(false)` — replace with the env-filter + JSON fmt pattern (same as above). These binaries don't need different filtering.
|
||||
|
||||
**Step 3: Verify `tracing-opentelemetry` is available to ml crate**
|
||||
|
||||
The `ml` crate depends on `common` (which has `tracing-opentelemetry`), but the examples need `tracing-opentelemetry` directly for the `layer()` call. Check if it's already a dev-dependency. If not, add it:
|
||||
|
||||
```toml
|
||||
# In crates/ml/Cargo.toml [dev-dependencies]
|
||||
tracing-opentelemetry.workspace = true
|
||||
```
|
||||
|
||||
Also verify `tracing-subscriber` has the `json` feature (for `.json()` layer). Check workspace config.
|
||||
|
||||
**Step 4: Verify all 6 compile**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --examples`
|
||||
Expected: compiles with 0 errors
|
||||
|
||||
**Step 5: Run tests**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- --test-threads=4`
|
||||
Expected: all 2390 tests pass
|
||||
|
||||
**Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml/examples/train_baseline_rl.rs \
|
||||
crates/ml/examples/train_baseline_supervised.rs \
|
||||
crates/ml/examples/evaluate_baseline.rs \
|
||||
crates/ml/examples/hyperopt_baseline_rl.rs \
|
||||
crates/ml/examples/hyperopt_baseline_supervised.rs \
|
||||
crates/ml/examples/evaluate_supervised.rs \
|
||||
crates/ml/Cargo.toml
|
||||
git commit -m "feat(ml): add OTLP tracing to all 6 training binaries
|
||||
|
||||
Replace tracing_subscriber::fmt() with JSON fmt + optional OTLP layer.
|
||||
When OTEL_EXPORTER_OTLP_ENDPOINT is set, spans flow to Tempo.
|
||||
When unset, logs-only mode (same as before).
|
||||
|
||||
No new runtime dependencies — tracing-opentelemetry already in workspace."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Add OTEL Env Var to K8s Training Job Template
|
||||
|
||||
**Files:**
|
||||
- Modify: `infra/k8s/training/job-template.yaml:140-146` (training container env block)
|
||||
|
||||
**Step 1: Add OTEL_EXPORTER_OTLP_ENDPOINT env var**
|
||||
|
||||
In the `training` container's `env` section (after `SQLX_OFFLINE`), add:
|
||||
|
||||
```yaml
|
||||
- name: OTEL_EXPORTER_OTLP_ENDPOINT
|
||||
value: "http://tempo.foxhunt.svc.cluster.local:4317"
|
||||
```
|
||||
|
||||
**Step 2: Verify YAML is valid**
|
||||
|
||||
Run: `python3 -c "import yaml; yaml.safe_load(open('infra/k8s/training/job-template.yaml'))"`
|
||||
Expected: no errors
|
||||
|
||||
**Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add infra/k8s/training/job-template.yaml
|
||||
git commit -m "feat(infra): add OTEL_EXPORTER_OTLP_ENDPOINT to training job template
|
||||
|
||||
Training binaries now export OTLP traces to Tempo when this env var is
|
||||
set. Traces appear in foxhunt-traces Grafana dashboard."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Verification Checklist
|
||||
|
||||
After all 3 tasks:
|
||||
|
||||
1. `SQLX_OFFLINE=true cargo check -p ml --examples` — all 6 binaries compile
|
||||
2. `SQLX_OFFLINE=true cargo test -p ml --lib -- --test-threads=4` — 2390 tests pass
|
||||
3. `SQLX_OFFLINE=true cargo clippy -p ml --examples -- -D warnings` — 0 warnings
|
||||
4. `python3 -c "import yaml; yaml.safe_load(open('infra/k8s/training/job-template.yaml'))"` — valid YAML
|
||||
5. Grafana dashboards already deployed — no action needed
|
||||
6. DCGM exporter already deployed — no action needed
|
||||
|
||||
## What This Enables
|
||||
|
||||
After deploying, a training job on the GPU node will:
|
||||
- Emit `foxhunt_training_epoch_loss`, `foxhunt_training_validation_loss`, `foxhunt_training_iteration_seconds` per fold → visible in Training Cockpit dashboard
|
||||
- DCGM provides `dcgm_gpu_utilization`, `dcgm_fb_used` → visible in GPU Training dashboard
|
||||
- OTLP traces flow to Tempo → visible in Traces dashboard
|
||||
- Cross-linking: Grafana links traces ↔ logs via `trace_id` field in JSON logs
|
||||
Reference in New Issue
Block a user