fix(python,training): validate PyO3 inputs; make the single-training-job guard atomic

Closes the two findings from the adversarial review that were verified but left
unfixed. Both now have proper tests, each proven to fail against the old code.

1. PyO3 bindings panicked / over-allocated on caller input (aether.rs).
   `EmbeddingExtractor(n_heads=0)` reached `d_model % n_heads` in the transformer
   and panicked (surfacing to Python as an opaque PanicException); a non-divisor
   head count tripped the native assert; and `AetherConfig(d_model=100_000)`
   allocated multi-gigabyte weight matrices that abort the interpreter. The
   binding did no validation.

   Now both constructors return `PyResult` and validate at the boundary —
   positive dims, `d_model % n_heads == 0`, and a generous MAX_DIM/MAX_LAYERS cap
   — raising `ValueError`. Proven: `n_heads=0` -> "n_heads must be positive",
   `d_model=100_000` -> a clean ValueError, both previously a
   PanicException / abort. +10 pytest cases (test_aether.py); a valid config
   still constructs and embeds.

2. The single-training-job guard was a TOCTOU race (training_api.rs).
   `spawn_training_job` checked `is_active()` under a `state` READ lock, released
   it, then set `active` later. A tokio RwLock read lock is SHARED, so two
   concurrent `POST /train/start` could both hold it, both see the slot free, and
   both spawn jobs — sharing/overwriting one status+cancel and orphaning a task
   handle.

   Extracted `claim_training_slot`, which does the check-and-set in ONE `status`
   mutex scope — the atomicity lives on the status mutex, not the coarse state
   lock — so concurrent starts serialise and exactly one wins. This also makes it
   unit-testable without a full AppState.

   Test: 32 threads hit a barrier and race to claim; asserts EXACTLY ONE wins.
   Mutation-proven — reverting to the split check-then-set makes it fail
   (`left: 3, right: 1`), and it returns to 1 with the fix.

Verified on aarch64/macOS: training_api 28 pass (26 existing + 2 new), full
python/tests suite 237 pass (227 + 10). The native module keeps its internal
assert as a defence-in-depth invariant; the binding now enforces it at the edge.

Co-Authored-By: Ruflo & AQE

python/src/bindings/aether.rs

@@ -37,6 +37,16 @@ use wifi_densepose_aether::embedding::{
 };
 use wifi_densepose_aether::graph_transformer::TransformerConfig;
 
+/// Upper bound on model/CSI dimensions accepted from Python. The transformer
+/// allocates weight matrices quadratic in these, so this caps a single
+/// construction well under a gigabyte and turns an accidental or malicious
+/// `d_model=100_000` into a `ValueError` instead of an allocation that aborts
+/// the interpreter. Generous relative to real configs (defaults 64/128); raise
+/// deliberately if a workload genuinely needs larger.
+const MAX_DIM: usize = 4096;
+/// Upper bound on GNN layer count — a sanity cap, not a modelling limit.
+const MAX_LAYERS: usize = 64;
+
 // ─── AetherConfig ────────────────────────────────────────────────────
 
 /// Configuration for the contrastive embedding model.
