From 1c8f5e96535652cd02532a675180e123495cf5eb Mon Sep 17 00:00:00 2001 From: Tanish Gudise <174986239+TanishGudise@users.noreply.github.com> Date: Sun, 11 Oct 2026 19:33:22 -0400 Subject: [PATCH 1/2] fix: don't load a partially written dataset file from the cache create_dataset_files wrote .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 --- .../base_request_processor.py | 26 ++++- tests/unittests/test_create_dataset_files.py | 106 ++++++++++++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) create mode 100644 tests/unittests/test_create_dataset_files.py diff --git a/src/bespokelabs/curator/request_processor/base_request_processor.py b/src/bespokelabs/curator/request_processor/base_request_processor.py index c3828af8..1067338b 100644 --- a/src/bespokelabs/curator/request_processor/base_request_processor.py +++ b/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 diff --git a/tests/unittests/test_create_dataset_files.py b/tests/unittests/test_create_dataset_files.py new file mode 100644 index 00000000..1db0fdfa --- /dev/null +++ b/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"] + ) From 2155c26cfa46847af78764829e5f23a7a0269d4a Mon Sep 17 00:00:00 2001 From: Tanish Gudise <174986239+TanishGudise@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:54:40 -0400 Subject: [PATCH 2/2] fix: remove the temporary dataset file when moving it into place fails If os.replace raised (for example a permission error, or a directory already at the dataset path), the completed temporary file was left in the cache directory. Run the move inside the try and remove the temporary file in finally, so it is cleaned up after any failure. Co-Authored-By: Claude Opus 5.5 --- .../request_processor/base_request_processor.py | 6 +++--- tests/unittests/test_create_dataset_files.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/bespokelabs/curator/request_processor/base_request_processor.py b/src/bespokelabs/curator/request_processor/base_request_processor.py index 1067338b..dec55768 100644 --- a/src/bespokelabs/curator/request_processor/base_request_processor.py +++ b/src/bespokelabs/curator/request_processor/base_request_processor.py @@ -41,11 +41,11 @@ def _write_atomically(path: str): tmp_path = f"{path}.{os.getpid()}.tmp" try: yield tmp_path - except BaseException: + os.replace(tmp_path, path) + finally: + # Already moved after a successful replace; otherwise the write or the move failed. if os.path.exists(tmp_path): os.remove(tmp_path) - raise - os.replace(tmp_path, path) class BaseRequestProcessor(ABC): diff --git a/tests/unittests/test_create_dataset_files.py b/tests/unittests/test_create_dataset_files.py index 1db0fdfa..878cf845 100644 --- a/tests/unittests/test_create_dataset_files.py +++ b/tests/unittests/test_create_dataset_files.py @@ -104,3 +104,19 @@ def test_dataset_is_rebuilt_from_cached_responses_after_interrupted_build(tmp_pa assert sorted(os.listdir(tmp_path)) == sorted( [f"{_PARSE_FUNC_HASH}.arrow", "failed_requests.jsonl", "metadata_0.json", "requests_0.jsonl", "responses_0.jsonl"] ) + + +def test_failed_move_into_place_leaves_no_temporary_file(tmp_path, monkeypatch): + """If the final rename fails (for example a permission error), the error surfaces and no dataset or temporary file is left behind.""" + _write_cached_run(tmp_path) + + def failing_replace(src, dst): + raise PermissionError(f"cannot move {src} to {dst}") + + with monkeypatch.context() as patch: + patch.setattr(os, "replace", failing_replace) + with pytest.raises(PermissionError): + _processor(tmp_path, _parse).create_dataset_files(_PARSE_FUNC_HASH) + + assert sorted(os.listdir(tmp_path)) == sorted(["failed_requests.jsonl", "metadata_0.json", "requests_0.jsonl", "responses_0.jsonl"]) + assert _processor(tmp_path, _parse).attempt_loading_cached_dataset(_PARSE_FUNC_HASH) is None