PR #624 Review — fix: harden resume checkpoint handling

Summary

This PR hardens the resume checkpoint pipeline to address several crash-recovery and correctness issues:

• Atomic metadata writes — metadata.json is now written via tmp file + fsync + os.replace, so a crash mid-write cannot leave a partially written checkpoint.
• Config-fingerprint in metadata — every metadata write is stamped with config_hash, config_hash_algo, and config_hash_version. Resume compatibility checks read metadata.json first and fall back to builder_config.json only for older artifacts.
• Single-writer lock — a new ArtifactStorage.dataset_lock() context manager creates an exclusive .data_designer.lock file (O_CREAT | O_EXCL) for the life of a run so two processes cannot write to the same dataset directory.
• Unsafe-resume rejection — resume now fails fast for allow_resize=True columns (row boundaries are non-deterministic) and for datasets that already completed process_after_generation() (schema/row-count may have changed).
• Real parquet row counts on resume — the async engine scans parquet-files/batch_*.parquet and reads parquet metadata to recover the actual persisted row count per row-group, which matters for partial-salvage scenarios.
• Better error on corrupt metadata — a json.JSONDecodeError while loading metadata.json now raises a dedicated DatasetGenerationError instead of being silently treated as "assume compatible."
• Docs updated in architecture/dataset-builders.md, the Fern docs (architecture-and-performance.mdx, processors.mdx).

Findings

🟠 Bug — outer lock caches resolved_dataset_name prematurely

In data_designer.py:268+, create() now acquires the lock before calling builder.build():

with builder.artifact_storage.dataset_lock():
    builder.build(num_records=num_records, resume=resume)

dataset_lock() calls mkdir_if_needed(self.base_dataset_path) (artifact_storage.py:252), which accesses base_dataset_path → the cached_property resolved_dataset_name. This is the exact premature-caching the existing _check_resume_config_compatibility comment warns against ("must NOT access base_dataset_path … which would cache resolved_dataset_name prematurely").

Inside build(), when resume == IF_POSSIBLE and the stored config fingerprint is incompatible, the code pops the cache and re-resolves to a new timestamped directory:

self.artifact_storage.resume = ResumeMode.NEVER
self.artifact_storage.__dict__.pop("resolved_dataset_name", None)

After this switch, base_dataset_path now points at foo_MM-DD-YYYY_HHMMSS/ — but the outer lock file was already created in the original foo/ directory. The new timestamped directory is unlocked, and the lock file in foo/ protects a directory the run never writes to.

In practice two concurrent runs will race to unique timestamps and not collide, so this is unlikely to cause data corruption today. But the lock semantics are broken in the IF_POSSIBLE + incompatible path, which is the single most interesting case for a lock to protect (both runs are targeting the same dataset_name). Worth fixing.

Suggested fix: either (a) move the outer lock inside build() after the compat-check / mode switch and re-acquire a narrower lock around profiling, or (b) expose a helper that resolves the final directory without creating it, lock that, and pass the path through. Option (a) is probably simpler — the inner dataset_lock() already exists.

🟡 Redundant nested lock when called from DataDesigner.create()

With the outer lock in data_designer.py plus the inner dataset_lock() inside DatasetBuilder.build(), the inner lock is effectively a no-op re-entry through _lock_depth in the common path. It only matters when build() is called directly (tests, future callers). That's fine, but the two-layer locking is easy to misread. A comment at either site pointing out the re-entry contract would save the next reader some time.

🟡 _lock_depth counter is not thread-safe

_lock_depth is a bare int; increments/decrements are not atomic. Today DataDesigner runs single-threaded control flow, so this doesn't matter — but if the same ArtifactStorage is ever shared across threads (e.g. a future profiler that reads metadata from a worker thread), re-entrant locks could desynchronize. Worth a one-line docstring note that dataset_lock() is re-entrant per-process but not thread-safe.

🟡 Stale lock requires manual cleanup

