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