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

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

Write to a temporary file (named with the process id) and move it into place
only when the build finishes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: TanishGudise

src/bespokelabs/curator/request_processor/base_request_processor.py

@@ -1,6 +1,7 @@
 """Base module for request processing functionality."""
 
 import asyncio
+import contextlib
 import functools
 import glob
 import json
@@ -30,6 +31,23 @@
     from datasets import Dataset
 
 
+@contextlib.contextmanager
+def _write_atomically(path: str):
+    """Yield a temporary path and move it to `path` only if the block finishes without raising.
+
+    An interrupted or failed write then never leaves a partial file at `path`. The temporary name
+    includes the process id so two runs writing the same file don't remove or move each other's.
+    """
+    tmp_path = f"{path}.{os.getpid()}.tmp"
+    try:
+        yield tmp_path
+    except BaseException:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+        raise
+    os.replace(tmp_path, path)
+
+
 class BaseRequestProcessor(ABC):
     """Base class for all request processors.
 
@@ -448,10 +466,12 @@ def create_dataset_files(
         total_responses_count = 0
         failed_responses_count = 0
         error_sample = []
-        dataset_file = os.path.join(self.working_dir, f"{parse_func_hash}.arrow")
+        final_dataset_file = os.path.join(self.working_dir, f"{parse_func_hash}.arrow")
         from datasets.arrow_writer import ArrowWriter
 
-        with ArrowWriter(path=dataset_file) as writer:
+        # A partial dataset file would be loaded as the cached result by the next run, so write
+        # to a temporary file and move it into place once the dataset is complete.
+        with _write_atomically(final_dataset_file) as dataset_file, ArrowWriter(path=dataset_file) as writer:
             for responses_file in responses_files:
                 with open(responses_file, "r") as f_in:
                     for generic_response_string in f_in:
@@ -551,7 +571,7 @@ def create_dataset_files(
                     os.remove(dataset_file)
                     raise ValueError("Some requests do not have responses and require_all_responses is True.")
 
-        return self._load_from_dataset_file(dataset_file)
+        return self._load_from_dataset_file(final_dataset_file)
 
     def _load_from_dataset_file(self, dataset_file: str) -> "Dataset":
         from datasets import Dataset

tests/unittests/test_create_dataset_files.py

@@ -0,0 +1,106 @@
+import datetime
+import json
+import os
+
+import pytest
+from datasets import config as datasets_config
+
+from bespokelabs.curator.llm.prompt_formatter import PromptFormatter
+from bespokelabs.curator.request_processor.config import OnlineRequestProcessorConfig
+from bespokelabs.curator.request_processor.online.openai_online_request_processor import OpenAIOnlineRequestProcessor
+from bespokelabs.curator.types.generic_request import GenericRequest
+from bespokelabs.curator.types.generic_response import GenericResponse
+
+# The dataset writer flushes rows to disk in batches of this size. Failing after the first batch
+# means part of the dataset is already in the file when the build stops.
+_BATCH = datasets_config.DEFAULT_MAX_BATCH_SIZE
+_N_ROWS = _BATCH + 50
+_FAIL_AT_ROW = _BATCH + 10
+_PARSE_FUNC_HASH = "parse_func_hash"
+
+
+def _write_cached_run(working_dir, corrupt_row=None):
+    """Write the request and response files a finished online run leaves in its cache directory."""
+    now = datetime.datetime.now()
+    with open(os.path.join(working_dir, "requests_0.jsonl"), "w") as requests_file, open(os.path.join(working_dir, "responses_0.jsonl"), "w") as responses_file:
+        for idx in range(_N_ROWS):
+            request = GenericRequest(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": f"question {idx}"}],
+                original_row={"question": f"question {idx}"},
+                original_row_idx=idx,
+            )
+            requests_file.write(json.dumps(request.model_dump(), default=str) + "\n")
+            if idx == corrupt_row:
+                responses_file.write('{"response_message": "unterminated\n')
+                continue
+            response = GenericResponse(
+                response_message=f"answer {idx}",
+                raw_response={},
+                generic_request=request,
+                created_at=now,
+                finished_at=now,
+            )
+            responses_file.write(json.dumps(response.model_dump(), default=str) + "\n")
+    with open(os.path.join(working_dir, "metadata_0.json"), "w") as f:
+        json.dump({"num_jobs": _N_ROWS}, f)
+
+
+def _processor(working_dir, parse_func):
+    # create_dataset_files only needs the config, working directory and prompt formatter, so skip
+    # __init__, which makes a network call to read rate limits.
+    processor = object.__new__(OpenAIOnlineRequestProcessor)
+    processor.config = OnlineRequestProcessorConfig(model="gpt-4o-mini")
+    processor.working_dir = str(working_dir)
+    processor.prompt_formatter = PromptFormatter(model_name="gpt-4o-mini", prompt_func=lambda row: row["question"], parse_func=parse_func)
+    return processor
+
+
+def _parse(row, response):
+    return {"question": row["question"], "answer": response}
+
+
+def _interrupting_parse():
+    calls = {"n": 0}
+
+    def parse(row, response):
+        calls["n"] += 1
+        if calls["n"] == _FAIL_AT_ROW:
+            raise KeyboardInterrupt
+        return _parse(row, response)
+
+    return parse
+
+
+@pytest.mark.parametrize(
+    "corrupt_row, parse_func, error",
+    [
+        pytest.param(None, _interrupting_parse, KeyboardInterrupt, id="keyboard_interrupt"),
+        pytest.param(_FAIL_AT_ROW, lambda: _parse, ValueError, id="unreadable_response_line"),
+    ],
+)
+def test_interrupted_dataset_build_is_not_loaded_as_cached_dataset(tmp_path, corrupt_row, parse_func, error):
+    """A dataset build that stops partway must not leave a dataset file that the next run loads from the cache."""
+    _write_cached_run(tmp_path, corrupt_row=corrupt_row)
+
+    with pytest.raises(error):
+        _processor(tmp_path, parse_func()).create_dataset_files(_PARSE_FUNC_HASH)
+
+    assert _processor(tmp_path, _parse).attempt_loading_cached_dataset(_PARSE_FUNC_HASH) is None
+    # No partial dataset or temporary file is left in the cache directory.
+    assert sorted(os.listdir(tmp_path)) == ["metadata_0.json", "requests_0.jsonl", "responses_0.jsonl"]
+
+
+def test_dataset_is_rebuilt_from_cached_responses_after_interrupted_build(tmp_path):
+    _write_cached_run(tmp_path)
+    with pytest.raises(KeyboardInterrupt):
+        _processor(tmp_path, _interrupting_parse()).create_dataset_files(_PARSE_FUNC_HASH)
+
+    dataset = _processor(tmp_path, _parse).create_dataset_files(_PARSE_FUNC_HASH)
+
+    assert len(dataset) == _N_ROWS
+    cached = _processor(tmp_path, _parse).attempt_loading_cached_dataset(_PARSE_FUNC_HASH)
+    assert cached["answer"] == [f"answer {idx}" for idx in range(_N_ROWS)]
+    assert sorted(os.listdir(tmp_path)) == sorted(
+        [f"{_PARSE_FUNC_HASH}.arrow", "failed_requests.jsonl", "metadata_0.json", "requests_0.jsonl", "responses_0.jsonl"]
+    )