dataset_lock() relies on unlink(missing_ok=True) in the finally block. If the process is SIGKILLed or the host crashes, the lock file remains and the next run fails with the "remove the stale lock if no process is active" message. The PID is written to the file, so a future enhancement could auto-detect staleness (e.g. psutil.pid_exists), but the current message is clear and actionable — flagging only as follow-up, not a blocker.

🟡 post_generation_processed flag is only set if there are AFTER_GENERATION processors

if self._processor_runner.has_processors_for(ProcessorStage.AFTER_GENERATION):
    self.artifact_storage.update_metadata({"post_generation_processed": True})

If a config has no AFTER_GENERATION processors, the flag is never set, so _post_generation_processed_resume_result never fires. That's correct (nothing was applied). But if a user re-runs the same dataset with a newly added AFTER_GENERATION processor, that processor will run on the completed dataset — which is what the PR description says is intended ("can change row counts, schemas, row-group boundaries"). Please confirm this is the desired behavior; a docstring or test pinning this contract would prevent future regressions.

🟢 _post_generation_processed_resume_result early-return bypasses task_traces initialization

When this method returns early, build() returns the final dataset path without entering the inner with block. data_designer.py:create() then accesses builder.task_traces, builder.actual_num_records, builder.first_non_retryable_error, and builder.early_shutdown on the builder. These need to be initialized to safe defaults (empty list, -1, None, False) before build() is called for the early-return path to work end-to-end. Worth adding an integration-style test that drives the whole DataDesigner.create() flow against a completed+post-processed dataset (the current test only patches _initialize_generators_and_graph on the builder).

🟢 read_metadata() used in two places, different exception-handling contracts

• _load_resume_state (dataset_builder.py:402) now catches json.JSONDecodeError and raises DatasetGenerationError with "metadata.json is corrupt."
• _post_generation_processed_resume_result catches (FileNotFoundError, json.JSONDecodeError) and silently returns None.

Catching JSONDecodeError silently in the post-generation check means a corrupt metadata slips past this gate and is only caught later in _load_resume_state. That's probably fine (the real resume path hits it), but it's asymmetric with the compat-check path which raises on JSONDecodeError. Consider routing all corrupt-metadata reads through a single helper so the behavior is consistent.

🟢 Style nit — typing.Iterator

artifact_storage.py imports Iterator from typing. The codebase already uses modern typing elsewhere; collections.abc.Iterator is preferred for runtime-usable generic aliases on 3.9+. Ruff's UP035 should flag this if it's in the rule set — not worth blocking on, but easy to fix.

🟢 Tests are solid

The new tests cover the important invariants:

• test_resume_rejects_allow_resize_columns
• test_build_resume_raises_on_corrupt_metadata
• test_build_resume_post_generation_processed_same_target_returns_existing_path
• test_build_resume_post_generation_processed_extension_raises
• test_initial_actual_num_records_uses_actual_parquet_rows_for_partial_row_group
• test_artifact_storage_write_metadata_includes_defaults
• test_artifact_storage_dataset_lock_rejects_second_writer

The parquet fixtures were upgraded from empty text files to real parquet with row counts, which is necessary for the new read_metadata(p).num_rows path. Good hygiene.

Missing coverage:

• End-to-end DataDesigner.create() flow against a completed+post-processed dataset (sanity-check for the early-return path through builder attributes).
• Lock behavior across the IF_POSSIBLE + incompatible-config mode switch (this is exactly the bug in the first finding).
• Thread-reentrancy behavior of _lock_depth, if the contract is "single-threaded" that should be documented/tested.

Docs

• architecture/dataset-builders.md — new "Resume Checkpointing" section is accurate and matches the code.
• fern/.../architecture-and-performance.mdx — clean summary of resume modes and invariants.
• fern/.../processors.mdx — the new warning about process_after_generation() being terminal is clear.

Verdict

