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
This commit is contained in:
Dragan Spiridonov
2026-07-24 14:45:55 +02:00
parent bc0e8fd031
commit 788d685b9f
3 changed files with 194 additions and 18 deletions
@@ -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 {