Training fault tolerance: stage-based resume, run manifest, DB supersede semantics

[zachtheyek]

Description

Training fault-tolerance overhaul: a persisted run manifest + explicit stage machine replace the checkpoint-hunting retry loop, and stale DB rows from failed attempts are flagged (never deleted) so plots stay clean across retries.

Stacked PR: base branch is feature/training-data-pipeline (#117). This PR contains only the fault-tolerance commits; review after #117 and rebase onto master once #117 merges.

Closes #119

What changed

• New src/aetherscan/run_state.py — TrainingRunState manifest persisted at {output_path}/run_state_{save_tag}.json (atomic tmp -> os.replace, mirroring round_data's .done protocol): run_start_time (wall clock of attempt 1), attempt counter, completed_rounds, stages_done, stages_failed. A corrupt manifest downgrades to a fresh run.
• train.py — run_training_pipeline() is now an explicit ordered stage machine (vae_rounds -> vae_plots -> rf_train -> rf_plots -> final_save) with skip-if-done + persist-on-success semantics (_execute_training_stages, duck-typed for unit testing):
  • TrainingPipeline.__init__ sets self.start_time from the manifest's run_start_time unconditionally — fixes the documented # BUG: at the old train_beta_vae() early return (unset start_time -> AttributeError in every plot, masking the real error). All DB queries/plots now span the whole run, not just the current attempt.
  • Round resume is manifest-driven: completed_rounds -> reload the round_{k:02d} checkpoint; each round records itself right after its checkpoint saves. Explicit --load-tag/--load-dir remain as an escape hatch (they override the manifest and clear downstream stage records).
  • vae_plots downgraded from raise-and-retry to record-and-continue: all four plots are attempted, failures land in stages_failed, and main.py exits nonzero at the very end if they never recover — a plot bug no longer costs a data regeneration, but lost artifacts stay loud. rf_plots failures are recorded the same way instead of being silently swallowed.
  • train_random_forest() persists random_forest_{tag}.joblib as soon as training completes; on resume it reloads model + eval artifacts instead of regenerating ~num_samples_rf cadences and retraining (gated on an actual resume so a fresh run reusing an old tag never skips training on stale artifacts).
  • The config JSON dump moved from train_command into the final_save stage (covered by retry).
• db/db.py — supersede semantics + minimal schema migration:
  • superseded INTEGER DEFAULT 0 on training_stats, injection_stats, latent_snapshots, inference_results (PR-08 in the v1 plan reuses this for inference resume).
  • Migration mechanism in _init_database gated on PRAGMA user_version with idempotent ALTER TABLE ... ADD COLUMN steps — minimally resolves the schema-evolution TODO.
  • Database.mark_superseded(table, tag, *, round_ge=None, npy_path=None) runs on the writer thread via a command tuple through the write queue (like the flush sentinel): buffered rows flush first, then the UPDATE — FIFO ordering guarantees rows written before the call are marked while later writes stay live. Runs inline when no writer thread exists.
  • query_* methods on those tables gain include_superseded: bool = False and filter superseded = 0 by default.
  • On resume of round k, (tag, round >= k) rows in the three training tables are marked before the rerun re-writes them — no duplicated epochs in loss curves.
• main.py — train_command retry loop simplified: no more get_latest_tag() checkpoint hunting or config.checkpoint.{load_dir,load_tag,start_round} mutation (the resolved # NOTE: block is removed); KeyboardInterrupt passthrough and max-retries accounting unchanged; exits nonzero when the manifest records permanently failed plot stages.

No CLI/config surface changed (no README CLI regeneration needed). preprocessing.py/inference.py untouched apart from the inference_results column + query filter (PR-06/08 own inference).

Deliberate deviations from the plan spec

• §6.2 suggested keeping get_latest_tag() as a retry-loop fallback when no manifest exists. The manifest always exists on retries (attempt 1 writes it in __init__), and hunting checkpoints without one risks silently resuming another run's round_XX checkpoints on a fresh tag — so the manifest is the single resume driver; get_latest_tag() remains as load_models()'s internal fallback and in train_random_forest's weight recovery.
• rf_train additionally persists random_forest_{tag}.joblib at stage end — required for the spec's own "load instead of retrain" resume path to have something to load before final_save runs.

Known limitation (deferred)

• RF-dataset injection stats are written with round_number=NULL, which round_ge marking can't target; an rf_train retry that regenerates the RF dataset can leave stale RF-generation rows live. Bounded (partial-generation rows only); PR-08/PR-10 own the follow-up.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

How has this been tested?

Unit tests (local, arm64 venv w/ TF 2.17.1)

PYTHONPATH=src pytest -m "not gpu and not cluster" -q -> 269 passed, 3 skipped. New coverage:

• tests/unit/test_run_state.py — manifest round-trip, atomic tmp->replace persistence (failed write leaves the previous manifest intact, no .tmp litter), corrupt-manifest downgrade, stage/round bookkeeping.
• tests/unit/test_db.py — supersede marking (round_ge, npy_path, tag scoping), queue-FIFO semantics (queued-before rows get marked, written-after rows stay live), default query filtering incl. aggregation queries, inline execution without a writer thread, filter validation; migration: old-schema (v0) DB fixture gains the column + user_version stamp, idempotent across reopens, fresh DBs created at current version.
• tests/unit/test_train_utils.py — stage machine against a stub pipeline: full fresh-run order, skip-if-done (incl. RF reload on skip), plot-failure record-and-continue, rf_train failure propagation + best-effort VAE save, reload-failure fallback to retrain, relaunch-retries-only-failed-stage.

Cluster kill/resume matrix (blpc3, 5 GPU, singularity)

Smoke config: --per-replica-batch-size 4 --per-replica-val-batch-size 4 --effective-batch-size 20 --num-samples-beta-vae 200 --num-samples-rf 200 --num-training-rounds 2 --epochs-per-round 2 --latent-viz-num-cadences-per-type 5, tag test_v23 (one continuous chain, relaunching the identical command after each SIGKILL); in-process retry run on tag test_v24.

#	SIGKILL point (verified in log)	Relaunch behavior (all asserted from logs + manifest + DB)	Result
1	Mid-round-2 epochs (3 s after round-2 epoch 1 completed)	Attempt 2: manifest completed_rounds [1] -> Resuming from checkpoint round_01 (per run manifest), loop starts at ROUND 2/2; round_ge=2 marked 14432 injection + 20 snapshot rows stale (the 21 round-2 training rows died unbuffered in the write queue — nothing to mark); stale round_02 data dir deleted + regenerated	PASS
2	During vae_plots (1 s into latent-GIF generation; 3 of 4 plots saved)	Attempt 3: vae_rounds skipped via manifest (stages_done ['vae_rounds'], vae_plots correctly unmarked), checkpoint round_02 reloaded, all four vae_plots re-ran to completion; round_ge=3 marked 0 rows (nothing stale)	PASS
3	During rf_train (2 s into RF data generation; kill also reaped the forked pool workers)	Attempt 4: vae_rounds + vae_plots both skipped; no RF artifacts had been persisted, so rf_train correctly re-ran from scratch; supersede no-op	PASS
4	During rf_plots (after 3 of 10 RF plots; rf_train had persisted random_forest_{tag}.joblib + rf_eval_artifacts_{tag}.joblib at stage end)	Attempt 5: vae_rounds/vae_plots skipped, Stage 'rf_train' already complete — loaded persisted RF model (no RF data regeneration), all ten rf_plots re-ran to completion, final_save wrote encoder/decoder/RF + config_{tag}.json, exit 0 with Training completed successfully!	PASS
5	(none — relaunch after clean completion)	Attempt 6: pure no-op — all five stages logged already complete — skipping (rf_train via RF reload), clean exit 0	PASS
R	Transient injection: plots/checkpoints chmod 555 during round 1, restored during the 60 s retry sleep (--max-retries 3, tag test_v24)	Attempt 1 died at PermissionError: .../plots/checkpoints/beta_vae_loss_curves_round_01.png (20:11:36); permissions restored during the retry sleep; Training attempt: 2/3 at 20:12:37 resumed via the manifest, marked 42 training / 14432 injection / 40 snapshot rows superseded at round_ge=1, re-ran everything, Training completed successfully! at 20:45 — exercising the in-process retry + per-attempt producer teardown end to end	PASS

Final-state assertions (both tags): manifest shows all five stages done with run_start_time unchanged since attempt 1 (spans every attempt — test_v23's covered 5 attempts over ~1 h 45 m); exactly one live total_loss row per (round, epoch) with all stale rows (superseded=1) filtered by default queries — the duplicated-epoch corruption is gone; encoder/decoder/RF/eval-artifacts/config all present; 56 final plots per tag; /dev/shm clean after every SIGKILL + relaunch cycle.

Bonus real-world validation: the cluster's live aetherscan.db (pre-existing schema, prior tags) was migrated in place on first open — user_version 0 -> 1, superseded added to all four tables, old rows visible with superseded=0, and every later launch skipped the migration cleanly.

Observed pre-existing quirk (not a regression, noted for PR-10): at smoke scale, train_random_forest's check_encoder_trained fallback loaded test_v21 weights via get_latest_tag — the already-documented # BUG: block about tag reuse; this PR's tag-scoped RF resume path is unaffected.

Deferred validation (bla0 unreachable this session)

• §6.3's original bla0 6-GPU smoke variant of the same matrix (--num-samples-beta-vae 3072 --num-samples-rf 960, apptainer/conda runtime).
• Migration behavior against bla0's existing aetherscan.db (blpc3's live DB was migrated in place during this matrix — see results).
• Full-scale resume rehearsal (default config) after a real OOM kill.

Checklist

• My code follows the style guidelines of this project (ruff lint + format)
• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (deferred: DATABASE.md documenting the migration mechanism is scoped to the docs-suite PR per the plan)
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

AI disclosure

This PR was implemented end-to-end with Claude Code: implementation, unit tests, docs/comments, cluster kill/resume verification, and this PR text (design from the maintainer's v1 plan). All cluster runs and assertions above were executed and checked as part of this session.

[claude]

Claude finished @zachtheyek's task in 1m 41s —— View job

---

Code Review: Training fault tolerance overhaul

• Read and analyze changed files
• Review run_state.py (new module)
• Review db/db.py changes (supersede semantics + migration)
• Review train.py changes (stage machine)
• Review main.py changes (simplified retry loop)
• Review test coverage
• Post detailed feedback

---

Summary

This is a well-designed fault-tolerance overhaul. The architecture — a persisted run manifest driving an explicit stage machine with skip-if-done semantics — cleanly separates concerns and replaces fragile checkpoint-hunting with deterministic resume logic. The cluster kill/resume matrix is impressively thorough testing.

Overall: looks good to merge. A few observations below, categorized by severity.

---

Issues

1. Minor: _writer_loop shutdown flush log is dead code (db/db.py:649)

# Final flush on shutdown
if self.buffer:
    self._flush_buffer()
    self.buffer.clear()
    logger.info(f"Flushed {len(self.buffer)} remaining data on shutdown")

The logger.info on line 649 will always print 0 because self.buffer.clear() runs before the log statement. This is pre-existing (not introduced by this PR), but worth noting since the surrounding code was touched.

2. Minor: _execute_training_stages clears RF caches only in the rf_plots branch (train.py)

If rf_plots is already done (skipped), _clear_rf_caches is never called. This is fine in practice (the caches won't be populated if the stage is skipped), but worth a comment to make the invariant explicit for future readers.

3. Nit: load_run_state silently downgrades on missing run_start_time (run_state.py:114)

run_start_time=float(payload["run_start_time"]),

If someone hand-edits the manifest and removes this field, KeyError triggers the catch-all except Exception and logs a generic "Failed to load" warning. Since this is a required field, the behavior is correct (fresh-run downgrade), but a dedicated KeyError branch with a more specific message would help debugging — low priority.

4. Observation: No index on superseded column

mark_superseded issues UPDATE ... WHERE superseded = 0 AND tag = ? (+ optional round_number >= ?). The existing composite indices lead with (tag, timestamp, ...), which covers the tag predicate, but the superseded = 0 filter must scan within that index range. At current data volumes this is fine (the cluster test marked 14k rows quickly), but if rows per tag grow to millions, a partial index WHERE superseded = 0 on (tag, round_number) would help. Noting for future consideration, not a blocking concern.

5. Observation: try_load_rf_for_resume doesn't load the eval-artifacts joblib

The method checks that rf_eval_artifacts_{tag}.joblib exists (gating the return True), and the comment says "rf_plots and final_save need a live model" — but the eval artifacts themselves aren't loaded into memory here. This works because plot_rf_diagnostics() calls _load_rf_eval_artifacts() internally, but the asymmetry (check existence but don't load) is worth a brief comment at the check site.

---

Positives

• Atomic manifest writes via tmp → os.replace mirror the existing round_data protocol — good consistency.
• Stage machine is duck-typed for testing — the _StageMachineStub in test_train_utils.py exercises all the interesting edge cases (skip, fail-and-continue, reload fallback) without needing TF.
• Supersede semantics (flag, never delete) are correct for an append-only analytics DB. The FIFO ordering guarantee via the writer thread is well-reasoned.
• Migration gated on PRAGMA user_version with idempotent column-existence checks — handles the "db file copied without journal" edge case gracefully.
• _run_plot_group tries all plots even when one fails — much better than the old raise-immediately behavior for non-critical stages.
• Test coverage is solid: manifest round-trip + atomicity, supersede queue semantics, stage machine logic all have dedicated tests.

---

Minor style nits

• run_state.py:72-73: The multiline docstring opening for record_stage_failure starts on the def line (violates the common single-line or next-line convention used elsewhere in this module). Cosmetic only.

---

No blocking issues found. The design is sound, the test matrix is thorough, and the code is production-ready.

[claude]

Nit: the method checks that rf_eval_artifacts_{tag}.joblib exists (and won't return True without it), but doesn't actually load the artifacts into memory — that's left to plot_rf_diagnostics() → _load_rf_eval_artifacts(). The existence check is correct (it gates "is resume viable?"), but a one-line comment explaining this asymmetry would help future readers understand why we check-but-don't-load:

Suggested change

	if self.rf_model is None:
	def try_load_rf_for_resume(self) -> bool:
	"""
	Attempt to restore the trained Random Forest persisted by a previous attempt of this
	run: both random_forest_{tag}.joblib and rf_eval_artifacts_{tag}.joblib must exist
	under model_path (the artifacts joblib is what every RF plot consumes, so resuming
	into rf_plots without it would be pointless — but it is loaded lazily by
	_load_rf_eval_artifacts(), not here). Returns True when the model is loaded
	and ready; False falls back to full RF training.
	"""

[zachtheyek]

Applied in e836642 — took the suggested wording: the docstring now states the eval-artifacts joblib is existence-checked here (gating resume viability) but loaded lazily by _load_rf_eval_artifacts().

[claude]

Ultra-minor style nit: the docstring for record_stage_failure below starts on the same line as def then wraps — elsewhere in this module (e.g. mark_stage_done) single-line docstrings are used. No functional impact, just a cosmetic inconsistency.

[zachtheyek]

Fixed in e836642 — collapsed to a single-line docstring matching mark_stage_done's style.

[zachtheyek]

Thanks for the review — all actionable notes addressed in e836642:

1. _writer_loop shutdown flush log always printing 0 — fixed inline (captured flushed_count before buffer.clear()). Pre-existing, but this PR touches _writer_loop (the mark-superseded command branch), so it was fair game.
2. _clear_rf_caches only in the rf_plots branch — added a comment making the invariant explicit: the caches are populated exclusively by the plot calls in that branch, so the skip-if-done path never has anything to clear.
3. load_run_state missing-required-field handling — added a dedicated KeyError branch that names the missing field in the warning (Run state at ... is missing required field 'run_start_time' — treating as a fresh run); behavior (fresh-run downgrade) unchanged. The existing test_missing_required_field_downgrades_to_none covers the new branch.
4. No index on superseded — agreed, acknowledged as a future consideration. At current volumes the (tag, ...)-leading composite indices carry the predicate fine (the cluster matrix marked 14k rows in ~60 ms); if per-tag row counts grow orders of magnitude, a partial index ON <table>(tag, round_number) WHERE superseded = 0 is the natural follow-up — a good fit for the benchmarking/DB work later in the v1 chain.
5. try_load_rf_for_resume check-but-don't-load asymmetry — applied your suggested docstring wording (see inline thread).
6. record_stage_failure docstring style — collapsed to a single line (see inline thread).

Affected unit tests re-run (89 passed) and all pre-commit hooks green.

[zachtheyek]

@claude please perform a second-pass code review of this PR. Re-review the current diff for any remaining correctness bugs, security issues, or significant quality problems, and post your findings. If the diff is clean, please say so explicitly. (Final verification round before merge.)

[claude]

Second-pass code review — PR #122

I've done a thorough re-review of the current diff across all four production files (run_state.py, db/db.py, train.py, main.py) and the three test files. The diff is clean. No correctness bugs, no security issues, and the code quality is solid throughout. Details below.

Areas examined

#	Area	Verdict
1	SQL injection in _execute_mark_superseded — table is f-string interpolated into the UPDATE. Verified that mark_superseded() validates against the _SUPERSEDE_TABLES allowlist before every code path (inline and writer-thread) reaches the interpolation. No unvalidated path exists.	Clean
2	_writer_loop sentinel detection — metric[0] is _MARK_SUPERSEDED_SENTINEL could fail if a non-tuple were dequeued, but all queue items are tuples: flush is (_FLUSH_SENTINEL, event), mark is (_MARK_SUPERSEDED_SENTINEL, payload, event), data is (table, values). No bare sentinels are ever enqueued.	Clean
3	Shutdown race in mark_superseded — TOCTOU window between the stop_event check and write_queue.put(). Benign: the caller's wait loop re-checks stop_event every 100 ms and returns False, so no deadlock. Same pattern as flush().	Minor (benign)
4	round_ge=1 supersede marking — When --start-round 1 is explicit, mark_superseded(..., round_ge=1) flags all rows for the tag. This is correct: "restart from round 1" means all prior rows are stale. The stages_done = [] clearing confirms the intent.	Clean (by design)
5	RF failure path in _execute_training_stages — After train_random_forest() raises, rf_model is set to None, then save_models() is called best-effort. Verified save_models guards RF save with if self.rf_model is not None; VAE weights are still saved.	Clean
6	self._resumed gating on try_load_rf_for_resume — _resumed is set in _init_run_state() (called from __init__), which completes before _execute_training_stages runs. Always set before it's read.	Clean
7	os.replace atomicity — Atomic on POSIX. Non-issue since the project targets Linux HPC exclusively.	N/A
8	pipeline variable when max_retries=0 — With max_retries=0, the for range(0) loop never executes, pipeline stays None, and the post-loop code would report "Training completed successfully!" without training. However, the >= 0 validation is pre-existing (cli.py:1421, cli.py:1823) and not introduced by this PR. The post-loop guard (if pipeline is not None else []) is actually a correct defensive addition by this PR.	Pre-existing (not a regression)
9	Migration ordering — _migrate_schema runs after CREATE TABLE IF NOT EXISTS. Fresh DBs already have superseded from the CREATE; the migration's PRAGMA table_info check finds it, skips the ALTER, and stamps user_version. Old DBs get the ALTER. Both converge correctly. Idempotent across reopens.	Clean
10	flushed_count fix — Old code logged len(self.buffer) after self.buffer.clear() (always 0). The fix captures the count before clearing. Correct.	Clean

Additional observations

• Test coverage is thorough: test_run_state.py covers round-trip, atomicity (failed-write leaves previous manifest intact), corrupt-file downgrade, and all bookkeeping helpers. test_db.py covers supersede marking (FIFO ordering, tag scoping, filter validation, inline execution, aggregation exclusion) and schema migration (old-schema upgrade, idempotency, fresh-DB version stamp). test_train_utils.py drives the stage machine against a duck-typed stub covering all branch paths.
• Style is consistent: from __future__ import annotations, field(default_factory=...) for mutable defaults, f-string loggers, no bare print().
• The _SUPERSEDE_TABLES allowlist + filter-validation pattern in mark_superseded is a clean design — it makes the table-name interpolation safe and the per-table filter constraints self-documenting.
• The atomic write pattern (tmp → os.replace, cleanup on failure) mirrors the existing round_data .done protocol — good consistency.

Verdict

The diff is clean — no correctness bugs, security issues, or significant quality problems found. Ready to merge (after #117).

Note: the max_retries >= 0 validation allowing a no-op --max-retries 0 is a pre-existing issue worth a follow-up fix (changing both validators in cli.py to >= 1), but it's outside this PR's scope.