Request changes (minor) before merge. The PR is a clear improvement in crash-recovery hardness, and most of the design is well-considered. Please address the outer-lock premature-caching issue (the IF_POSSIBLE + incompatible path leaves the run unlocked) before merging — that's the only finding I'd treat as blocking. The rest are follow-ups or style nits.

Greptile Summary

This PR hardens the resume/checkpoint system for both the sync and async dataset-builder engines. It unifies progress recovery — both engines now scan parquet-files/batch_*.parquet for actual row counts instead of reading stale metadata counters — and closes the narrow crash window where a complete dataset could silently skip process_after_generation.

• Atomic metadata writes replace the previous in-place json.dump: a tmp-file + fsync + os.replace sequence prevents corrupt checkpoints on crash mid-write, and config fingerprints are injected into every write via _metadata_defaults so compatibility checks no longer require deserializing builder_config.json.
• Unified _recover_progress_from_disk now drives both engines; the sync engine additionally enforces contiguous batch IDs and rejects allow_resize=True columns, while the async engine tolerates out-of-order completion.
• _post_generation_processed_resume_result intercepts the terminal post-processed state before any generation begins, surfacing clear DatasetGenerationError messages for the "already complete", "smaller target", and "extension" cases.

Important Files Changed

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["build(num_records, resume)"] --> B{resume mode?}
    B -- IF_POSSIBLE --> C["_check_resume_config_compatibility()"]
    C -- COMPATIBLE --> D["resume = ALWAYS"]
    C -- INCOMPATIBLE/NO_PRIOR --> E["resume = NEVER"]
    B -- ALWAYS --> F
    B -- NEVER --> F
    D --> F["_set_metadata_defaults()\n(inject config fingerprint into every write)"]
    E --> F
    F --> G["_post_generation_processed_resume_result()"]
    G -- "post_gen_started → raise\nnot post_gen processed → None" --> H["_write_builder_config()"]
    G -- "post_gen complete + same target → return final_dataset_path" --> Z["✅ no-op return"]
    G -- "post_gen complete + wrong target → raise" --> ERR["🛑 DatasetGenerationError"]
    H --> I{allow_resize?}
    I -- "Yes + ALWAYS → raise" --> ERR
    I -- No --> J["_resolve_async_compatibility()"]
    J --> K{engine?}
    K -- async --> L["_build_async()"]
    K -- sync + ALWAYS --> M["_build_with_resume()"]
    K -- sync + NEVER --> N["_build fresh batches"]
    L --> O["_load_resume_state()\n→ _recover_progress_from_disk(allow_holes=True)\n→ _find_completed_row_groups()"]
    M --> O
    O --> P{all batches done?}
    P -- Yes --> Q["log 'already complete', return"]
    P -- No --> R["generate remaining batches"]
    Q --> S
    R --> S["has_after_generation_processors?"]
    S -- No --> T["log_model_usage → return final_path"]
    S -- Yes --> U["update_metadata(started)\nrun_after_generation()\nupdate_metadata(complete)"]
    U --> T

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into fix/resume-hard..." | Re-trigger Greptile}

Thanks for putting this together, @nabinchha!

Summary

This PR hardens resume checkpoints by making metadata writes atomic, persisting config fingerprints, deriving resume progress from parquet files, and treating after-generation-processed datasets as terminal for resume. The implementation matches the stated intent, and the docs updates line up with the new resume model.

Findings

No findings.

What Looks Good

• The split between metadata as configuration state and parquet files as progress state is clear and closes the stale-metadata crash window.
• The post-generation processor guard is conservative in the right way: it marks the risky rewrite phase before mutation and avoids re-running processors on already-complete resumed datasets.
• The regression coverage is strong around config compatibility, missing/corrupt metadata, non-contiguous sync batches, async row counts from parquet metadata, and allow-resize rejection.

Verdict

Ship it. I verified the changed Python files with uv run ruff check and uv run ruff format --check, plus the targeted engine resume/storage tests and the touched interface tests. I could not run fern check because the Fern CLI is not installed in this environment.

This review was generated by an AI assistant.