@@ -56,15 +66,30 @@ pub struct PyAetherConfig {
 impl PyAetherConfig {
     #[new]
     #[pyo3(signature = (d_model=64, d_proj=128, temperature=0.07, normalize=true))]
-    fn new(d_model: usize, d_proj: usize, temperature: f32, normalize: bool) -> Self {
-        Self {
+    fn new(d_model: usize, d_proj: usize, temperature: f32, normalize: bool) -> PyResult<Self> {
+        // Validate at the boundary and raise ValueError. The native constructor
+        // allocates weight matrices quadratic in these dims and (elsewhere)
+        // divides by them, so zero or absurd values would otherwise reach Rust
+        // as a panic (surfacing to Python as an opaque PanicException) or a
+        // multi-gigabyte allocation that aborts the interpreter.
+        if d_model == 0 || d_proj == 0 {
+            return Err(PyValueError::new_err(
+                "d_model and d_proj must be positive",
+            ));
+        }
+        if d_model > MAX_DIM || d_proj > MAX_DIM {
+            return Err(PyValueError::new_err(format!(
+                "d_model ({d_model}) and d_proj ({d_proj}) must be <= {MAX_DIM}"
+            )));
+        }
+        Ok(Self {
             inner: EmbeddingConfig {
                 d_model,
                 d_proj,
                 temperature,
                 normalize,
             },
-        }
+        })
     }
 
     #[getter]
@@ -169,8 +194,30 @@ impl PyEmbeddingExtractor {
         n_keypoints: usize,
         n_heads: usize,
         n_gnn_layers: usize,
-    ) -> Self {
+    ) -> PyResult<Self> {
         let e_config = config.inner.clone();
+        // n_heads == 0 reaches `d_model % n_heads` in the transformer and panics
+        // (divide-by-zero); a non-divisor trips the native `assert!`. Both would
+        // surface to Python as a PanicException. Reject cleanly instead.
+        if n_heads == 0 {
+            return Err(PyValueError::new_err("n_heads must be positive"));
+        }
+        if e_config.d_model % n_heads != 0 {
+            return Err(PyValueError::new_err(format!(
+                "d_model ({}) must be divisible by n_heads ({n_heads})",
+                e_config.d_model
+            )));
+        }
+        if n_subcarriers == 0 || n_keypoints == 0 {
+            return Err(PyValueError::new_err(
+                "n_subcarriers and n_keypoints must be positive",
+            ));
+        }
+        if n_subcarriers > MAX_DIM || n_keypoints > MAX_DIM || n_gnn_layers > MAX_LAYERS {
+            return Err(PyValueError::new_err(format!(
+                "n_subcarriers/n_keypoints must be <= {MAX_DIM} and n_gnn_layers <= {MAX_LAYERS}"
+            )));
+        }
         let t_config = TransformerConfig {
             n_subcarriers,
             n_keypoints,
@@ -179,10 +226,10 @@ impl PyEmbeddingExtractor {
             n_gnn_layers,
         };
         let embedding_dim = e_config.d_proj;
-        Self {
+        Ok(Self {
             inner: EmbeddingExtractor::new(t_config, e_config),
             embedding_dim,
-        }
+        })
     }
 
     /// Extract an embedding from a CSI window (frames × subcarriers).

python/tests/test_aether.py

@@ -102,6 +102,56 @@ def test_config_roundtrips_fields() -> None:
     assert cfg.normalize is True
 
 
+# ─── Constructor input validation (raise ValueError, never panic) ─────────
+# Bad dimensions used to reach Rust and either panic (surfacing as an opaque
+# PanicException) or allocate multi-gigabyte matrices that abort the interpreter.
+
+@pytest.mark.parametrize("kwargs", [
+    {"d_model": 0, "d_proj": 128},
+    {"d_model": 64, "d_proj": 0},
+    {"d_model": 100_000, "d_proj": 128},   # unbounded allocation guard
+    {"d_model": 64, "d_proj": 100_000},
+])
+def test_aether_config_rejects_bad_dims(kwargs: dict) -> None:
+    with pytest.raises(ValueError):
+        aether.AetherConfig(**kwargs)
+
+
+def test_extractor_rejects_zero_heads_instead_of_panicking() -> None:
+    # THE crash codex flagged: n_heads=0 -> `d_model % n_heads` -> panic.
+    cfg = aether.AetherConfig(d_model=64, d_proj=128)
+    with pytest.raises(ValueError):
+        aether.EmbeddingExtractor(n_subcarriers=56, config=cfg, n_heads=0)
+
+
+def test_extractor_rejects_indivisible_head_count() -> None:
+    # 64 % 5 != 0 trips the native assert; must be a clean ValueError.
+    cfg = aether.AetherConfig(d_model=64, d_proj=128)
+    with pytest.raises(ValueError):
+        aether.EmbeddingExtractor(n_subcarriers=56, config=cfg, n_heads=5)
+
+
+@pytest.mark.parametrize("kwargs", [
+    {"n_subcarriers": 0},
+    {"n_subcarriers": 100_000},
+    {"n_keypoints": 0},
+])
+def test_extractor_rejects_bad_shape(kwargs: dict) -> None:
+    cfg = aether.AetherConfig(d_model=64, d_proj=128)
+    base = {"n_subcarriers": 56, "config": cfg}
+    base.update(kwargs)
+    with pytest.raises(ValueError):
+        aether.EmbeddingExtractor(**base)
+
+
+def test_valid_extractor_still_constructs() -> None:
+    # The negatives above must not pass by making the constructor reject
+    # everything: a valid config still builds and embeds.
+    cfg = aether.AetherConfig(d_model=64, d_proj=128)
+    ext = aether.EmbeddingExtractor(n_subcarriers=56, config=cfg, n_heads=4)
+    assert len(ext.embed(load_input())) == 128
+
+
 def test_embedding_shape_and_unit_norm() -> None:
     emb = build_extractor().embed(load_input())
     assert len(emb) == 128

v2/crates/wifi-densepose-sensing-server/src/training_api.rs

@@ -1698,17 +1698,48 @@ fn active_error(snap: &TrainingStatus) -> serde_json::Value {
 ///
 /// Centralises the single-job guard + spawn used by the supervised, pretrain,
 /// and LoRA start handlers so they cannot diverge.
+/// Atomically claim the single training slot.
+///
+/// Checks `active` and sets it `true` **in one `status` lock scope**, so two
+/// concurrent callers cannot both observe the slot free — the first claims it,
+/// the second gets `Err(current_status)`. Returns the seeded status on success.
+///
+/// This is the fix for a TOCTOU race: the previous code checked `is_active()`
+/// under a `state` READ lock, released it, and only afterward set `active`.
+/// A `tokio::RwLock` read lock is shared, so two starts could both hold it, both
+/// see the slot inactive, both proceed — spawning two jobs that then share and
+/// overwrite one status/cancel and orphan a task handle. The claim's atomicity
+/// lives on the `status` mutex, not the coarse `state` lock, which also keeps it
+/// unit-testable without a full `AppState`.
+fn claim_training_slot(
+    status: &Mutex<TrainingStatus>,
+    config: &TrainingConfig,
+) -> Result<(), TrainingStatus> {
+    let mut st = status.lock().unwrap();
+    if st.active {
+        return Err(st.clone());
+    }
+    *st = TrainingStatus {
+        active: true,
+        total_epochs: config.epochs,
+        lr: config.learning_rate,
+        patience_remaining: config.early_stopping_patience,
+        phase: "initializing".to_string(),
+        ..Default::default()
+    };
+    Ok(())
+}
+
 async fn spawn_training_job(
     state: &AppState,
     config: TrainingConfig,
     dataset_ids: Vec<String>,
     training_type: &'static str,
 ) -> Result<(), TrainingStatus> {
+    // Grab the shared handles under a read lock; the RwLock is only guarding
+    // access to the Arcs, not the single-job decision.
     let (progress_tx, status, cancel, history_snapshot) = {
         let s = state.read().await;
-        if s.training_state.is_active() {
-            return Err(s.training_state.snapshot());
-        }
         (
             s.training_progress_tx.clone(),
             s.training_state.status.clone(),
@@ -1717,16 +1748,10 @@ async fn spawn_training_job(
         )
     };
 
-    // Clear any prior stop request and seed the initial status snapshot.
+    // Atomic check-and-set on the status mutex. This — not the read lock above —
+    // is what serialises concurrent starts (see `claim_training_slot`).
+    claim_training_slot(&status, &config)?;
     cancel.store(false, Ordering::Relaxed);
-    *status.lock().unwrap() = TrainingStatus {
-        active: true,
-        total_epochs: config.epochs,
-        lr: config.learning_rate,
-        patience_remaining: config.early_stopping_patience,
-        phase: "initializing".to_string(),
-        ..Default::default()
-    };
 
     let handle = tokio::spawn(async move {
         run_training_job(
@@ -1950,6 +1975,60 @@ mod tests {
         assert_eq!(status.phase, "idle");
     }
 
+    #[test]
+    fn claim_training_slot_admits_exactly_one_concurrent_start() {
+        // Regression test for the single-job TOCTOU race. Many threads race to
+        // claim one slot at the same instant (a barrier maximises contention);
+        // the status mutex must admit EXACTLY ONE. A split check-then-set (the
+        // old shape) would let several through under load — verified by
+        // temporarily reverting the atomicity, which drops this from 1.
+        use std::sync::atomic::{AtomicUsize, Ordering as O};
+        use std::sync::{Arc, Barrier};
+
+        let status = Arc::new(Mutex::new(TrainingStatus::default()));
+        let config = TrainingConfig::default();
+        let winners = Arc::new(AtomicUsize::new(0));
+
+        const N: usize = 32;
+        let barrier = Arc::new(Barrier::new(N));
+        let mut handles = Vec::with_capacity(N);
+        for _ in 0..N {
+            let status = status.clone();
+            let config = config.clone();
+            let winners = winners.clone();
+            let barrier = barrier.clone();
+            handles.push(std::thread::spawn(move || {
+                barrier.wait();
+                if claim_training_slot(&status, &config).is_ok() {
+                    winners.fetch_add(1, O::SeqCst);
+                }
+            }));
+        }
+        for h in handles {
+            h.join().unwrap();
+        }
+
+        assert_eq!(
+            winners.load(O::SeqCst),
+            1,
+            "exactly one concurrent start may claim the single training slot"
+        );
+        assert!(
+            status.lock().unwrap().active,
+            "the slot must be marked active after a successful claim"
+        );
+    }
+
+    #[test]
+    fn claim_training_slot_rejects_when_already_active() {
+        let status = Arc::new(Mutex::new(TrainingStatus::default()));
+        let config = TrainingConfig::default();
+        assert!(claim_training_slot(&status, &config).is_ok(), "first claim wins");
+        let err = claim_training_slot(&status, &config)
+            .expect_err("second claim must be refused while active");
+        assert!(err.active, "the rejection carries the active status");
+    }
+
     #[test]
     fn training_progress_serializes() {
         let progress = TrainingProgress {