S5 verification runbook + GNN health endpoint (#102)

Author: hyperpolymath

Justfile

@@ -349,6 +349,13 @@ train:
 train-cpu:
     ECHIDNA_MAX_PROOF_STATES=2000 ECHIDNA_NUM_EPOCHS=2 julia --project=src/julia src/julia/run_training_cpu.jl
 
+# Evaluate trained model on held-out validation split (MRR / top-k).
+# Requires `just train-cpu` (or `just train`) to have produced models/neural/.
+# Appends one JSONL row to training_data/eval_results.jsonl and exits non-zero
+# if MRR < 0.66 (below the cosine baseline).
+eval:
+    julia --project=src/julia src/julia/eval_held_out.jl
+
 # End-to-end pipeline: provision → extract → merge → align → retrain.
 # Use `ECHIDNA_MAX_PROOF_STATES=0 just corpus-refresh` to lift the sample cap.
 corpus-refresh: provision-corpora extract-corpora merge-corpora align-premises retrain

docs/handover/S5-VERIFICATION-RUNBOOK.md

@@ -0,0 +1,107 @@
+# S5 Verification Runbook: Trained-Weights Verification Flow
+
+## Purpose
+
+This runbook proves end-to-end that once `just train-cpu` produces real weights
+in `models/neural/`, (a) the Julia GNN server picks them up and stops using cosine
+similarity, (b) the Rust backends that call `gnn_augment_tactics` receive
+model-derived scores via the `/gnn/rank` wire format, and (c) MRR on a held-out
+validation split can be measured and compared against the 0.66 cosine baseline.
+Steps 1 and 2 can run before training has finished; only Step 3 requires weights.
+
+---
+
+## Prerequisite
+
+`just train-cpu` (or `just train` on GPU) has completed and
+`models/neural/gnn_ranker/` (or `best_model/` or `final_model/`) exists.
+
+---
+
+## Step 1 — Confirm wire format (no training needed)
+
+```bash
+cargo test --test gnn_augment_integration
+```
+
+Spawns an in-process mock HTTP server and asserts:
+
+- `GnnClient::health_status()` returns the richer payload (model_path, vocab_size,
+  training_records_received).
+- For each of rocq, lean, agda, isabelle, z3: `suggest_tactics` returns
+  `Tactic::Custom { command: "apply", args: ["lemma_foo"] }` as the first tactic,
+  proving the `/gnn/rank` wire format is consumed correctly by every backend.
+
+Expected output: `test result: ok. 6 passed; 0 failed`.
+
+---
+
+## Step 2 — Confirm model load
+
+Start the Julia GNN server:
+
+```bash
+julia --project=src/julia src/julia/api_server.jl
+```
+
+(or however you normally start it — check `src/julia/run_server.jl` for args).
+Then query health:
+
+```bash
+curl -s http://localhost:8090/gnn/health | jq .
+```
+
+Expected response when weights are loaded:
+
+```json
+{
+  "status": "ok",
+  "gnn_model_loaded": true,
+  "model_path": "/absolute/path/to/models/neural/gnn_ranker",
+  "vocab_size": <non-zero>,
+  "training_records_received": 0,
+  "num_gnn_layers": 4,
+  "service": "echidna-gnn",
+  "version": "2.1.0"
+}
+```
+
+`gnn_model_loaded: true` and `vocab_size > 0` confirm the server picked up real weights
+and is no longer falling back to cosine similarity.
+
+---
+
+## Step 3 — Measure lift
+
+```bash
+just eval
+```
+
+Loads the trained model, runs `compute_metrics` on the held-out validation split
+(20% of `training_data/`), and appends one JSONL row to
+`training_data/eval_results.jsonl`.  Read the result:
+
+```bash
+tail -1 training_data/eval_results.jsonl | jq '{mrr, top1, top5, top10, val_size}'
+```
+
+---
+
+## Acceptance
+
+| Result | Action |
+|--------|--------|
+| MRR ≥ 0.66 | Baseline met. Log the row. Run `just train` on GPU for full training. |
+| MRR < 0.66 | **STOP.** File a bug. Do not paper over it — this signals architecture mismatch, label leakage, or vocabulary miss. |
+
+---
+
+## Troubleshooting
+
+| Symptom | Check |
+|---------|-------|
+| `gnn_model_loaded: false` after `just train-cpu` | Confirm `models/neural/gnn_ranker/` exists and contains `model.bson`, `vocabulary.bson`, `config.bson`. Restart the Julia server. |
+| Integration test fails on `connection refused` | Re-run: `cargo test --test gnn_augment_integration`. The mock binds `127.0.0.1:0`; transient OS port exhaustion is extremely rare but retry once. |
+| `just eval` errors with `No trained model found` | Run `just train-cpu` first. |
+| `just eval` errors on a missing solver method | File a bug in the `EchidnaML` Julia module. Do not patch the eval script to bypass. |
+| MRR unchanged from 0.66 after training | Check `models/model_metadata.txt` — if it still reads `0 words, 0 classes` the training run did not write weights. Check training logs. |

src/julia/api/gnn_endpoint.jl

@@ -29,6 +29,7 @@ using LinearAlgebra
 
 # Reference to global GNN model (loaded at server startup)
 const GNN_MODEL = Ref{Any}(nothing)
+const GNN_MODEL_PATH = Ref{Union{Nothing,String}}(nothing)
 const GNN_VOCAB = Ref{Any}(nothing)
 
 # Per-(prover, domain) success-rate weights pushed from Rust via
@@ -64,6 +65,7 @@ function load_gnn_model(models_dir::String)
         try
             solver = load_solver(model_path)
             GNN_MODEL[] = solver
+            GNN_MODEL_PATH[] = model_path
             @info "GNN model loaded successfully from $model_path"
             return true
         catch e
@@ -74,6 +76,7 @@ function load_gnn_model(models_dir::String)
     # All candidates exhausted — cosine fallback is genuine missing-model path.
     @warn "No trained GNN model found in $(joinpath(models_dir, "neural")) — ranking will use cosine similarity until weights are trained (run: just train-cpu)"
     GNN_MODEL[] = nothing
+    GNN_MODEL_PATH[] = nothing
     return false
 end
 
@@ -360,9 +363,12 @@ function handle_gnn_health(req::HTTP.Request)
     response = Dict(
         "status" => "ok",
         "gnn_model_loaded" => model_loaded,
+        "model_path" => model_loaded ? GNN_MODEL_PATH[] : nothing,
+        "vocab_size" => model_loaded ? length(GNN_VOCAB[]) : 0,
+        "training_records_received" => TOTAL_TRAINING_RECORDS[],
         "num_gnn_layers" => model_loaded ? 4 : 0,
         "service" => "echidna-gnn",
-        "version" => "2.1.0"
+        "version" => "2.1.0",
     )
 
     return HTTP.Response(200, JSON3.write(response))

src/julia/eval_held_out.jl

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: MPL-2.0
+# SPDX-FileCopyrightText: 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
+#
+# eval_held_out.jl — Evaluate a trained GNN model on the held-out validation split.
+#
+# Measures MRR, top-1, top-5, top-10 and appends one JSONL row to
+# training_data/eval_results.jsonl.  Compares against the 0.66 cosine baseline.
+#
+# Usage:
+#   just eval
+# or:
+#   cd src/julia && julia --project=. eval_held_out.jl
+
+using Pkg
+Pkg.activate(joinpath(@__DIR__))
+
+using EchidnaML
+using JSON3, Dates, Statistics, BSON
+
+function main()
+    repo_root = abspath(joinpath(@__DIR__, "..", ".."))
+    data_dir  = joinpath(repo_root, "training_data")
+
+    # Try candidate model directories in the same order as load_gnn_model.
+    candidate_dirs = [
+        joinpath(repo_root, "models", "neural", "gnn_ranker"),
+        joinpath(repo_root, "models", "neural", "best_model"),
+        joinpath(repo_root, "models", "neural", "final_model"),
+    ]
+
+    models_dir = nothing
+    for d in candidate_dirs
+        if isdir(d)
+            models_dir = d
+            break
+        end
+    end
+
+    if models_dir === nothing
+        error("No trained model found — run 'just train-cpu' first. " *
+              "Checked: $(join(candidate_dirs, ", "))")
+    end
+
+    @info "Loading solver from $models_dir"
+    solver = load_solver(models_dir)
+
+    @info "Loading validation data from $data_dir"
+    t_start = time()
+    (_train_ds, val_ds, vocab) = load_training_data(data_dir; train_split=0.8f0)
+    @info "Validation split size: $(length(val_ds.examples))  vocab: $(vocab.vocab_size)"
+
+    if isempty(val_ds.examples)
+        error("Validation dataset is empty — check training_data/ has proof_states and premises files.")
+    end
+
+    # Evaluate using the same scoring path as train.jl::compute_metrics.
+    # compute_metrics returns (precision, recall, mrr) using cosine over text embeddings;
+    # when a trained model is loaded, the solver's text_encoder reflects learned weights.
+    val_metrics = compute_metrics(solver, val_ds; k=10)
+    mrr   = val_metrics.mrr
+    top1  = val_metrics.precision   # precision@1 is top-1 accuracy at k=1
+    top5  = val_metrics.recall      # secondary; full top-k breakdown below
+    top10 = val_metrics.precision   # reuse — compute_metrics gives precision@k
+
+    # Finer breakdown: compute_metrics at k=1 and k=5.
+    val_k1  = compute_metrics(solver, val_ds; k=1)
+    val_k5  = compute_metrics(solver, val_ds; k=5)
+    top1  = val_k1.precision
+    top5  = val_k5.precision
+    top10 = val_metrics.precision
+
+    duration = time() - t_start
+
+    git_sha = try
+        strip(read(`git -C $repo_root rev-parse HEAD`, String))
+    catch
+        "unknown"
+    end
+
+    metrics = Dict(
+        "timestamp"               => string(now()),
+        "git_sha"                 => git_sha,
+        "model_path"              => models_dir,
+        "val_size"                => length(val_ds.examples),
+        "vocab_size"              => vocab.vocab_size,
+        "mrr"                     => mrr,
+        "top1"                    => top1,
+        "top5"                    => top5,
+        "top10"                   => top10,
+        "cosine_baseline_mrr"     => 0.66,
+        "eval_duration_seconds"   => duration,
+    )
+
+    out_path = joinpath(data_dir, "eval_results.jsonl")
+    open(out_path, "a") do io
+        println(io, JSON3.write(metrics))
+    end
+
+    @info "Eval complete: MRR=$mrr  top1=$top1  top5=$top5  top10=$top10"
+    @info "Results appended to $out_path"
+
+    if mrr >= 0.66f0
+        @info "PASS: MRR=$mrr ≥ 0.66 cosine baseline — weights are improving retrieval."
+    else
+        @warn "FAIL: MRR=$mrr < 0.66 cosine baseline — stop, file a bug before proceeding."
+        exit(1)
+    end
+end
+
+main()

src/rust/gnn/client.rs

@@ -138,7 +138,7 @@ struct GnnRankResponse {
     error: Option<String>,
 }
 
-/// GNN health check response.
+/// GNN health check response (internal, minimal — kept for legacy deserialization path).
 #[derive(Debug, Deserialize)]
 struct GnnHealthResponse {
     /// Server status
@@ -150,6 +150,25 @@ struct GnnHealthResponse {
     num_gnn_layers: usize,
 }
 
+/// Richer health payload returned by `/gnn/health` after the S5 extension.
+///
+/// All fields present when `gnn_model_loaded` is true; `model_path`,
+/// `vocab_size`, and `training_records_received` are the new additions.
+#[derive(Debug, Clone, Deserialize)]
+pub struct GnnHealth {
+    pub status: String,
+    pub gnn_model_loaded: bool,
+    pub model_path: Option<String>,
+    #[serde(default)]
+    pub vocab_size: usize,
+    #[serde(default)]
+    pub training_records_received: u64,
+    #[serde(default)]
+    pub num_gnn_layers: usize,
+    pub service: String,
+    pub version: String,
+}
+
 /// Result of a GNN inference call.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GnnInferenceResult {
@@ -199,6 +218,28 @@ impl GnnClient {
         }
     }
 
+    /// Fetch the richer `/gnn/health` payload introduced in S5.
+    ///
+    /// Returns the full `GnnHealth` struct (model path, vocab size,
+    /// training record count, etc.).  Does NOT update `server_available`
+    /// — callers that need the availability gate should use `check_health`.
+    pub async fn health_status(&self) -> Result<GnnHealth> {
+        let url = format!("{}/gnn/health", self.config.api_url);
+        let response = self
+            .client
+            .get(&url)
+            .send()
+            .await
+            .context("Failed to reach GNN health endpoint")?;
+        if !response.status().is_success() {
+            anyhow::bail!("GNN health returned HTTP {}", response.status());
+        }
+        response
+            .json::<GnnHealth>()
+            .await
+            .context("Failed to parse GnnHealth response")
+    }
+
     /// Check if the GNN server is available and has a loaded model.
     ///
     /// Updates `server_available` state. Call this before making

src/rust/gnn/mod.rs

@@ -49,7 +49,7 @@ pub mod fallback_monitor;
 pub mod graph;
 pub mod guided_search;
 
-pub use client::{GnnClient, GnnConfig, GnnInferenceResult};
+pub use client::{GnnClient, GnnConfig, GnnHealth, GnnInferenceResult};
 pub use embeddings::{TermEmbedding, TermFeatureExtractor};
 pub use fallback_monitor::{FallbackMetrics, FallbackMonitor, FallbackSlaConfig};
 pub use graph::{EdgeKind, NodeKind, ProofGraph, ProofGraphBuilder};

src/rust/provers/mod.rs

@@ -1516,7 +1516,7 @@ pub(crate) async fn gnn_augment_tactics(
     use crate::gnn::graph::ProofGraphBuilder;
 
     let graph = ProofGraphBuilder::new(4).build_from_proof_state(state);
-    let gnn = GnnClient::with_config(GnnConfig {
+    let mut gnn = GnnClient::with_config(GnnConfig {
         api_url: url,
         timeout_ms: 2000,
         top_k: 8,
@@ -1530,6 +1530,11 @@ pub(crate) async fn gnn_augment_tactics(
     // so Julia's /gnn/rank can apply per-domain weights from prior outcomes.
     // When metadata has no "aspects" key (e.g. direct REPL invocation), aspects
     // is empty and behaviour is identical to the no-hint path.
+    // Ping the server so rank_premises_with_aspects sees server_available = true.
+    // Graceful: if check_health errors (server down), gnn.server_available stays
+    // false and rank_premises_with_aspects returns empty — no disruption to callers.
+    let _ = gnn.check_health().await;
+
     let aspects: Vec<String> = state
         .metadata
         .get("aspects")