fix: don't load a partially written dataset file from the cache

[TanishGudise]

Closes #745

create_dataset_files wrote <parse_func_hash>.arrow in place. When the build stopped partway (Ctrl-C, or an exception such as an unreadable line in a responses file), ArrowWriter closed the file cleanly, so it held the rows flushed so far (in batches of 1,000) and still loaded. The next run returned that partial dataset as the cached result.

This writes the dataset to a temporary file, <parse_func_hash>.arrow.<pid>.tmp (the process id keeps two processes running the same job from touching each other's file), and moves it into place with os.replace only after the build finishes. If the build fails or is interrupted, or the final move fails, the temporary file is removed and the exception is re-raised. The existing error paths (all requests failed, require_all_responses, invalid rows) behave as before; they now remove the temporary file instead of the final one.

Tests (tests/unittests/test_create_dataset_files.py)

• An interrupted build (KeyboardInterrupt from parse_func, and an unreadable response line) leaves no dataset file and no temporary file, so attempt_loading_cached_dataset returns None. On main both cases load a 1,000-row partial dataset and fail.
• After an interrupted build, the next build returns every row, writes <hash>.arrow and leaves no temporary file behind (passes on main and with the fix).
• If moving the file into place fails (for example a permission error), the error is raised and neither the dataset file nor the temporary file is left behind (test_failed_move_into_place_leaves_no_temporary_file, added after review; fails on 1c8f5e9).

Each part of the fix removed on its own makes a test fail: no cleanup of the temporary file, moving it into place even when the build fails, and loading the temporary path at the end.

Manual check: end to end against a local OpenAI-compatible stub (2,500 rows, interrupted at row 2,200 of the build), the next identical run returned 2,000 rows on main and 2,500 with this change, with 0 new requests in both.

This doesn't repair a partial <hash>.arrow left by an earlier version; deleting that file makes the next run rebuild it from the cached responses. CodeExecutionBackend.create_dataset_files (code_executor/code_execution_backend/base_backend.py) writes and reloads its .arrow file the same way; I left it out to keep this PR small and can add it here or in a follow-up.

Without network access, the full test suite has no new failures compared with main: the offline failures are tests that need the network, plus test_simple_code_execution_local, which also fails locally on main.

I also checked this against the open #730: it merges without conflicts, and with both changes applied the tests from both PRs pass. #730's change on its own doesn't make the two new failing tests here pass.

CI note: CI on main is currently red for reasons unrelated to this PR: test_basic_reasoning and test_batch_reasoning fail because litellm 1.83.7 doesn't map claude-3-7-sonnet-20250219, and total coverage is 78%, under --cov-fail-under=80 (run 37956501158 on 31abdb7). Once CI runs here (first-time contributor runs wait for approval), those failures will show on this PR too until main is fixed.

🤖 Generated with Claude Code

[strix-security]

Strix Security Review

Warning

This pull request has 1 commit after the last Strix review (1c8f5e9). Strix has not reviewed these changes.
Automatic review on push is off for this repository. To review the latest changes, tag @strix-security in a comment, or turn on re-review on push.

No security issues found.

Review summary

Reviewed the atomic dataset-write change in base_request_processor.py and the accompanying unit tests. The new _write_atomically context manager writes the dataset to a per-process temporary file and moves it into place with os.replace only after a successful build, correctly removing the temporary file on failure or interruption (including KeyboardInterrupt), and the existing error paths were updated to clean up the temporary file. I traced the file-path inputs (parse_func_hash is an xxh64 hex digest; the working directory defaults under the user's cache directory) and confirmed no injection, authorization, deserialization, path-traversal, or credential issues are introduced. The only theoretical consideration — the predictable PID-based temporary filename without exclusive creation — does not present a concrete exploit in this local-library threat model, since an attacker would already require write access to the victim's own cache directory. No security issues requiring changes were found.

_{Updated for 1c8f5e9.}

---

Reviewed by Strix
_{Re-run review · Configure security review settings}

[entelligence-ai-pr-reviews]

EntelligenceAI PR Summary

Dataset creation now writes Arrow output to a process-specific temporary file and atomically moves it into the cache only after the build completes successfully, preventing interrupted or failed builds from being reused as valid cached datasets. Temporary files are cleaned up on write or rename failures, and new tests cover interruptions, malformed responses, successful rebuilds, and failed final moves.

flowchart TD
    classDef newBehavior fill:#dcfce7,stroke:#16a34a,color:#14532d;
    A["Cached responses"] --> B["Write dataset to process-specific temporary Arrow file"]
    B --> C["Parse responses and build dataset"]
    C --> D["Atomic replace temporary file with final cache file"]
    D --> E["Load completed dataset from cache"]
    C --> F["Build fails or is interrupted"]
    F --> G["Remove temporary file; leave no final dataset"]
    D --> H["Rename fails"]
    H --> G
    class B,C,D,E,F,G,H newBehavior;

Loading

🟢 Green = new or changed in this PR

---

Review Scorecard

Dimension	Rating	Basis
Code Quality	●●●●● 5/5 — Excellent	no critical/significant findings — rated Excellent
Blast Radius	Low	changed symbols are referenced only within their own file(s); no high-impact surface touched, 1 file(s) / ~26 line(s) changed (size only — not a blast signal)
Merge Confidence	●●●●● 5/5 — Safe to Merge	code quality 5/5 × Low blast radius

Safe to merge with a low blast radius; code quality rated Excellent (5/5): the review found nothing to flag.

Key Findings:

• No critical or significant issues found

Evaluated against

• 2/2 changed files reviewed
• criteria: correctness, security & access control, robustness & error handling, concurrency & data integrity, repo conventions / steering docs
• steering docs: none found in repo

[entelligence-ai-pr-reviews]

Walkthrough

This PR adds atomic persistence for generated Arrow datasets, ensuring interrupted or failed writes cannot leave reusable partial cache files. It updates dataset creation to finalize files safely and reload the committed path, with comprehensive tests covering parsing failures, corrupt responses, cache rebuilds, rename failures, cleanup, and error propagation.

Changes

Files

• src/bespokelabs/curator/request_processor/base_request_processor.py

Added the _write_atomically context manager, which writes to a process-specific temporary file and atomically replaces the destination only after successful completion. Updated create_dataset_files to use atomic writes for Arrow dataset generation, clean up failed artifacts, and load the finalized dataset path after persistence.

Files

• tests/unittests/test_create_dataset_files.py

Added pytest coverage for interrupted parsing, corrupt response lines, successful dataset rebuilding from cached responses, failures during atomic renames, cleanup of temporary and partial files, rejection of incomplete cache hits, dataset reloadability, and error propagation.

[TanishGudise]

Closing as superseded by #753, which fixes #745. I ran this PR's regression tests against main after #753 (interrupted build, unreadable response line, rebuild from cached responses, failed move into place) and all four pass. Thanks!