Files
foxhunt/docs/plans/2026-03-01-training-observability-implementation.md
jgrusewski 7524406e3b 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>
2026-03-01 22:31:54 +01:00

12 KiB

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

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

// 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

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:

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:

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:

# 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

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:

            - 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

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