feat(querygen): orchestrate checkpoint resume with drift gate#242
Conversation
Adds the gen_queries orchestration: per-stage LLM fingerprints, the Stage 1 frozen-result skip and per-batch planning resume, Stage 2 per-batch realization resume, and the fail-fast-on-drift gate (raises QueryGenDriftError unless --force, with corrupt-file self-heal). Exports QueryGenDriftError on the public pragmata.querygen surface; CLI translates it to a clean typer.Exit. Larger than the 500-line target (~+800 net) because the orchestration api and its integration tests are tightly coupled; the alternative was an artificial binary-Stage-1-only intermediate that is not historically meaningful. Reviewable as one coherent unit.
…n planning-summary export (#237) **Stack (6 open PRs):** #237 → #238 → #239 → #240 → #241 → #242 **Chain goal:** Add per-batch checkpoint/resume to querygen so a re-run with the same `run_id` skips already-completed work. Nests resumable state under `<run_dir>/checkpoints/`, writes a frozen Stage 1 result (planning + dedup) plus per-batch Stage 1 and Stage 2 artifacts, and orchestrates resume behind a fail-fast config-drift gate (`--force` overrides) so checkpoints are never silently reused under changed settings. **This PR (1/6):** Extracts the tmpfile + atomic-rename idiom into a shared `core/atomic_io.py` helper and adopts it for querygen's planning-summary write — the atomic-write primitive every later checkpoint artifact (#239–242) writes through. --- ## What - `core/atomic_io.py` — `atomic_write_text` context manager: PID+uuid-uniquified tempfile, `Path.replace` for the atomic rename, no fsync (matches house style). - Querygen `export.py` adopts it for `export_planning_summary` (indented JSON for human-scannable persisted artifacts). - **Note:** annotation `record_builder`/`export_runner` should adopt `atomic_write_text` in a follow-up — kept out here to keep this PR's dependency surface to querygen only. ## Why Extracts the tmpfile + atomic-rename idiom (previously duplicated inline in two annotation sites — `record_builder.write_partition_manifest` and `export_runner.write_export_csv`) into one shared core helper, then adopts it for querygen's planning-summary write. Reaches the rule-of-3 with the upcoming querygen artifact writes in #239–242. ## Test plan - [x] `tests/unit/core/test_atomic_io.py` — round-trip + temp-cleanup-on-exception.
saschagobel
left a comment
There was a problem hiding this comment.
I do not think this PR is mergeable in its current shape.
The issue is not the idea of batch-level crash recovery itself. The issue is that this PR substantially changes the architecture and control flow of api/querygen.py. It does not simply add crash resumability to the existing flow; it refactors gen_queries() around Stage 1 results, per-batch resume, per-stage fingerprints, partial recomputation, force, damaged-checkpoint recovery, stale checkpoint cleanup, and checkpoint projection into the final CSV.
That is too much API-layer machinery for this feature. I’d like to keep gen_queries() structurally close to the current straight-line orchestration and only inject the minimal calls required for crash recovery at the existing planning, deduplication, and realization boundaries. Checkpoint policy, compatibility checks, recovery mechanics, and resume semantics should live in small core helpers/wrappers, while lower-level helpers such as run_planning_stage() and run_realization_stage() stay pure.
For this feature, I would expect the api/querygen.py diff to be fairly small: resolve the run compatibility state and delegate checkpoint-aware behavior to core. If the API-layer diff is medium-sized or larger on its own, I would treat that as a design smell and a sign that too much checkpoint/workflow logic is sitting at the wrong layer.
In particular, I would remove or move out of the API module:
_shorten()_drift_message()_read_reusable_checkpoint()_resume_or_run_planning_batch()_resume_or_run_realization_batch()- the per-stage
planning_llm_fingerprint/realization_llm_fingerprintlogic - damaged-checkpoint recovery policy
- stale realization checkpoint cleanup
forceparameter/behavior
I would also defer the broader reuse behavior. For v1, this should be crash/abort resumability only: same run_id, same effective settings, semantic batch checkpoints, atomic writes, and a strict same-run compatibility check. If a checkpoint is incompatible with the current invocation, failing clearly is enough; users can start a new run or delete/recreate the run directory.
I also noticed the added try/except handling for unreadable cross-run planning-summary artifacts. That may be a reasonable robustness improvement separately, but I do not think it belongs in this resumability stack.
So my suggestion would be to redraw this PR much more narrowly rather than polishing it in place: keep the API thin and straight, move the checkpoint/resume mechanics into core, and reduce the first implementation to crash-resumability only.
Stack (6 open PRs): #236 → #237 → #238 → #239 → #240 → #241 → #242
Chain goal: Add per-batch checkpoint/resume to querygen so a re-run with the same
run_idskips already-completed work. Nests resumable state under<run_dir>/checkpoints/, writes a frozen Stage 1 result (planning + dedup) plus per-batch Stage 1 and Stage 2 artifacts, and orchestrates resume behind a fail-fast config-drift gate (--forceoverrides) so checkpoints are never silently reused under changed settings.This PR (6/6): The orchestration layer that ties the foundation + three artifact slices together: per-batch resume for both stages, a frozen Stage 1 result that lets a same-
run_idrerun skip all of Stage 1, and a fail-fast-on-config-drift gate (--forceto override) that prevents silent reuse of checkpoints produced under different settings.What
gen_queries: Stage 1 frozen-result skip + per-batch planning resume (contiguous-prefix;resume_exhaustedsticks once any batch reruns); end-of-Stage-1 cross-run summary write; Stage 2 per-batch realization resume; CSV-as-projection of frozen result + realization artifacts._read_reusable_checkpointresolver is the single gate every checkpoint read goes through — clean match → reuse; drift +force=False→ raiseQueryGenDriftError(fail fast, message names old → current per field); drift +force=True→ recompute; damaged file → self-heal (silent recompute).QueryGenDriftErrorre-exported on the publicpragmata.querygensurface (advertised in the docstring; callers canexcept querygen.QueryGenDriftError).--force/--no-forceflag (per-invocation control, not a settings field); translatesQueryGenDriftErrorinto a cleantyper.Exit(1)with the message on stderr (not a Python traceback).This PR is large (~+800 net) because the orchestration api and its integration tests are inseparable.
Note on CLI scope
This PR adds only
--force+ theQueryGenDriftError→typer.Exittranslation to the existing CLI. The unmerged #222 runtime-knobs block (--requests-per-secondetc.) is intentionally not included — it is independent of this stack.Test plan
tests/unit/api/test_querygen.py— full resume + drift + force + self-heal + tolerance-rededup + planning-batch-drift integration tests.tests/unit/cli/test_cli_querygen.py— drift → exit 1 + message (no traceback).