From 84e6e3480be22b26ae9365ef93122dff663741cb Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 08:44:16 -0400 Subject: [PATCH 01/14] fix(anthropic): stop sending invalid output_config.format.name (HTTP 400) The Anthropic client placed a `name` key inside output_config.format, which Anthropic's API forbids, so every native-schema call (parse / chat_structured / response(response_format=...) and the batch structured path) failed with HTTP 400. Build the request via the SDK's typed OutputConfigParam and drop `name`, which is still required by the OpenAI wire object and the neutral spec. Flip the unit test into a regression guard and add a batch-path regression test. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 1 + src/prkit/core/model_clients/anthropic.py | 12 +++++-- .../core/model_clients/test_anthropic.py | 32 ++++++++++++++++++- 3 files changed, 41 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c8d86b..424a4b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,7 @@ Production releases follow semantic versioning. TestPyPI validation builds use P - **`DatasetHub` registration-ordering bug** — calling `DatasetHub.register(name, Loader)` before any built-in dataset was touched caused all built-in loaders and downloaders to be permanently suppressed. Built-ins are now seeded idempotently (via `setdefault`) at the start of every public mutating method, so external registrations can safely happen in any order. - JEEBench loader handling for numeric answer categories and retained metadata. - Workflow module behavior in domain assessment, theorem review, and workflow composition paths. +- **Anthropic structured output 400** — the Anthropic client no longer sends an invalid `name` key inside `output_config.format` (a key Anthropic's API forbids), which previously made every native-schema call — `parse()`, `chat_structured()`, `response(response_format=...)`, and the batch structured path — fail with HTTP 400. The request is now built via the SDK's typed `OutputConfigParam`, so future schema drift surfaces as a type error rather than a runtime 400. OpenAI (which requires `name`) and Gemini are unaffected. ### Deprecated diff --git a/src/prkit/core/model_clients/anthropic.py b/src/prkit/core/model_clients/anthropic.py index 700c325..94e4ba2 100644 --- a/src/prkit/core/model_clients/anthropic.py +++ b/src/prkit/core/model_clients/anthropic.py @@ -5,7 +5,7 @@ import os import re from collections.abc import Callable, Iterator, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any from pydantic import BaseModel @@ -34,6 +34,9 @@ except ImportError: # pragma: no cover - tested via runtime error path pass +if TYPE_CHECKING: + from anthropic.types import OutputConfigParam + TOOL_NAME = "emit_structured_output" _TOOL_NAME_RE = re.compile(r"[^A-Za-z0-9_-]+") ANTHROPIC_OPTIONAL_PARAMETER_LIMIT = 24 @@ -329,13 +332,16 @@ def _build_messages_params( params["system"] = instructions if response_format is not None: normalized = normalize_response_format(response_format) - params["output_config"] = { + # Anthropic's output_config.format accepts only ``type`` and ``schema``; + # a ``name`` key (an OpenAI-ism) makes the API reject the request with a + # 400. The typed annotation makes any future stray key a mypy error. + output_config: OutputConfigParam = { "format": { "type": "json_schema", - "name": normalized["name"], "schema": normalized["schema"], } } + params["output_config"] = output_config if extra: params.update(extra) return params diff --git a/tests/prkit/core/model_clients/test_anthropic.py b/tests/prkit/core/model_clients/test_anthropic.py index 2e5f341..a589fa5 100644 --- a/tests/prkit/core/model_clients/test_anthropic.py +++ b/tests/prkit/core/model_clients/test_anthropic.py @@ -14,6 +14,7 @@ _parse_data_url, ) from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS +from prkit.core.model_clients.structured_output import coerce_structured_output_spec ANTHROPIC_TEST_MODEL = "claude-sonnet-4-6" @@ -165,10 +166,39 @@ class ExampleResponse(BaseModel): assert response == '{"answer":"Response"}' call_kwargs = mock_client.messages.create.call_args.kwargs assert call_kwargs["output_config"]["format"]["type"] == "json_schema" - assert call_kwargs["output_config"]["format"]["name"] == "ExampleResponse" + # Anthropic forbids a ``name`` key in output_config.format (400 otherwise). + assert "name" not in call_kwargs["output_config"]["format"] assert call_kwargs["output_config"]["format"]["schema"]["type"] == "object" assert "tools" not in call_kwargs + @patch("prkit.core.model_clients.anthropic.Anthropic") + def test_batch_structured_request_omits_name(self, mock_anthropic_class): + """Batch structured requests must also omit Anthropic's forbidden `name` key.""" + mock_client = MagicMock() + mock_anthropic_class.return_value = mock_client + + class ExampleResponse(BaseModel): + answer: str + + client = AnthropicModel(ANTHROPIC_TEST_MODEL) + spec = coerce_structured_output_spec(ExampleResponse) + plan = client._resolve_structured_output_plan( + spec, structured_policy="native_required" + ) + request = client._build_batch_structured_request( + request_id="r1", + user_prompt="hi", + response_model=ExampleResponse, + image_paths=(), + max_output_tokens=None, + plan=plan, + ) + + fmt = request["params"]["output_config"]["format"] + assert fmt["type"] == "json_schema" + assert "name" not in fmt # Anthropic forbids it (400 otherwise) + assert fmt["schema"]["type"] == "object" + def test_extract_tool_use_json_requires_exactly_one_tool_block(self): with pytest.raises(ValueError, match="exactly one tool_use block"): _extract_tool_use_json([{"type": "text", "text": "not structured"}]) From 3065ba1d060c0a9dcf801c814f87407befc91a97 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 11:49:52 -0400 Subject: [PATCH 02/14] Add prkit.batch submit module and BaseModelClient batch facade Introduce the N4 Stage 1 (submit) surface: a new leaf module prkit.batch with submit_batch_physics_reasoning(client, dataset_or_problems, ...) -> list[BatchSubmission], the dumps_batch_jsonl/write_batch_jsonl/validate_batch_requests helpers, and the BatchInputError/BatchSubmitError exceptions. The orchestrator preprocesses each problem like the synchronous path, splits the dataset into provider batches, writes provider-correct JSONL under one run folder per call, submits each batch via the existing client.submit_batch, and writes a run-level metadata.json before returning typed receipts. Add two methods to BaseModelClient: build_problem_batch_request (free-text batch analogue of solve_physics_problem, reusing build_plain_question_prompt for batch == sync prompt parity) and a thin submit_batch_physics_reasoning facade that lazily imports prkit.batch. The leaf imports only prkit.core.domain + stdlib at module load; the client is duck-typed. Tests run fully offline (MagicMock+patch the provider SDKs) and cover splitting, JSONL artifacts, submission, id validation/correlation, surrogate ids, partial failure, the facade, the Anthropic inline path, run-folder layout, metadata.json content, id round-trips, defaults/overrides, and import isolation. Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 362 +++++++++++++++ src/prkit/core/model_clients/base.py | 70 ++- tests/prkit/batch/__init__.py | 0 .../batch/test_build_problem_batch_request.py | 120 +++++ tests/prkit/batch/test_import_isolation.py | 55 +++ .../batch/test_submit_physics_reasoning.py | 416 ++++++++++++++++++ 6 files changed, 1022 insertions(+), 1 deletion(-) create mode 100644 src/prkit/batch/__init__.py create mode 100644 tests/prkit/batch/__init__.py create mode 100644 tests/prkit/batch/test_build_problem_batch_request.py create mode 100644 tests/prkit/batch/test_import_isolation.py create mode 100644 tests/prkit/batch/test_submit_physics_reasoning.py diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py new file mode 100644 index 0000000..0865c24 --- /dev/null +++ b/src/prkit/batch/__init__.py @@ -0,0 +1,362 @@ +"""Batch-mode *submit* for physics-reasoning runs (N4 Stage 1). + +This leaf turns a set of :class:`~prkit.core.domain.PhysicsProblem`\\s into +submitted provider batch jobs and returns one typed receipt +(:class:`BatchSubmission`) per batch. It preprocesses each problem **exactly** +like the synchronous ``solve_physics_problem`` path (via the client's +``build_problem_batch_request``), splits the dataset into provider batches of +``batch_size`` problems, writes each batch's requests as provider-correct JSONL +under one run folder per call, submits each batch, and writes a run-level +``metadata.json`` audit record **before** returning. + +It is a *bounded submitter*: it stops at submitted receipts. Fetching / +polling / scoring / pricing is a later milestone that consumes these receipts. + +Import discipline: at module load this imports only :mod:`prkit.core.domain` +and the standard library — never ``prkit.api``, the dataset hub, a scorer, the +cost meter, or a provider SDK. The model client is duck-typed. +""" + +from __future__ import annotations + +import json +import re +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from prkit.core.domain import PhysicsDataset, PhysicsProblem + +__all__ = [ + "BatchSubmission", + "BatchInputError", + "BatchSubmitError", + "submit_batch_physics_reasoning", + "dumps_batch_jsonl", + "write_batch_jsonl", + "validate_batch_requests", +] + +# Providers whose batch line correlates by ``key`` rather than ``custom_id``. +_KEY_ID_PROVIDERS = frozenset({"google", "gemini"}) +# Anthropic restricts custom ids to this charset/length. +_ANTHROPIC_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") +_MAX_ID_LEN = 64 + + +@dataclass(frozen=True) +class BatchSubmission: + """One submission receipt = one batch = one provider batch job. + + The normalized provider id is in :attr:`batch_id` (empty string when submit + failed — the batch is then resumable via :attr:`input_file_path`, whose JSONL + is always written). :attr:`id_map` maps each wire ``custom_id`` / ``key`` back + to its ``problem_id`` so a later fetch step correlates results to problems + even when surrogate ids are used. + """ + + provider: str + model: str + batch_id: str + input_file_path: str + submitted_at: datetime + num_requests: int + batch_index: int + id_map: dict[str, str] + endpoint: str | None = None + completion_window: str | None = None + metadata_path: str | None = None + metadata: dict[str, str] = field(default_factory=dict) + error: str | None = None + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable dict (``submitted_at`` rendered ISO 8601).""" + return { + "provider": self.provider, + "model": self.model, + "batch_id": self.batch_id, + "input_file_path": self.input_file_path, + "submitted_at": self.submitted_at.isoformat(), + "num_requests": self.num_requests, + "batch_index": self.batch_index, + "id_map": dict(self.id_map), + "endpoint": self.endpoint, + "completion_window": self.completion_window, + "metadata_path": self.metadata_path, + "metadata": dict(self.metadata), + "error": self.error, + } + + +class BatchInputError(ValueError): + """Invalid submit input: empty problem set, duplicate/illegal id, or a + pre-existing non-empty run folder (when ``overwrite`` is False).""" + + +class BatchSubmitError(RuntimeError): + """A batch submission failed. + + Carries the receipts that succeeded before the failure plus the failed + receipt so a consumer can resume by re-submitting the failed batch (its + input JSONL already exists). The Stage-1 submit loop records per-batch + failures on the receipts and returns the full list rather than raising; this + type is the resume-carrying error reserved for that tooling. + """ + + def __init__( + self, + message: str, + *, + successes: list[BatchSubmission], + failed: BatchSubmission, + ) -> None: + super().__init__(message) + self.successes = successes + self.failed = failed + + +def _slug(text: str) -> str: + """Lowercase, collapse non-``[a-z0-9._-]`` runs to ``-``, trim edges. + + Falls back to ``"run"`` when the result would be empty. + """ + slugged = re.sub(r"[^a-z0-9._-]+", "-", text.strip().lower()).strip("-") + return slugged or "run" + + +def dumps_batch_jsonl(requests: Sequence[dict[str, Any]]) -> str: + """Serialize *requests* as JSON Lines (one object per line). + + Uses ``ensure_ascii=False`` so non-ASCII (e.g. Chinese) physics problems stay + readable in the artifact. This artifact is for inspection / resume; the + provider upload re-serializes from the request list, so byte-identity is not + required. + """ + return "\n".join(json.dumps(request, ensure_ascii=False) for request in requests) + + +def write_batch_jsonl(requests: Sequence[dict[str, Any]], path: str | Path) -> Path: + """Write *requests* as JSONL to *path* (creating parents) and return the path.""" + out = Path(path) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(dumps_batch_jsonl(requests) + "\n", encoding="utf-8") + return out + + +def validate_batch_requests( + requests: Sequence[dict[str, Any]], *, provider: str +) -> None: + """Fail-fast validation of correlation ids before any file/network I/O. + + Raises :class:`BatchInputError` on an empty request set, a missing/empty id, + a duplicate id across the whole run, or an id that violates the active + provider's length/charset rule (OpenAI ≤ 64 chars; Anthropic + ``^[a-zA-Z0-9_-]{1,64}$``). Gemini and unknown providers require only a + non-empty unique id. + """ + if not requests: + raise BatchInputError("No batch requests to submit (empty input).") + + id_field = "key" if provider in _KEY_ID_PROVIDERS else "custom_id" + seen: set[str] = set() + for index, request in enumerate(requests): + raw = request.get(id_field) + if not isinstance(raw, str) or not raw: + raise BatchInputError( + f"Batch request at index {index} has a missing/empty {id_field!r}." + ) + if raw in seen: + raise BatchInputError( + f"Duplicate correlation id {raw!r} across the run; ids must be unique." + ) + seen.add(raw) + + if provider == "anthropic": + if not _ANTHROPIC_ID_RE.match(raw): + raise BatchInputError( + f"Id {raw!r} is invalid for Anthropic; must match " + r"^[a-zA-Z0-9_-]{1,64}$." + ) + elif provider == "openai": + if len(raw) > _MAX_ID_LEN: + raise BatchInputError( + f"Id {raw!r} exceeds OpenAI's {_MAX_ID_LEN}-char limit " + f"(len={len(raw)})." + ) + + +def _chunked(items: list[Any], size: int) -> list[list[Any]]: + """Split *items* into consecutive chunks of at most *size*.""" + if size <= 0: + raise BatchInputError(f"batch_size must be positive; got {size}.") + return [items[i : i + size] for i in range(0, len(items), size)] + + +def _prkit_api_version() -> str: + """Return the contract version, read lazily to keep this module leaf-light.""" + try: + from prkit.api import API_VERSION + + return API_VERSION + except Exception: # pragma: no cover - defensive; api import should not fail + return "1.0" + + +def submit_batch_physics_reasoning( + client: Any, + problems: PhysicsDataset | Sequence[PhysicsProblem], + *, + output_dir: str | Path = "batch_runs", + run_name: str | None = None, + batch_size: int = 500, + instructions: str | None = None, + max_output_tokens: int | None = None, + temperature: float | None = None, + custom_id_fn: Callable[[PhysicsProblem], str] | None = None, + display_name: str | None = None, + metadata: dict[str, str] | None = None, + overwrite: bool = False, +) -> list[BatchSubmission]: + """Preprocess, split, write, and submit *problems* as provider batch jobs. + + Builds one free-text request per problem (mirroring the synchronous + ``solve_physics_problem`` preprocessing via ``client.build_problem_batch_request``), + splits them into batches of ``batch_size``, writes each batch's provider-correct + JSONL under ``//inputs/batch_NNNN.jsonl``, submits each + batch sequentially via ``client.submit_batch``, writes a run-level + ``metadata.json`` (run header + submissions) **before** returning, and returns + one :class:`BatchSubmission` per batch. + + A batch whose submit raises is recorded with ``error`` set and ``batch_id=""`` + (its input file remains for resume); the returned list always has one entry + per batch. + + Raises: + BatchInputError: empty input, an invalid/duplicate id, a non-positive + ``batch_size``, or a pre-existing non-empty run folder when + ``overwrite`` is False. + """ + if isinstance(problems, PhysicsDataset): + dataset_name: str | None = problems.name + dataset_version: Any = problems.get_info().get("version") + problem_list: list[PhysicsProblem] = list(problems) + else: + dataset_name = None + dataset_version = None + problem_list = list(problems) + + if not problem_list: + raise BatchInputError("No problems to submit (empty dataset/sequence).") + + now = datetime.now(timezone.utc) + timestamp = now.strftime("%Y%m%d-%H%M%S") + if run_name is None: + base = ( + f"{dataset_name}-{client.model}-{timestamp}" + if dataset_name + else f"{client.model}-{timestamp}" + ) + run_name = _slug(base) + else: + run_name = _slug(run_name) + display_name = display_name or run_name + + run_dir = Path(output_dir) / run_name + if run_dir.exists() and any(run_dir.iterdir()): + if not overwrite: + raise BatchInputError( + f"Run folder {str(run_dir)!r} already exists and is non-empty; " + "pass overwrite=True to write into it." + ) + for stale in (run_dir / "inputs").glob("batch_*.jsonl"): + stale.unlink() + inputs_dir = run_dir / "inputs" + inputs_dir.mkdir(parents=True, exist_ok=True) + + # Build one request per problem, keeping ids parallel for id_map + validation. + requests: list[dict[str, Any]] = [] + request_ids: list[str] = [] + problem_ids: list[str] = [] + for problem in problem_list: + request_id = custom_id_fn(problem) if custom_id_fn else problem.problem_id + requests.append( + client.build_problem_batch_request( + problem, + request_id=request_id, + instructions=instructions, + max_output_tokens=max_output_tokens, + temperature=temperature, + ) + ) + request_ids.append(request_id) + problem_ids.append(problem.problem_id) + + provider = client.provider + validate_batch_requests(requests, provider=provider) + + request_chunks = _chunked(requests, batch_size) + rid_chunks = _chunked(request_ids, batch_size) + pid_chunks = _chunked(problem_ids, batch_size) + + # Write all JSONL artifacts first (cheap, local), then submit sequentially. + file_paths: list[Path] = [] + for index, chunk in enumerate(request_chunks): + file_paths.append( + write_batch_jsonl(chunk, inputs_dir / f"batch_{index:04d}.jsonl") + ) + + metadata_path = run_dir / "metadata.json" + merged_metadata = {**(metadata or {}), "display_name": display_name} + + submissions: list[BatchSubmission] = [] + for index, chunk in enumerate(request_chunks): + id_map = dict(zip(rid_chunks[index], pid_chunks[index])) + submitted_at = datetime.now(timezone.utc) + batch_id = "" + error: str | None = None + try: + batch_id = client.submit_batch(chunk, metadata=merged_metadata) + except Exception as exc: # noqa: BLE001 - recorded on the receipt for resume + error = f"{type(exc).__name__}: {exc}" + submissions.append( + BatchSubmission( + provider=provider, + model=client.model, + batch_id=batch_id, + input_file_path=str(file_paths[index]), + submitted_at=submitted_at, + num_requests=len(chunk), + batch_index=index, + id_map=id_map, + metadata_path=str(metadata_path), + metadata=merged_metadata, + error=error, + ) + ) + + dataset_field: dict[str, Any] | None + if dataset_name is not None: + dataset_field = {"name": dataset_name, "version": dataset_version} + else: + dataset_field = None + + run_record = { + "run_id": run_name, + "created_at": now.isoformat(), + "provider": provider, + "model": client.model, + "dataset": dataset_field, + "batch_size": batch_size, + "total_problems": len(problem_list), + "num_batches": len(request_chunks), + "request_kind": "free_text", + "prkit_api_version": _prkit_api_version(), + "submissions": [submission.to_dict() for submission in submissions], + } + metadata_path.write_text( + json.dumps(run_record, indent=2, ensure_ascii=False), encoding="utf-8" + ) + + return submissions diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index b259acb..3b213c5 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -10,7 +10,7 @@ from pydantic import BaseModel, ValidationError -from ..domain import PhysicsProblem +from ..domain import PhysicsDataset, PhysicsProblem from ..logging_config import PRKitLogger from ..project_env import load_project_dotenv from .batch_types import BatchResult, BatchStatus @@ -28,6 +28,7 @@ ) if TYPE_CHECKING: + from prkit.batch import BatchSubmission from prkit.semantics.schema import PhysicsQuestionSemantics T = TypeVar("T", bound=BaseModel) @@ -381,6 +382,73 @@ def build_batch_request( **kwargs, ) + def build_problem_batch_request( + self, + problem: PhysicsProblem | str, + *, + request_id: str | None = None, + instructions: str | None = None, + max_output_tokens: int | None = None, + temperature: float | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Build a free-text batch request line for *problem*. + + Batch analogue of :meth:`solve_physics_problem`, dispatching on the + runtime type of *problem* the same way: a ``str`` is treated as the + question text (no images, and *request_id* is required since there is no + ``problem_id``); a :class:`~prkit.core.domain.PhysicsProblem` is rendered + with :func:`build_plain_question_prompt` and its ``image_path`` images, + defaulting *request_id* to ``problem.problem_id``. Reusing the same prompt + builder and image path as the synchronous solver guarantees batch ≡ sync + prompts. Free-text only, matching the ANSWER_TEXT-only sync path; delegates + to :meth:`build_batch_request`. + """ + if isinstance(problem, str): + if request_id is None: + raise ValueError( + "request_id is required when problem is a raw question string." + ) + prompt = problem.strip() + image_paths: list[str] | None = None + elif isinstance(problem, PhysicsProblem): + prompt = build_plain_question_prompt(problem) + image_paths = problem.image_path or None + if request_id is None: + request_id = problem.problem_id + else: + raise TypeError( + f"Unsupported problem input type: {type(problem)!r}. Expected a " + "PhysicsProblem or a question string." + ) + + return self.build_batch_request( + request_id=request_id, + input=prompt, + instructions=instructions, + image_paths=image_paths, + max_output_tokens=max_output_tokens, + temperature=temperature, + **kwargs, + ) + + def submit_batch_physics_reasoning( + self, + problems: PhysicsDataset | Sequence[PhysicsProblem], + **kwargs: Any, + ) -> list[BatchSubmission]: + """Submit *problems* as provider batch jobs; return one receipt per batch. + + Thin one-line facade mirroring :meth:`solve_physics_problem`: lazily imports + :mod:`prkit.batch` (kept off the import-light path, like the lazy + ``prkit.semantics`` import in :meth:`solve_physics_problem`) and delegates to + :func:`prkit.batch.submit_batch_physics_reasoning`. See that function for the + keyword arguments (``output_dir`` / ``run_name`` / ``batch_size`` / …). + """ + from prkit.batch import submit_batch_physics_reasoning as _submit_batch + + return _submit_batch(self, problems, **kwargs) + def _build_batch_request( self, *, diff --git a/tests/prkit/batch/__init__.py b/tests/prkit/batch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/prkit/batch/test_build_problem_batch_request.py b/tests/prkit/batch/test_build_problem_batch_request.py new file mode 100644 index 0000000..6febdfe --- /dev/null +++ b/tests/prkit/batch/test_build_problem_batch_request.py @@ -0,0 +1,120 @@ +"""Tests for ``BaseModelClient.build_problem_batch_request`` (batch ≡ sync prompts). + +Provider SDKs are faked with ``MagicMock`` so these run fully offline, mirroring +the pattern in ``tests/prkit/core/model_clients/test_batch.py``. The parity tests +assert that the batch line carries the SAME prompt the synchronous +``format_problem_context`` builds, across the OpenAI / Anthropic / Gemini line shapes. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.core.domain import PhysicsProblem +from prkit.core.model_clients import format_problem_context + + +def _openai_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _anthropic_client(model: str = "claude-opus-4-8"): + with patch("prkit.core.model_clients.anthropic.Anthropic") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.anthropic import AnthropicModel + + return AnthropicModel(model) + + +def _gemini_client(model: str = "gemini-3.5-flash"): + with patch("prkit.core.model_clients.gemini.genai.Client") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.gemini import GeminiModel + + return GeminiModel(model) + + +def _problem() -> PhysicsProblem: + return PhysicsProblem( + problem_id="prob-1", + question="A block slides down a frictionless incline. Find its acceleration.", + problem_type="OE", + domain="mechanics", + ) + + +class TestPromptParity: + def test_openai_prompt_equals_format_problem_context(self): + problem = _problem() + req = _openai_client().build_problem_batch_request(problem, instructions="") + text = req["body"]["input"][0]["content"][0]["text"] + assert text == format_problem_context(problem) + + def test_anthropic_prompt_equals_format_problem_context(self): + problem = _problem() + req = _anthropic_client().build_problem_batch_request(problem, instructions="") + text = req["params"]["messages"][0]["content"][0]["text"] + assert text == format_problem_context(problem) + + def test_gemini_prompt_equals_format_problem_context(self): + problem = _problem() + req = _gemini_client().build_problem_batch_request(problem, instructions="") + text = req["request"]["contents"][0]["parts"][0]["text"] + assert text == format_problem_context(problem) + + +class TestCorrelationAndImages: + def test_request_id_defaults_to_problem_id(self): + problem = _problem() + req = _openai_client().build_problem_batch_request(problem, instructions="") + assert req["custom_id"] == "prob-1" + + def test_explicit_request_id_overrides(self): + problem = _problem() + req = _openai_client().build_problem_batch_request( + problem, request_id="surrogate-9", instructions="" + ) + assert req["custom_id"] == "surrogate-9" + + def test_image_paths_forwarded_from_problem(self): + # Patch the delegate so we assert the forwarded image_paths without any + # file I/O (the provider image encoders would otherwise read the file). + client = _openai_client() + client.build_batch_request = MagicMock(return_value={}) + problem = PhysicsProblem( + problem_id="img-1", question="Q", image_path=["/tmp/a.png", "/tmp/b.png"] + ) + client.build_problem_batch_request(problem) + _, kwargs = client.build_batch_request.call_args + assert kwargs["image_paths"] == ["/tmp/a.png", "/tmp/b.png"] + assert kwargs["input"] == format_problem_context(problem) + + def test_no_images_forwards_none(self): + client = _openai_client() + client.build_batch_request = MagicMock(return_value={}) + client.build_problem_batch_request(_problem()) + _, kwargs = client.build_batch_request.call_args + assert kwargs["image_paths"] is None + + +class TestStringAndUnsupported: + def test_str_input_requires_request_id(self): + with pytest.raises(ValueError, match="request_id is required"): + _openai_client().build_problem_batch_request("just a question") + + def test_str_input_uses_stripped_text(self): + req = _openai_client().build_problem_batch_request( + " hello world ", request_id="r1", instructions="" + ) + assert req["custom_id"] == "r1" + assert req["body"]["input"][0]["content"][0]["text"] == "hello world" + + def test_unsupported_type_raises_type_error(self): + with pytest.raises(TypeError, match="Unsupported problem input type"): + _openai_client().build_problem_batch_request(123, request_id="r1") diff --git a/tests/prkit/batch/test_import_isolation.py b/tests/prkit/batch/test_import_isolation.py new file mode 100644 index 0000000..9f6f690 --- /dev/null +++ b/tests/prkit/batch/test_import_isolation.py @@ -0,0 +1,55 @@ +"""Import-boundary guard: ``prkit.batch`` must stay a light leaf at module load. + +Runs in a fresh subprocess (so the host test process's own imports cannot mask a +leak) and asserts that merely importing ``prkit.batch`` never pulls in a provider +SDK, the dataset hub, the semantics/api packages, ``pint``, ``pandas``, or even +``prkit.core.model_clients``. The submit orchestrator depends only on +``prkit.core.domain`` + stdlib at load; the client is duck-typed and the +``prkit.api`` version read is lazy (call-time, not import-time). +""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + +# Heavy/optional modules that must NOT be importable as a side effect of merely +# importing the leaf. ``google.genai`` (the provider SDK) is checked, never bare +# ``google`` (a namespace-package ``.pth`` artifact present at interpreter start). +_FORBIDDEN = [ + "anthropic", + "openai", + "google.genai", + "pint", + "datasets", + "pandas", + "prkit.api", + "prkit.datasets", + "prkit.semantics", + "prkit.core.model_clients", +] + + +def test_batch_import_does_not_pull_heavy_deps(): + code = textwrap.dedent(f""" + import sys + import prkit.batch + from prkit.batch import submit_batch_physics_reasoning, BatchSubmission + + forbidden = {_FORBIDDEN!r} + leaked = [name for name in forbidden if name in sys.modules] + if leaked: + print("LEAKED:" + ",".join(leaked)) + raise SystemExit(1) + raise SystemExit(0) + """) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + "heavy deps leaked into the prkit.batch import path:\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) diff --git a/tests/prkit/batch/test_submit_physics_reasoning.py b/tests/prkit/batch/test_submit_physics_reasoning.py new file mode 100644 index 0000000..45eb2ef --- /dev/null +++ b/tests/prkit/batch/test_submit_physics_reasoning.py @@ -0,0 +1,416 @@ +"""Tests for ``prkit.batch.submit_batch_physics_reasoning`` (the submit orchestrator). + +Provider SDKs are faked with ``MagicMock`` so these run fully offline (mirroring +``tests/prkit/core/model_clients/test_batch.py``); artifacts are written under +pytest's ``tmp_path``. The tests cover splitting, JSONL artifacts, submission, +validation, surrogate ids, partial failure, the facade, the Anthropic inline +path, the run-folder layout, ``metadata.json`` content, id round-trips, defaults, +and overwrite behavior. +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.batch import ( + BatchInputError, + BatchSubmission, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem + + +# --------------------------------------------------------------------------- # +# Offline client builders (SDK constructor patched -> no key / network needed) # +# --------------------------------------------------------------------------- # +def _openai_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + client = OpenAIModel(model) + client.client.files.create.return_value = MagicMock(id="file_1") + client.client.batches.create.return_value = MagicMock(id="batch_abc") + return client + + +def _anthropic_client(model: str = "claude-opus-4-8"): + with patch("prkit.core.model_clients.anthropic.Anthropic") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.anthropic import AnthropicModel + + client = AnthropicModel(model) + client.client.messages.batches.create.return_value = MagicMock(id="msgbatch_1") + return client + + +def _gemini_client(model: str = "gemini-3.5-flash"): + with patch("prkit.core.model_clients.gemini.genai.Client") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.gemini import GeminiModel + + client = GeminiModel(model) + uploaded = MagicMock() + uploaded.name = "files/in" + client.genai_client.files.upload.return_value = uploaded + job = MagicMock() + job.name = "batches/xyz" + client.genai_client.batches.create.return_value = job + return client + + +def _dataset(n: int, *, name: str = "physreason", version: str = "1.0"): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": name, "version": version}) + + +def _problems(n: int): + return [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + + +# --------------------------------------------------------------------------- # +class TestSplitting: + def test_1200_problems_split_into_three_batches(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(1200), output_dir=tmp_path, batch_size=500 + ) + assert [s.num_requests for s in subs] == [500, 500, 200] + assert [s.batch_index for s in subs] == [0, 1, 2] + run_dir = tmp_path / subs[0].input_file_path.split("/")[-3] + files = sorted(p.name for p in (run_dir / "inputs").glob("*.jsonl")) + assert files == ["batch_0000.jsonl", "batch_0001.jsonl", "batch_0002.jsonl"] + + +class TestJsonlArtifact: + def test_one_json_object_per_line_roundtrips(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, batch_size=2 + ) + lines = ( + (tmp_path / subs[0].input_file_path.split("/")[-3] / "inputs") + .joinpath("batch_0000.jsonl") + .read_text() + .splitlines() + ) + assert len(lines) == 2 + objs = [json.loads(line) for line in lines] + assert [o["custom_id"] for o in objs] == ["p0", "p1"] + assert all("body" in o for o in objs) + + +class TestSubmission: + def test_submit_called_once_per_batch_with_its_chunk(self, tmp_path): + client = _openai_client() + client.submit_batch = MagicMock(side_effect=["b0", "b1", "b2"]) + submit_batch_physics_reasoning( + client, _dataset(5), output_dir=tmp_path, batch_size=2 + ) + assert client.submit_batch.call_count == 3 + chunk_ids = [ + [r["custom_id"] for r in call.args[0]] + for call in client.submit_batch.call_args_list + ] + assert chunk_ids == [["p0", "p1"], ["p2", "p3"], ["p4"]] + + def test_receipt_carries_stub_fields(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, batch_size=2 + ) + (s,) = subs + assert s.batch_id == "batch_abc" + assert s.num_requests == 2 + assert s.provider == "openai" + assert s.model == "gpt-5.1" + assert s.error is None + assert isinstance(s.submitted_at, datetime) + assert s.submitted_at.tzinfo == timezone.utc + assert s.id_map == {"p0": "p0", "p1": "p1"} + + def test_display_name_routed_through_submit_metadata(self, tmp_path): + client = _openai_client() + client.submit_batch = MagicMock(return_value="b0") + subs = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="my-run" + ) + _, kwargs = client.submit_batch.call_args + assert kwargs["metadata"]["display_name"] == "my-run" + assert subs[0].metadata["display_name"] == "my-run" + + +class TestValidation: + def test_duplicate_problem_id_raises_before_submit(self, tmp_path): + client = _openai_client() + client.submit_batch = MagicMock() + problems = [ + PhysicsProblem(problem_id="dup", question="A"), + PhysicsProblem(problem_id="dup", question="B"), + ] + with pytest.raises(BatchInputError, match="Duplicate"): + submit_batch_physics_reasoning(client, problems, output_dir=tmp_path) + client.submit_batch.assert_not_called() + + def test_over_64_char_id_openai_raises(self, tmp_path): + client = _openai_client() + long_id = "x" * 65 + problems = [PhysicsProblem(problem_id=long_id, question="A")] + with pytest.raises(BatchInputError, match="64-char"): + submit_batch_physics_reasoning(client, problems, output_dir=tmp_path) + + def test_illegal_charset_id_anthropic_raises(self, tmp_path): + client = _anthropic_client() + problems = [PhysicsProblem(problem_id="has space", question="A")] + with pytest.raises(BatchInputError, match="Anthropic"): + submit_batch_physics_reasoning(client, problems, output_dir=tmp_path) + + def test_empty_input_raises(self, tmp_path): + client = _openai_client() + with pytest.raises(BatchInputError, match="empty"): + submit_batch_physics_reasoning(client, [], output_dir=tmp_path) + + def test_non_positive_batch_size_raises(self, tmp_path): + client = _openai_client() + with pytest.raises(BatchInputError, match="batch_size must be positive"): + submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, batch_size=0 + ) + + +class TestBatchSubmitError: + def test_carries_successes_and_failed_for_resume(self): + from prkit.batch import BatchSubmitError + + ok = BatchSubmission( + provider="openai", + model="gpt-5.1", + batch_id="batch_0", + input_file_path="/x/inputs/batch_0000.jsonl", + submitted_at=datetime.now(timezone.utc), + num_requests=1, + batch_index=0, + id_map={"p0": "p0"}, + ) + failed = BatchSubmission( + provider="openai", + model="gpt-5.1", + batch_id="", + input_file_path="/x/inputs/batch_0001.jsonl", + submitted_at=datetime.now(timezone.utc), + num_requests=1, + batch_index=1, + id_map={"p1": "p1"}, + error="RuntimeError: boom", + ) + err = BatchSubmitError("submit failed", successes=[ok], failed=failed) + assert err.successes == [ok] + assert err.failed is failed + assert isinstance(err, RuntimeError) + + +class TestSurrogateIds: + def test_custom_id_fn_surrogates_recovered_via_id_map(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, + _dataset(2), + output_dir=tmp_path, + batch_size=2, + custom_id_fn=lambda p: f"sur-{p.problem_id}", + ) + (s,) = subs + assert s.id_map == {"sur-p0": "p0", "sur-p1": "p1"} + # The wire custom_id is the surrogate. + line = json.loads( + (tmp_path / s.input_file_path.split("/")[-3] / "inputs") + .joinpath("batch_0000.jsonl") + .read_text() + .splitlines()[0] + ) + assert line["custom_id"] == "sur-p0" + + +class TestPartialFailure: + def test_middle_batch_failure_recorded_others_succeed(self, tmp_path): + client = _openai_client() + client.client.batches.create.side_effect = [ + MagicMock(id="batch_0"), + RuntimeError("boom"), + MagicMock(id="batch_2"), + ] + subs = submit_batch_physics_reasoning( + client, _dataset(5), output_dir=tmp_path, batch_size=2 + ) + assert len(subs) == 3 + assert subs[0].error is None and subs[0].batch_id == "batch_0" + assert subs[1].error is not None and "boom" in subs[1].error + assert subs[1].batch_id == "" + assert subs[2].error is None and subs[2].batch_id == "batch_2" + # Failed batch's input file still exists for resume. + for s in subs: + assert ( + tmp_path / s.input_file_path.split(str(tmp_path) + "/")[-1] + ).exists() + + +class TestFacade: + def test_client_facade_delegates_to_module(self, tmp_path): + client = _openai_client() + ds = _dataset(1) + with patch("prkit.batch.submit_batch_physics_reasoning") as mock_submit: + mock_submit.return_value = [] + client.submit_batch_physics_reasoning(ds, output_dir=tmp_path) + mock_submit.assert_called_once() + args, kwargs = mock_submit.call_args + assert args[0] is client and args[1] is ds + assert kwargs["output_dir"] == tmp_path + + +class TestAnthropicInline: + def test_artifact_written_and_inline_list_submitted(self, tmp_path): + client = _anthropic_client() + subs = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, batch_size=2 + ) + (s,) = subs + # The .jsonl artifact is still written for Anthropic (audit/resume). + artifact = ( + tmp_path / s.input_file_path.split("/")[-3] / "inputs" / "batch_0000.jsonl" + ) + assert artifact.exists() + # submit_batch received the inline list of request dicts (no file upload). + _, kwargs = client.client.messages.batches.create.call_args + inline = kwargs["requests"] + assert [r["custom_id"] for r in inline] == ["p0", "p1"] + client.client.files.create.assert_not_called() + + +class TestRunFolderLayout: + def test_layout_and_paths(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, run_name="run-x", batch_size=2 + ) + run_dir = tmp_path / "run-x" + assert (run_dir / "metadata.json").is_file() + assert (run_dir / "inputs" / "batch_0000.jsonl").is_file() + for s in subs: + assert s.input_file_path.startswith(str(run_dir / "inputs")) + assert s.metadata_path == str(run_dir / "metadata.json") + + +class TestMetadataContent: + def test_metadata_written_before_return_with_header_and_submissions(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, run_name="run-x", batch_size=2 + ) + record = json.loads((tmp_path / "run-x" / "metadata.json").read_text()) + assert record["run_id"] == "run-x" + assert record["provider"] == "openai" + assert record["model"] == "gpt-5.1" + assert record["batch_size"] == 2 + assert record["total_problems"] == 3 + assert record["num_batches"] == 2 + assert record["request_kind"] == "free_text" + assert record["prkit_api_version"] + assert record["dataset"] == {"name": "physreason", "version": "1.0"} + assert record["submissions"] == [s.to_dict() for s in subs] + + def test_bare_list_omits_dataset_field_and_segment(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _problems(2), output_dir=tmp_path, batch_size=2 + ) + record = json.loads( + ( + tmp_path / subs[0].metadata_path.split("/")[-2] / "metadata.json" + ).read_text() + ) + assert record["dataset"] is None + # run name omits the segment -> starts with the slugified model. + assert record["run_id"].startswith("gpt-5.1-") + + +class TestIdNormalizationRoundTrip: + @pytest.mark.parametrize( + "builder,expected", + [ + (_openai_client, "batch_abc"), + (_anthropic_client, "msgbatch_1"), + (_gemini_client, "batches/xyz"), + ], + ) + def test_wire_id_lands_verbatim(self, builder, expected, tmp_path): + client = builder() + subs = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, batch_size=1 + ) + assert subs[0].batch_id == expected + + def test_openai_id_accepted_by_poll_batch(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, batch_size=1 + ) + client.client.batches.retrieve.return_value = MagicMock( + status="completed", + output_file_id=None, + error_file_id=None, + request_counts=MagicMock(total=1, completed=1, failed=0), + ) + status = client.poll_batch(subs[0].batch_id) + assert status.batch_id == "batch_abc" + + +class TestDefaultsAndOverrides: + def test_default_output_dir_and_display_name(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + client = _openai_client() + subs = submit_batch_physics_reasoning(client, _dataset(1)) + run_id = subs[0].metadata_path.split("/")[-2] + assert (tmp_path / "batch_runs" / run_id / "metadata.json").is_file() + assert subs[0].metadata["display_name"] == run_id + + def test_explicit_run_name_is_slugified(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="My Eval 001!" + ) + assert (tmp_path / "my-eval-001").is_dir() + assert subs[0].metadata_path == str(tmp_path / "my-eval-001" / "metadata.json") + + +class TestOverwrite: + def test_existing_nonempty_folder_refused_by_default(self, tmp_path): + client = _openai_client() + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x" + ) + with pytest.raises(BatchInputError, match="already exists"): + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x" + ) + + def test_overwrite_true_succeeds_and_clears_stale_inputs(self, tmp_path): + client = _openai_client() + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x" + ) + stale = tmp_path / "run-x" / "inputs" / "batch_0005.jsonl" + stale.write_text("stale\n") + subs = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x", overwrite=True + ) + assert len(subs) == 1 + assert not stale.exists() + + def test_returns_list_of_batch_submission(self, tmp_path): + client = _openai_client() + subs = submit_batch_physics_reasoning(client, _dataset(1), output_dir=tmp_path) + assert all(isinstance(s, BatchSubmission) for s in subs) From 2ebdc17dbeb335f9cac5a1720a8fe31337452f31 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 11:57:12 -0400 Subject: [PATCH 03/14] Remove dead structured-batch request wrappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Delete the 100%-unused structured-output batch path: build_batch_structured_request and parse_batch_structured_response plus the _build_batch_structured_request default on BaseModelClient, and the four per-provider overrides (openai/anthropic/gemini/xai). These had zero callers anywhere — free-text batch submit mirrors the free-text-only solve_physics_problem and needs none of it; structured batch can return as a follow-up when structured output lands on the sync path too. xAI's only batch method was this override; it now has no batch surface (it is not a supported batch provider — submit/poll/retrieve already raise NotImplementedError). Prune the imports orphaned by the deletions and fix the build_batch_request docstring cross-reference. The structured-output engine (StructuredOutputPlan / .parse() / _resolve_structured_output_plan) and the retained batch lifecycle (build_batch_request / submit_batch / poll_batch / retrieve_batch_results / _parse_*_result_line) are untouched. Co-Authored-By: Claude Opus 4.8 --- src/prkit/core/model_clients/anthropic.py | 28 ------- src/prkit/core/model_clients/base.py | 82 +------------------ src/prkit/core/model_clients/gemini.py | 46 ----------- src/prkit/core/model_clients/openai.py | 31 ------- src/prkit/core/model_clients/xai.py | 54 ------------ .../core/model_clients/test_anthropic.py | 29 ------- 6 files changed, 1 insertion(+), 269 deletions(-) diff --git a/src/prkit/core/model_clients/anthropic.py b/src/prkit/core/model_clients/anthropic.py index 94e4ba2..53d5219 100644 --- a/src/prkit/core/model_clients/anthropic.py +++ b/src/prkit/core/model_clients/anthropic.py @@ -444,34 +444,6 @@ def _resolve_structured_output_plan( response_format=_anthropic_transformed_response_format(spec), ) - def _build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[BaseModel], - image_paths: tuple[str, ...], - max_output_tokens: int | None, - plan: StructuredOutputPlan, - **kwargs: Any, - ) -> dict[str, Any]: - del response_model, kwargs - if plan.mode != "json_schema": - raise ValueError( - "Anthropic batch structured requests require json_schema mode. " - f"Got {plan.mode!r}." - ) - params = self._build_messages_params( - input=user_prompt, - instructions=None, - image_paths=image_paths, - max_output_tokens=( - max_output_tokens if max_output_tokens is not None else 4096 - ), - response_format=plan.response_format or {}, - ) - return {"custom_id": request_id, "params": params} - def _parse_anthropic_result_entry(entry: Any) -> BatchResult: """Parse one streamed Message Batch result entry into a ``BatchResult``.""" diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index 3b213c5..a57f873 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -294,61 +294,6 @@ def solve_physics_problem( **kwargs, ) - def build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[T], - image_paths: Sequence[str] | None = None, - max_output_tokens: int | None = None, - structured_policy: StructuredOutputPolicy = "native_required", - **kwargs: Any, - ) -> dict[str, Any]: - """Build a provider-specific batch request dict for structured output.""" - plan = self.resolve_structured_output_plan( - response_model, - structured_policy=structured_policy, - ) - return self._build_batch_structured_request( - request_id=request_id, - user_prompt=user_prompt + (plan.prompt_suffix or ""), - response_model=response_model, - image_paths=tuple(image_paths or ()), - max_output_tokens=max_output_tokens, - plan=plan, - **kwargs, - ) - - def parse_batch_structured_response( - self, - *, - response_model: type[T], - raw_response_text: str, - structured_policy: StructuredOutputPolicy = "native_required", - structured_output_strategy: str | None = None, - ) -> StructuredCallResult[T]: - """Parse a raw batch response text into a typed ``StructuredCallResult``.""" - plan = self.resolve_structured_output_plan( - response_model, - structured_policy=structured_policy, - ) - if structured_output_strategy is not None: - plan = StructuredOutputPlan( - mode=plan.mode, - strategy=structured_output_strategy, - native_schema_enforced=plan.native_schema_enforced, - accepted_artifact_modes=plan.accepted_artifact_modes, - accepted_artifact_strategies=plan.accepted_artifact_strategies, - response_format=plan.response_format, - prompt_suffix=plan.prompt_suffix, - ) - return self._build_structured_call_result( - response_model=response_model, - raw_text=raw_response_text, - plan=plan, - ) - # ------------------------------------------------------------------ # Free-text batch requests + the asynchronous job lifecycle # ------------------------------------------------------------------ @@ -369,8 +314,7 @@ def build_batch_request( no ``response_format``. Passing ``instructions=""`` suppresses the system prompt on every provider (see :meth:`_resolve_instructions`), so the resulting request matches a synchronous - ``response(input=..., instructions="")`` call. Use - :meth:`build_batch_structured_request` instead for schema-enforced output. + ``response(input=..., instructions="")`` call. """ return self._build_batch_request( request_id=request_id, @@ -596,27 +540,3 @@ def _extract_structured_payload( ) -> dict[str, Any] | list[Any] | None: del plan return extract_json_payload(raw_text) - - def _build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[T], - image_paths: tuple[str, ...], - max_output_tokens: int | None, - plan: StructuredOutputPlan, - **kwargs: Any, - ) -> dict[str, Any]: - del ( - request_id, - user_prompt, - response_model, - image_paths, - max_output_tokens, - plan, - kwargs, - ) - raise NotImplementedError( - f"Batch structured requests are not implemented for provider={self._provider_name()!r}." - ) diff --git a/src/prkit/core/model_clients/gemini.py b/src/prkit/core/model_clients/gemini.py index 67cded0..a652d7e 100644 --- a/src/prkit/core/model_clients/gemini.py +++ b/src/prkit/core/model_clients/gemini.py @@ -10,7 +10,6 @@ import PIL.Image from google import genai from google.genai import types -from pydantic import BaseModel from .base import BaseModelClient from .batch_types import BatchItemStatus, BatchResult, BatchState, BatchStatus @@ -165,51 +164,6 @@ def _resolve_structured_output_plan( ), ) - def _build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[BaseModel], - image_paths: tuple[str, ...], - max_output_tokens: int | None, - plan: StructuredOutputPlan, - **kwargs: Any, - ) -> dict[str, Any]: - del response_model, kwargs - if plan.mode != "json_schema": - raise ValueError( - f"Gemini batch structured requests require json_schema mode. Got {plan.mode!r}." - ) - - parts: list[dict[str, Any]] = [] - for image_path in image_paths: - parts.append( - { - "inline_data": { - "mime_type": _guess_mime_type(image_path), - "data": _encode_image_file_base64(image_path), - } - } - ) - parts.append({"text": user_prompt}) - - normalized = normalize_response_format(plan.response_format or {}) - generation_config: dict[str, Any] = { - "response_mime_type": "application/json", - "response_json_schema": extract_schema_for_gemini(normalized), - } - if max_output_tokens is not None: - generation_config["max_output_tokens"] = max_output_tokens - - return { - "key": request_id, - "request": { - "contents": [{"parts": parts, "role": "user"}], - "generation_config": generation_config, - }, - } - def _build_batch_request( self, *, diff --git a/src/prkit/core/model_clients/openai.py b/src/prkit/core/model_clients/openai.py index 6c96c74..ab6aa9e 100644 --- a/src/prkit/core/model_clients/openai.py +++ b/src/prkit/core/model_clients/openai.py @@ -17,7 +17,6 @@ from typing import Any from openai import OpenAI -from pydantic import BaseModel from ..project_env import ensure_openai_api_key from .base import BaseModelClient @@ -497,36 +496,6 @@ def _resolve_structured_output_plan( ), ) - def _build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[BaseModel], - image_paths: tuple[str, ...], - max_output_tokens: int | None, - plan: StructuredOutputPlan, - **kwargs: Any, - ) -> dict[str, Any]: - del response_model, kwargs - if plan.mode != "json_schema": - raise ValueError( - f"OpenAI batch structured requests require json_schema mode. Got {plan.mode!r}." - ) - body = self._build_responses_body( - input=user_prompt, - instructions=None, - image_paths=image_paths, - max_output_tokens=max_output_tokens, - response_format=plan.response_format or {}, - ) - return { - "custom_id": request_id, - "method": "POST", - "url": "/v1/responses", - "body": body, - } - def _iter_jsonl_lines(content: Any) -> Iterator[str]: """Yield non-empty lines from an OpenAI ``files.content`` payload. diff --git a/src/prkit/core/model_clients/xai.py b/src/prkit/core/model_clients/xai.py index 7c15709..65b88db 100644 --- a/src/prkit/core/model_clients/xai.py +++ b/src/prkit/core/model_clients/xai.py @@ -6,9 +6,6 @@ from typing import Any -from pydantic import BaseModel - -from .openai import prepare_image_url_from_image_path from .openai_compatible_chat import OpenAICompatibleChatModel from .structured_output import ( StructuredOutputPlan, @@ -95,54 +92,3 @@ def _resolve_structured_output_plan( accepted_artifact_strategies=("xai_chat_json_schema",), response_format=_xai_transformed_response_format(spec), ) - - def _build_batch_structured_request( - self, - *, - request_id: str, - user_prompt: str, - response_model: type[BaseModel], - image_paths: tuple[str, ...], - max_output_tokens: int | None, - plan: StructuredOutputPlan, - **kwargs: Any, - ) -> dict[str, Any]: - del response_model, kwargs - if plan.mode != "json_schema": - raise ValueError( - f"xAI batch structured requests require json_schema mode. Got {plan.mode!r}." - ) - - normalized = normalize_response_format(plan.response_format or {}) - text_format: dict[str, Any] = { - "type": "json_schema", - "name": normalized["name"], - "schema": normalized["schema"], - "strict": normalized.get("strict", True), - } - if normalized.get("description") is not None: - text_format["description"] = normalized["description"] - - content: list[dict[str, Any]] = [{"type": "input_text", "text": user_prompt}] - for image_path in image_paths: - content.append( - { - "type": "input_image", - "image_url": prepare_image_url_from_image_path(image_path), - } - ) - - body: dict[str, Any] = { - "model": self.model, - "input": [{"role": "user", "content": content}], - "text": {"format": text_format}, - } - if max_output_tokens is not None: - body["max_output_tokens"] = max_output_tokens - - return { - "custom_id": request_id, - "method": "POST", - "url": "/v1/responses", - "body": body, - } diff --git a/tests/prkit/core/model_clients/test_anthropic.py b/tests/prkit/core/model_clients/test_anthropic.py index a589fa5..42a04be 100644 --- a/tests/prkit/core/model_clients/test_anthropic.py +++ b/tests/prkit/core/model_clients/test_anthropic.py @@ -14,7 +14,6 @@ _parse_data_url, ) from prkit.core.model_clients.base import DEFAULT_INSTRUCTIONS -from prkit.core.model_clients.structured_output import coerce_structured_output_spec ANTHROPIC_TEST_MODEL = "claude-sonnet-4-6" @@ -171,34 +170,6 @@ class ExampleResponse(BaseModel): assert call_kwargs["output_config"]["format"]["schema"]["type"] == "object" assert "tools" not in call_kwargs - @patch("prkit.core.model_clients.anthropic.Anthropic") - def test_batch_structured_request_omits_name(self, mock_anthropic_class): - """Batch structured requests must also omit Anthropic's forbidden `name` key.""" - mock_client = MagicMock() - mock_anthropic_class.return_value = mock_client - - class ExampleResponse(BaseModel): - answer: str - - client = AnthropicModel(ANTHROPIC_TEST_MODEL) - spec = coerce_structured_output_spec(ExampleResponse) - plan = client._resolve_structured_output_plan( - spec, structured_policy="native_required" - ) - request = client._build_batch_structured_request( - request_id="r1", - user_prompt="hi", - response_model=ExampleResponse, - image_paths=(), - max_output_tokens=None, - plan=plan, - ) - - fmt = request["params"]["output_config"]["format"] - assert fmt["type"] == "json_schema" - assert "name" not in fmt # Anthropic forbids it (400 otherwise) - assert fmt["schema"]["type"] == "object" - def test_extract_tool_use_json_requires_exactly_one_tool_block(self): with pytest.raises(ValueError, match="exactly one tool_use block"): _extract_tool_use_json([{"type": "text", "text": "not structured"}]) From 7a85fa0eb15c5de17576853c74dfab0bcf06dfd2 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 12:01:29 -0400 Subject: [PATCH 04/14] Remove reserved Runner Protocol from the public contract Delete the Runner structural Protocol from prkit.api (and prkit.api.__all__), the docstring mention in prkit/__init__.py, the reserved row in CONTRACT.md, and its references in tests/prkit/test_api.py. Runner had no implementation; its docstring reserved it "for roadmap N4", which is now batch-mode submit shipped as a bounded submitter (prkit.batch.submit_batch_physics_reasoning + the BaseModelClient facade), not a Runner noun. The contract now pins three Protocol nouns (DatasetProvider / ModelClient / Scorer) plus Verdict. The unrelated annotation Runner (a Callable alias) and the llm_judge OpenAIJudgeRunner are different symbols and untouched. Per the API_VERSION policy this provisional-1.0 removal is recorded in CONTRACT.md, not signalled by a bump. Co-Authored-By: Claude Opus 4.8 --- src/prkit/CONTRACT.md | 10 ++++++++-- src/prkit/__init__.py | 2 +- src/prkit/api.py | 19 ------------------- tests/prkit/test_api.py | 3 --- 4 files changed, 9 insertions(+), 25 deletions(-) diff --git a/src/prkit/CONTRACT.md b/src/prkit/CONTRACT.md index da665fe..1523dfb 100644 --- a/src/prkit/CONTRACT.md +++ b/src/prkit/CONTRACT.md @@ -32,14 +32,13 @@ and returns the same canonical `Verdict`. `check_model_client`, `ConformanceTestMixin`) is a stable companion: use it to verify your own loader/scorer/client satisfies the contract. -The contract pins four structural (`typing.Protocol`) nouns plus one result type: +The contract pins three structural (`typing.Protocol`) nouns plus one result type: | Noun | Protocol | Reference implementation | |------|----------|--------------------------| | Dataset loader | `DatasetProvider` | `BaseDatasetLoader` subclasses | | Inference client | `ModelClient` | `BaseModelClient` subclasses | | Scorer | `Scorer` | `prkit.scoring.SemanticsScorer` (binary); `EedScorer`/`SeedScorer` (vendor edit-distance baselines); `SemanticsEedScorer`/`SemanticsSeedScorer` (our-semantics edit distance, graded); `LLMJudgeScorer` (model-graded) | -| Runner | `Runner` | *(reserved; no implementation yet)* | | Result | `Verdict` | `prkit.core.verdict.Verdict` | > **Note on `@runtime_checkable`:** `isinstance(x, Scorer)` only verifies that the @@ -135,3 +134,10 @@ changes to those are documented in the package release notes, not the contract v compatibility shim is deleted; import from `prkit.semantics.build` directly. - `prkit.evaluation.llm_judge` (model-graded scoring) is **retained** — it is a distinct capability, not a duplicate of the deterministic scoring path. +- **`Runner` Protocol removed.** The reserved `Runner` structural Protocol (it had no + implementation; its docstring reserved it "for roadmap N4") is removed from + `prkit.api` / `prkit.api.__all__`. Batch-mode runs (N4) ship as a bounded *submitter* + (`prkit.batch.submit_batch_physics_reasoning` + the `BaseModelClient` facade), not as a + `Runner` noun. The contract now pins **three** Protocol nouns plus `Verdict`. Per the + `API_VERSION` policy this provisional-1.0 removal is tracked here, **not** signalled by a + major bump. diff --git a/src/prkit/__init__.py b/src/prkit/__init__.py index c50d0f0..487b7e0 100644 --- a/src/prkit/__init__.py +++ b/src/prkit/__init__.py @@ -9,7 +9,7 @@ Integration entry point: :mod:`prkit.api` is the frozen, semver-governed public contract (the - ``DatasetProvider``/``ModelClient``/``Scorer``/``Runner`` protocols and the + ``DatasetProvider``/``ModelClient``/``Scorer`` protocols and the ``Verdict`` result). Integrate against it rather than reaching into subpackages. See ``prkit/CONTRACT.md`` for the stability and version policy. diff --git a/src/prkit/api.py b/src/prkit/api.py index 3a58206..a4c4732 100644 --- a/src/prkit/api.py +++ b/src/prkit/api.py @@ -101,31 +101,12 @@ def score( def get_info(self) -> dict[str, Any]: ... # MUST include "version" -@runtime_checkable -class Runner(Protocol): - """Orchestration noun: drive a :class:`ModelClient` over a - :class:`PhysicsDataset` and score with a :class:`Scorer`. - - No implementation ships today; the contract is reserved for a later - orchestration item (roadmap N4). - """ - - def run( - self, - dataset: PhysicsDataset, - model: ModelClient, - scorer: Scorer, - **kwargs: Any, - ) -> Any: ... - - __all__ = [ "API_VERSION", # the four structural nouns + the canonical result "DatasetProvider", "ModelClient", "Scorer", - "Runner", "Verdict", # canonical answer ontology "AnswerObjectKind", diff --git a/tests/prkit/test_api.py b/tests/prkit/test_api.py index 4ca92f6..c6b6cfc 100644 --- a/tests/prkit/test_api.py +++ b/tests/prkit/test_api.py @@ -6,7 +6,6 @@ from prkit.api import ( DatasetProvider, ModelClient, - Runner, Scorer, Verdict, create_model_client, @@ -31,7 +30,6 @@ def test_all_is_frozen_surface(self): "DatasetProvider", "ModelClient", "Scorer", - "Runner", "Verdict", "AnswerObjectKind", "AnswerStructure", @@ -124,4 +122,3 @@ def test_protocols_reject_unrelated_object(self): assert not isinstance(object(), DatasetProvider) assert not isinstance(object(), ModelClient) assert not isinstance(object(), Scorer) - assert not isinstance(object(), Runner) From d5dab67ed0a0669d308666972b1e65f89f88c99f Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 15:29:38 -0400 Subject: [PATCH 05/14] Reshape batch submit into a mutable whole-batch ledger Replace the frozen per-minibatch BatchSubmission receipt with one mutable whole-batch ledger per run, in preparation for a resumable fetch step that advances it. Batch-level facts (provider, model, created_at, dataset, ...) are stored once; each provider job becomes a dict in `minibatches` carrying a mutable `status` plus its batch_id / id_map / num_requests / counts / output_path. - BatchSubmission gains `from_dict`, `save`, `load`, and the pure status helpers `minibatches_to_fetch` / `set_status` / `is_complete` / `status_counts`; it is now a dataclass (was frozen). `from_dict` parses `created_at` with `datetime.fromisoformat` and compares tz by equality (not identity). - `submit_batch_physics_reasoning` now saves the ledger to `/metadata.json` and returns the run-folder path (str) instead of `list[BatchSubmission]`; disk is the source of truth across the ~24h provider window. The `BaseModelClient` facade return type follows. - Rename batch->minibatch throughout the unit sense: `batch_size`->`minibatch_size`, `num_batches`->`minibatch_count`, `batch_index`->minibatch `index`, `batch_*.jsonl`->`minibatch_*.jsonl`. Add the status-string vocabulary (submitted/running/.../fetch_error) the ledger records. - Remove `BatchSubmitError`: it carried per-minibatch BatchSubmission receipts for resume, a shape the single-ledger reshape obsoletes (a failed submit is now a `submit_error` minibatch in the ledger). It was never raised. Stage-1 submit tests rewritten to load the ledger and assert on its fields; new test_batch_submission_ledger.py covers the status helpers + save/load round-trip. Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 368 ++++++++++++------ src/prkit/core/model_clients/base.py | 9 +- .../batch/test_batch_submission_ledger.py | 147 +++++++ .../batch/test_submit_physics_reasoning.py | 252 ++++++------ 4 files changed, 522 insertions(+), 254 deletions(-) create mode 100644 tests/prkit/batch/test_batch_submission_ledger.py diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py index 0865c24..dec6b00 100644 --- a/src/prkit/batch/__init__.py +++ b/src/prkit/batch/__init__.py @@ -1,16 +1,25 @@ -"""Batch-mode *submit* for physics-reasoning runs (N4 Stage 1). +"""Batch-mode *submit* for physics-reasoning runs (N4). This leaf turns a set of :class:`~prkit.core.domain.PhysicsProblem`\\s into -submitted provider batch jobs and returns one typed receipt -(:class:`BatchSubmission`) per batch. It preprocesses each problem **exactly** -like the synchronous ``solve_physics_problem`` path (via the client's -``build_problem_batch_request``), splits the dataset into provider batches of -``batch_size`` problems, writes each batch's requests as provider-correct JSONL -under one run folder per call, submits each batch, and writes a run-level -``metadata.json`` audit record **before** returning. - -It is a *bounded submitter*: it stops at submitted receipts. Fetching / -polling / scoring / pricing is a later milestone that consumes these receipts. +submitted provider batch jobs and one whole-batch :class:`BatchSubmission` ledger. + +**Vocabulary** (owner-set): a **batch** is the whole thing a user triggers over a +dataset (one :func:`submit_batch_physics_reasoning` call → one *run folder* → one +:class:`BatchSubmission` ledger). A **minibatch** is one ``minibatch_size``-problem +group = one provider batch job = one ``minibatch_XXXX.jsonl`` = one element of +:attr:`BatchSubmission.minibatches`. ``prkit.batch`` / ``submit_batch_*`` keep +"batch" because it names the *batch lane*, not a unit. + +:func:`submit_batch_physics_reasoning` preprocesses each problem exactly like the +synchronous ``solve_physics_problem`` path (via the client's +``build_problem_batch_request``), splits the dataset into minibatches of +``minibatch_size`` problems, writes each minibatch's requests as provider-correct +JSONL under one run folder, submits each minibatch, saves the consolidated +``metadata.json`` ledger, and **returns the run-folder path** (a ``str``). The +ledger is mutable on purpose — it is the resume state store (each minibatch carries +a ``status``); disk is the source of truth across the ~24h provider window, so the +object is reconstructed on demand via :meth:`BatchSubmission.load`. Polling / +downloading / correlating the results is a later milestone that consumes the ledger. Import discipline: at module load this imports only :mod:`prkit.core.domain` and the standard library — never ``prkit.api``, the dataset hub, a scorer, the @@ -32,13 +41,46 @@ __all__ = [ "BatchSubmission", "BatchInputError", - "BatchSubmitError", "submit_batch_physics_reasoning", "dumps_batch_jsonl", "write_batch_jsonl", "validate_batch_requests", + # minibatch status constants (plain strings; kept off the BatchState enum to + # keep the leaf light) + "SUBMITTED", + "RUNNING", + "COMPLETED", + "EXPIRED", + "FAILED", + "CANCELLED", + "FETCHED", + "SUBMIT_ERROR", + "FETCH_ERROR", ] +# --------------------------------------------------------------------------- # +# Minibatch status constants # +# --------------------------------------------------------------------------- # +SUBMITTED = "submitted" # accepted by the provider, not yet polled to a result +RUNNING = "running" # provider is processing (pending / in-progress) +COMPLETED = "completed" # provider finished; download not yet persisted +EXPIRED = "expired" # window elapsed and nothing was retrievable +FAILED = "failed" # batch-level provider failure +CANCELLED = "cancelled" # cancelled at the provider +FETCHED = "fetched" # results downloaded + persisted to outputs/ (terminal-good) +SUBMIT_ERROR = "submit_error" # never submitted (batch_id == ""); re-submit via Stage 1 +FETCH_ERROR = "fetch_error" # retrieve raised; non-final, retried on the next pass + +# Minibatches in these statuses are done with the fetch loop and skipped on the +# next pass (idempotent resume). EXPIRED is deliberately NOT here: an expired job +# is re-polled (cheap snapshot) until its results window truly closes — matches +# the design's skip set exactly. +_SKIP_FETCH_STATUSES = frozenset({FETCHED, SUBMIT_ERROR, FAILED, CANCELLED}) + +# Minibatches in these statuses are terminal for ``is_complete()`` — fetched, or +# terminal-failed with nothing left to retrieve. +_COMPLETE_STATUSES = frozenset({FETCHED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED}) + # Providers whose batch line correlates by ``key`` rather than ``custom_id``. _KEY_ID_PROVIDERS = frozenset({"google", "gemini"}) # Anthropic restricts custom ids to this charset/length. @@ -46,75 +88,163 @@ _MAX_ID_LEN = 64 -@dataclass(frozen=True) +@dataclass class BatchSubmission: - """One submission receipt = one batch = one provider batch job. - - The normalized provider id is in :attr:`batch_id` (empty string when submit - failed — the batch is then resumable via :attr:`input_file_path`, whose JSONL - is always written). :attr:`id_map` maps each wire ``custom_id`` / ``key`` back - to its ``problem_id`` so a later fetch step correlates results to problems - even when surrogate ids are used. + """The whole-batch ledger = one run folder = one ``metadata.json``. + + Mutable on purpose: the fetch step advances each minibatch's ``status`` and + persists the ledger, which is what makes fetch idempotent/resumable across + processes. Batch-level facts are stored once; the per-minibatch fetch state + lives in :attr:`minibatches`, a list of plain dicts (one per provider job). + + Each ``minibatches[i]`` dict carries:: + + { + "index": int, # minibatch ordinal (0-based) + "batch_id": str, # provider job id ("" => submit failed) + "status": str, # one of the module status constants + "input_file_path": str, # /inputs/minibatch_XXXX.jsonl + "num_requests": int, # problems in this minibatch + "id_map": dict[str, str], # wire custom_id / key -> problem_id + "submitted_at": str | None, # ISO 8601 + "error": str | None, # local submit error (with SUBMIT_ERROR) + "output_path": str | None, # /outputs/... (set on fetch) + "fetched_at": str | None, # ISO 8601 (set on fetch) + "counts": dict[str, int], # last poll's request_counts + "endpoint": str | None, # audit-only (provider parity) + "completion_window": str | None, # audit-only (provider parity) + } """ - provider: str + # ---- batch-level (stored once; never repeated per minibatch) ---- + run_name: str + provider: str # "openai" | "anthropic" | "google" model: str - batch_id: str - input_file_path: str - submitted_at: datetime - num_requests: int - batch_index: int - id_map: dict[str, str] - endpoint: str | None = None - completion_window: str | None = None - metadata_path: str | None = None + created_at: datetime # UTC + minibatch_size: int # record cap per minibatch + minibatch_count: int # number of minibatches + total_problems: int + dataset: dict[str, Any] | None # {"name", "version"} or None for a bare list + request_kind: str # "free_text" + prkit_api_version: str + display_name: str + run_dir: str # "/" metadata: dict[str, str] = field(default_factory=dict) - error: str | None = None - + # ---- per-minibatch ledger (the mutable, fetch-updated part) ---- + minibatches: list[dict[str, Any]] = field(default_factory=list) + + # ---- pure state helpers (NO network) ---- + def minibatches_to_fetch(self) -> list[dict[str, Any]]: + """Minibatches still in the fetch loop (status not in the skip set).""" + return [ + mb for mb in self.minibatches if mb["status"] not in _SKIP_FETCH_STATUSES + ] + + def set_status(self, index: int, status: str, **fields: Any) -> None: + """Set ``minibatches[index]['status']`` (and any extra fields) in place. + + Looks up by the minibatch's own ``index`` field (not list position, though + they coincide for ledgers built by :func:`submit_batch_physics_reasoning`). + """ + for mb in self.minibatches: + if mb["index"] == index: + mb["status"] = status + mb.update(fields) + return + raise KeyError(f"No minibatch with index {index}.") + + def is_complete(self) -> bool: + """True once every minibatch is fetched or terminal-failed.""" + return all(mb["status"] in _COMPLETE_STATUSES for mb in self.minibatches) + + def status_counts(self) -> dict[str, int]: + """Return ``{status: count}`` over the minibatches (backs the summary).""" + counts: dict[str, int] = {} + for mb in self.minibatches: + counts[mb["status"]] = counts.get(mb["status"], 0) + 1 + return counts + + # ---- serialization ---- def to_dict(self) -> dict[str, Any]: - """Return a JSON-serializable dict (``submitted_at`` rendered ISO 8601).""" + """Return a JSON-serializable dict (``created_at`` rendered ISO 8601). + + Per-minibatch timestamps are already stored as ISO strings, so only the + batch-level ``created_at`` needs rendering. + """ return { + "run_name": self.run_name, "provider": self.provider, "model": self.model, - "batch_id": self.batch_id, - "input_file_path": self.input_file_path, - "submitted_at": self.submitted_at.isoformat(), - "num_requests": self.num_requests, - "batch_index": self.batch_index, - "id_map": dict(self.id_map), - "endpoint": self.endpoint, - "completion_window": self.completion_window, - "metadata_path": self.metadata_path, + "created_at": self.created_at.isoformat(), + "minibatch_size": self.minibatch_size, + "minibatch_count": self.minibatch_count, + "total_problems": self.total_problems, + "dataset": self.dataset, + "request_kind": self.request_kind, + "prkit_api_version": self.prkit_api_version, + "display_name": self.display_name, + "run_dir": self.run_dir, "metadata": dict(self.metadata), - "error": self.error, + "minibatches": [dict(mb) for mb in self.minibatches], } + @classmethod + def from_dict(cls, data: dict[str, Any]) -> BatchSubmission: + """Reconstruct a ledger from :meth:`to_dict` output. + + ``created_at`` is parsed with :meth:`datetime.fromisoformat`; the resulting + tzinfo compares **equal** to ``timezone.utc`` (``+00:00`` round-trips by + equality, not identity), which is the contract Stage-1 tests assert. + """ + created_raw = data["created_at"] + created_at = ( + datetime.fromisoformat(created_raw) + if isinstance(created_raw, str) + else created_raw + ) + return cls( + run_name=data["run_name"], + provider=data["provider"], + model=data["model"], + created_at=created_at, + minibatch_size=data["minibatch_size"], + minibatch_count=data["minibatch_count"], + total_problems=data["total_problems"], + dataset=data.get("dataset"), + request_kind=data.get("request_kind", "free_text"), + prkit_api_version=data.get("prkit_api_version", ""), + display_name=data.get("display_name", data["run_name"]), + run_dir=data["run_dir"], + metadata=dict(data.get("metadata") or {}), + minibatches=[dict(mb) for mb in data.get("minibatches", [])], + ) -class BatchInputError(ValueError): - """Invalid submit input: empty problem set, duplicate/illegal id, or a - pre-existing non-empty run folder (when ``overwrite`` is False).""" - + def save(self, run_dir: str | Path | None = None) -> Path: + """Write ``/metadata.json`` and return its path. + + Defaults to :attr:`run_dir`. Creates the folder if missing so a fresh + ledger can be persisted before any inputs are laid down. + """ + target = Path(run_dir) if run_dir is not None else Path(self.run_dir) + target.mkdir(parents=True, exist_ok=True) + path = target / "metadata.json" + path.write_text( + json.dumps(self.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8" + ) + return path -class BatchSubmitError(RuntimeError): - """A batch submission failed. + @classmethod + def load(cls, run_dir_or_metadata_path: str | Path) -> BatchSubmission: + """Read a ledger back from a run folder or a ``metadata.json`` path.""" + p = Path(run_dir_or_metadata_path) + metadata_path = p if p.name == "metadata.json" else p / "metadata.json" + data = json.loads(metadata_path.read_text(encoding="utf-8")) + return cls.from_dict(data) - Carries the receipts that succeeded before the failure plus the failed - receipt so a consumer can resume by re-submitting the failed batch (its - input JSONL already exists). The Stage-1 submit loop records per-batch - failures on the receipts and returns the full list rather than raising; this - type is the resume-carrying error reserved for that tooling. - """ - def __init__( - self, - message: str, - *, - successes: list[BatchSubmission], - failed: BatchSubmission, - ) -> None: - super().__init__(message) - self.successes = successes - self.failed = failed +class BatchInputError(ValueError): + """Invalid submit input: empty problem set, duplicate/illegal id, or a + pre-existing non-empty run folder (when ``overwrite`` is False).""" def _slug(text: str) -> str: @@ -190,7 +320,7 @@ def validate_batch_requests( def _chunked(items: list[Any], size: int) -> list[list[Any]]: """Split *items* into consecutive chunks of at most *size*.""" if size <= 0: - raise BatchInputError(f"batch_size must be positive; got {size}.") + raise BatchInputError(f"minibatch_size must be positive; got {size}.") return [items[i : i + size] for i in range(0, len(items), size)] @@ -210,7 +340,7 @@ def submit_batch_physics_reasoning( *, output_dir: str | Path = "batch_runs", run_name: str | None = None, - batch_size: int = 500, + minibatch_size: int = 500, instructions: str | None = None, max_output_tokens: int | None = None, temperature: float | None = None, @@ -218,24 +348,28 @@ def submit_batch_physics_reasoning( display_name: str | None = None, metadata: dict[str, str] | None = None, overwrite: bool = False, -) -> list[BatchSubmission]: +) -> str: """Preprocess, split, write, and submit *problems* as provider batch jobs. Builds one free-text request per problem (mirroring the synchronous ``solve_physics_problem`` preprocessing via ``client.build_problem_batch_request``), - splits them into batches of ``batch_size``, writes each batch's provider-correct - JSONL under ``//inputs/batch_NNNN.jsonl``, submits each - batch sequentially via ``client.submit_batch``, writes a run-level - ``metadata.json`` (run header + submissions) **before** returning, and returns - one :class:`BatchSubmission` per batch. + splits them into minibatches of ``minibatch_size``, writes each minibatch's + provider-correct JSONL under ``//inputs/minibatch_NNNN.jsonl``, + submits each minibatch sequentially via ``client.submit_batch``, saves the + consolidated :class:`BatchSubmission` ledger to ``/metadata.json``, and + **returns the run-folder path** (a ``str``). + + Disk is the source of truth across the ~24h provider window, so the ledger is + persisted rather than handed back as an object: reconstruct it on demand at + fetch time via :meth:`BatchSubmission.load`. - A batch whose submit raises is recorded with ``error`` set and ``batch_id=""`` - (its input file remains for resume); the returned list always has one entry - per batch. + A minibatch whose submit raises is recorded with ``status=SUBMIT_ERROR``, + ``error`` set, and ``batch_id=""`` (its input file remains for re-submission via + Stage 1); the ledger always has one ``minibatches`` entry per minibatch. Raises: BatchInputError: empty input, an invalid/duplicate id, a non-positive - ``batch_size``, or a pre-existing non-empty run folder when + ``minibatch_size``, or a pre-existing non-empty run folder when ``overwrite`` is False. """ if isinstance(problems, PhysicsDataset): @@ -270,7 +404,9 @@ def submit_batch_physics_reasoning( f"Run folder {str(run_dir)!r} already exists and is non-empty; " "pass overwrite=True to write into it." ) - for stale in (run_dir / "inputs").glob("batch_*.jsonl"): + for stale in (run_dir / "inputs").glob("minibatch_*.jsonl"): + stale.unlink() + for stale in (run_dir / "outputs").glob("minibatch_*.jsonl"): stale.unlink() inputs_dir = run_dir / "inputs" inputs_dir.mkdir(parents=True, exist_ok=True) @@ -296,21 +432,20 @@ def submit_batch_physics_reasoning( provider = client.provider validate_batch_requests(requests, provider=provider) - request_chunks = _chunked(requests, batch_size) - rid_chunks = _chunked(request_ids, batch_size) - pid_chunks = _chunked(problem_ids, batch_size) + request_chunks = _chunked(requests, minibatch_size) + rid_chunks = _chunked(request_ids, minibatch_size) + pid_chunks = _chunked(problem_ids, minibatch_size) # Write all JSONL artifacts first (cheap, local), then submit sequentially. file_paths: list[Path] = [] for index, chunk in enumerate(request_chunks): file_paths.append( - write_batch_jsonl(chunk, inputs_dir / f"batch_{index:04d}.jsonl") + write_batch_jsonl(chunk, inputs_dir / f"minibatch_{index:04d}.jsonl") ) - metadata_path = run_dir / "metadata.json" merged_metadata = {**(metadata or {}), "display_name": display_name} - submissions: list[BatchSubmission] = [] + minibatches: list[dict[str, Any]] = [] for index, chunk in enumerate(request_chunks): id_map = dict(zip(rid_chunks[index], pid_chunks[index])) submitted_at = datetime.now(timezone.utc) @@ -318,22 +453,24 @@ def submit_batch_physics_reasoning( error: str | None = None try: batch_id = client.submit_batch(chunk, metadata=merged_metadata) - except Exception as exc: # noqa: BLE001 - recorded on the receipt for resume + except Exception as exc: # noqa: BLE001 - recorded on the ledger for resume error = f"{type(exc).__name__}: {exc}" - submissions.append( - BatchSubmission( - provider=provider, - model=client.model, - batch_id=batch_id, - input_file_path=str(file_paths[index]), - submitted_at=submitted_at, - num_requests=len(chunk), - batch_index=index, - id_map=id_map, - metadata_path=str(metadata_path), - metadata=merged_metadata, - error=error, - ) + minibatches.append( + { + "index": index, + "batch_id": batch_id, + "status": SUBMITTED if batch_id else SUBMIT_ERROR, + "input_file_path": str(file_paths[index]), + "num_requests": len(chunk), + "id_map": id_map, + "submitted_at": submitted_at.isoformat(), + "error": error, + "output_path": None, + "fetched_at": None, + "counts": {}, + "endpoint": None, + "completion_window": None, + } ) dataset_field: dict[str, Any] | None @@ -342,21 +479,22 @@ def submit_batch_physics_reasoning( else: dataset_field = None - run_record = { - "run_id": run_name, - "created_at": now.isoformat(), - "provider": provider, - "model": client.model, - "dataset": dataset_field, - "batch_size": batch_size, - "total_problems": len(problem_list), - "num_batches": len(request_chunks), - "request_kind": "free_text", - "prkit_api_version": _prkit_api_version(), - "submissions": [submission.to_dict() for submission in submissions], - } - metadata_path.write_text( - json.dumps(run_record, indent=2, ensure_ascii=False), encoding="utf-8" + submission = BatchSubmission( + run_name=run_name, + provider=provider, + model=client.model, + created_at=now, + minibatch_size=minibatch_size, + minibatch_count=len(request_chunks), + total_problems=len(problem_list), + dataset=dataset_field, + request_kind="free_text", + prkit_api_version=_prkit_api_version(), + display_name=display_name, + run_dir=str(run_dir), + metadata=dict(metadata or {}), + minibatches=minibatches, ) + submission.save() - return submissions + return str(run_dir) diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index a57f873..8727de9 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -28,7 +28,6 @@ ) if TYPE_CHECKING: - from prkit.batch import BatchSubmission from prkit.semantics.schema import PhysicsQuestionSemantics T = TypeVar("T", bound=BaseModel) @@ -380,14 +379,16 @@ def submit_batch_physics_reasoning( self, problems: PhysicsDataset | Sequence[PhysicsProblem], **kwargs: Any, - ) -> list[BatchSubmission]: - """Submit *problems* as provider batch jobs; return one receipt per batch. + ) -> str: + """Submit *problems* as provider batch jobs; return the run-folder path. Thin one-line facade mirroring :meth:`solve_physics_problem`: lazily imports :mod:`prkit.batch` (kept off the import-light path, like the lazy ``prkit.semantics`` import in :meth:`solve_physics_problem`) and delegates to :func:`prkit.batch.submit_batch_physics_reasoning`. See that function for the - keyword arguments (``output_dir`` / ``run_name`` / ``batch_size`` / …). + keyword arguments (``output_dir`` / ``run_name`` / ``minibatch_size`` / …). + The ledger is saved to ``/metadata.json``; reconstruct it later via + :meth:`fetch_batch_physics_reasoning` (or ``BatchSubmission.load(run_dir)``). """ from prkit.batch import submit_batch_physics_reasoning as _submit_batch diff --git a/tests/prkit/batch/test_batch_submission_ledger.py b/tests/prkit/batch/test_batch_submission_ledger.py new file mode 100644 index 0000000..c7572f1 --- /dev/null +++ b/tests/prkit/batch/test_batch_submission_ledger.py @@ -0,0 +1,147 @@ +"""Tests for the ``BatchSubmission`` whole-batch ledger (state + serialization). + +These exercise the mutable ledger in isolation — no client, no network: the pure +status helpers (``minibatches_to_fetch`` / ``set_status`` / ``is_complete`` / +``status_counts``) and the ``to_dict`` / ``from_dict`` / ``save`` / ``load`` +round-trip, including the timezone-by-equality contract for ``created_at``. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest + +from prkit.batch import ( + CANCELLED, + COMPLETED, + EXPIRED, + FAILED, + FETCH_ERROR, + FETCHED, + RUNNING, + SUBMIT_ERROR, + SUBMITTED, + BatchSubmission, +) + + +def _minibatch(index: int, status: str) -> dict: + return { + "index": index, + "batch_id": f"b{index}", + "status": status, + "input_file_path": f"/tmp/run/inputs/minibatch_{index:04d}.jsonl", + "num_requests": 2, + "id_map": {f"p{index}a": f"p{index}a", f"p{index}b": f"p{index}b"}, + "submitted_at": "2026-06-22T00:00:00+00:00", + "error": None, + "output_path": None, + "fetched_at": None, + "counts": {}, + "endpoint": None, + "completion_window": None, + } + + +def _ledger(statuses: list[str], *, run_dir: str = "/tmp/run") -> BatchSubmission: + minibatches = [_minibatch(i, status) for i, status in enumerate(statuses)] + return BatchSubmission( + run_name="run", + provider="openai", + model="gpt-5.1", + created_at=datetime(2026, 6, 22, 1, 2, 3, 456789, tzinfo=timezone.utc), + minibatch_size=2, + minibatch_count=len(statuses), + total_problems=2 * len(statuses), + dataset={"name": "d", "version": "1.0"}, + request_kind="free_text", + prkit_api_version="1.0", + display_name="run", + run_dir=run_dir, + metadata={"k": "v"}, + minibatches=minibatches, + ) + + +class TestStatusHelpers: + def test_status_counts(self): + sub = _ledger([SUBMITTED, SUBMITTED, FETCHED, FAILED]) + assert sub.status_counts() == {SUBMITTED: 2, FETCHED: 1, FAILED: 1} + + def test_minibatches_to_fetch_skips_only_terminal_done(self): + sub = _ledger( + [ + SUBMITTED, + RUNNING, + COMPLETED, + EXPIRED, + FETCH_ERROR, # all of the above are still in the fetch loop + FETCHED, + SUBMIT_ERROR, + FAILED, + CANCELLED, # these four are skipped + ] + ) + indices = [mb["index"] for mb in sub.minibatches_to_fetch()] + assert indices == [0, 1, 2, 3, 4] + + def test_set_status_updates_in_place_with_extra_fields(self): + sub = _ledger([SUBMITTED]) + sub.set_status(0, FETCHED, output_path="/tmp/o.jsonl", fetched_at="T") + mb = sub.minibatches[0] + assert mb["status"] == FETCHED + assert mb["output_path"] == "/tmp/o.jsonl" + assert mb["fetched_at"] == "T" + + def test_set_status_unknown_index_raises(self): + sub = _ledger([SUBMITTED]) + with pytest.raises(KeyError): + sub.set_status(99, FETCHED) + + def test_is_complete_true_when_all_terminal(self): + sub = _ledger([FETCHED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED]) + assert sub.is_complete() + + def test_is_complete_false_with_any_in_flight(self): + for in_flight in (SUBMITTED, RUNNING, COMPLETED, FETCH_ERROR): + sub = _ledger([FETCHED, in_flight]) + assert not sub.is_complete() + + +class TestSerialization: + def test_to_dict_from_dict_round_trip_equal(self): + sub = _ledger([SUBMITTED, FETCHED]) + rebuilt = BatchSubmission.from_dict(sub.to_dict()) + assert rebuilt == sub + # created_at survives by value with tz compared by equality, not identity. + assert rebuilt.created_at == sub.created_at + assert rebuilt.created_at.tzinfo == timezone.utc + + def test_save_load_round_trip_equal(self, tmp_path): + sub = _ledger([SUBMITTED, RUNNING], run_dir=str(tmp_path / "run")) + path = sub.save() + assert path == tmp_path / "run" / "metadata.json" + assert BatchSubmission.load(str(tmp_path / "run")) == sub + + def test_load_accepts_run_dir_or_metadata_path(self, tmp_path): + sub = _ledger([FETCHED], run_dir=str(tmp_path / "run")) + sub.save() + from_dir = BatchSubmission.load(str(tmp_path / "run")) + from_file = BatchSubmission.load(str(tmp_path / "run" / "metadata.json")) + assert from_dir == from_file == sub + + def test_save_to_explicit_dir(self, tmp_path): + sub = _ledger([SUBMITTED], run_dir="/nonexistent/placeholder") + path = sub.save(tmp_path / "elsewhere") + assert path == tmp_path / "elsewhere" / "metadata.json" + assert BatchSubmission.load(tmp_path / "elsewhere") == sub + + def test_from_dict_parses_iso_timestamp_to_utc_equal(self): + sub = _ledger([SUBMITTED]) + data = sub.to_dict() + assert isinstance(data["created_at"], str) + assert data["created_at"].endswith("+00:00") + rebuilt = BatchSubmission.from_dict(data) + # fromisoformat yields timezone(timedelta(0)) which == utc but is not it. + assert rebuilt.created_at.tzinfo == timezone.utc diff --git a/tests/prkit/batch/test_submit_physics_reasoning.py b/tests/prkit/batch/test_submit_physics_reasoning.py index 45eb2ef..b9f978a 100644 --- a/tests/prkit/batch/test_submit_physics_reasoning.py +++ b/tests/prkit/batch/test_submit_physics_reasoning.py @@ -2,21 +2,27 @@ Provider SDKs are faked with ``MagicMock`` so these run fully offline (mirroring ``tests/prkit/core/model_clients/test_batch.py``); artifacts are written under -pytest's ``tmp_path``. The tests cover splitting, JSONL artifacts, submission, -validation, surrogate ids, partial failure, the facade, the Anthropic inline -path, the run-folder layout, ``metadata.json`` content, id round-trips, defaults, -and overwrite behavior. +pytest's ``tmp_path``. Submit now returns the **run-folder path (str)** and saves +one whole-batch :class:`BatchSubmission` ledger to ``metadata.json``; the tests +reconstruct it via ``BatchSubmission.load(run_dir)`` and assert on the ledger's +batch-level fields + per-minibatch dicts. They cover splitting, JSONL artifacts, +submission, validation, surrogate ids, partial failure, the facade, the Anthropic +inline path, the run-folder layout, ``metadata.json`` content, id round-trips, +defaults, and overwrite behavior. """ from __future__ import annotations import json from datetime import datetime, timezone +from pathlib import Path from unittest.mock import MagicMock, patch import pytest from prkit.batch import ( + SUBMIT_ERROR, + SUBMITTED, BatchInputError, BatchSubmission, submit_batch_physics_reasoning, @@ -74,29 +80,31 @@ def _problems(n: int): # --------------------------------------------------------------------------- # class TestSplitting: - def test_1200_problems_split_into_three_batches(self, tmp_path): + def test_1200_problems_split_into_three_minibatches(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(1200), output_dir=tmp_path, batch_size=500 + run_dir = submit_batch_physics_reasoning( + client, _dataset(1200), output_dir=tmp_path, minibatch_size=500 ) - assert [s.num_requests for s in subs] == [500, 500, 200] - assert [s.batch_index for s in subs] == [0, 1, 2] - run_dir = tmp_path / subs[0].input_file_path.split("/")[-3] - files = sorted(p.name for p in (run_dir / "inputs").glob("*.jsonl")) - assert files == ["batch_0000.jsonl", "batch_0001.jsonl", "batch_0002.jsonl"] + sub = BatchSubmission.load(run_dir) + assert [mb["num_requests"] for mb in sub.minibatches] == [500, 500, 200] + assert [mb["index"] for mb in sub.minibatches] == [0, 1, 2] + assert sub.minibatch_count == 3 + files = sorted(p.name for p in (Path(run_dir) / "inputs").glob("*.jsonl")) + assert files == [ + "minibatch_0000.jsonl", + "minibatch_0001.jsonl", + "minibatch_0002.jsonl", + ] class TestJsonlArtifact: def test_one_json_object_per_line_roundtrips(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(3), output_dir=tmp_path, batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, minibatch_size=2 ) lines = ( - (tmp_path / subs[0].input_file_path.split("/")[-3] / "inputs") - .joinpath("batch_0000.jsonl") - .read_text() - .splitlines() + (Path(run_dir) / "inputs" / "minibatch_0000.jsonl").read_text().splitlines() ) assert len(lines) == 2 objs = [json.loads(line) for line in lines] @@ -105,11 +113,11 @@ def test_one_json_object_per_line_roundtrips(self, tmp_path): class TestSubmission: - def test_submit_called_once_per_batch_with_its_chunk(self, tmp_path): + def test_submit_called_once_per_minibatch_with_its_chunk(self, tmp_path): client = _openai_client() client.submit_batch = MagicMock(side_effect=["b0", "b1", "b2"]) submit_batch_physics_reasoning( - client, _dataset(5), output_dir=tmp_path, batch_size=2 + client, _dataset(5), output_dir=tmp_path, minibatch_size=2 ) assert client.submit_batch.call_count == 3 chunk_ids = [ @@ -118,30 +126,32 @@ def test_submit_called_once_per_batch_with_its_chunk(self, tmp_path): ] assert chunk_ids == [["p0", "p1"], ["p2", "p3"], ["p4"]] - def test_receipt_carries_stub_fields(self, tmp_path): + def test_ledger_carries_minibatch_fields(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(2), output_dir=tmp_path, batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, minibatch_size=2 ) - (s,) = subs - assert s.batch_id == "batch_abc" - assert s.num_requests == 2 - assert s.provider == "openai" - assert s.model == "gpt-5.1" - assert s.error is None - assert isinstance(s.submitted_at, datetime) - assert s.submitted_at.tzinfo == timezone.utc - assert s.id_map == {"p0": "p0", "p1": "p1"} + sub = BatchSubmission.load(run_dir) + assert sub.provider == "openai" + assert sub.model == "gpt-5.1" + assert isinstance(sub.created_at, datetime) + assert sub.created_at.tzinfo == timezone.utc + (mb,) = sub.minibatches + assert mb["batch_id"] == "batch_abc" + assert mb["num_requests"] == 2 + assert mb["status"] == SUBMITTED + assert mb["error"] is None + assert mb["id_map"] == {"p0": "p0", "p1": "p1"} def test_display_name_routed_through_submit_metadata(self, tmp_path): client = _openai_client() client.submit_batch = MagicMock(return_value="b0") - subs = submit_batch_physics_reasoning( + run_dir = submit_batch_physics_reasoning( client, _dataset(1), output_dir=tmp_path, run_name="my-run" ) _, kwargs = client.submit_batch.call_args assert kwargs["metadata"]["display_name"] == "my-run" - assert subs[0].metadata["display_name"] == "my-run" + assert BatchSubmission.load(run_dir).display_name == "my-run" class TestValidation: @@ -174,61 +184,30 @@ def test_empty_input_raises(self, tmp_path): with pytest.raises(BatchInputError, match="empty"): submit_batch_physics_reasoning(client, [], output_dir=tmp_path) - def test_non_positive_batch_size_raises(self, tmp_path): + def test_non_positive_minibatch_size_raises(self, tmp_path): client = _openai_client() - with pytest.raises(BatchInputError, match="batch_size must be positive"): + with pytest.raises(BatchInputError, match="minibatch_size must be positive"): submit_batch_physics_reasoning( - client, _dataset(2), output_dir=tmp_path, batch_size=0 + client, _dataset(2), output_dir=tmp_path, minibatch_size=0 ) -class TestBatchSubmitError: - def test_carries_successes_and_failed_for_resume(self): - from prkit.batch import BatchSubmitError - - ok = BatchSubmission( - provider="openai", - model="gpt-5.1", - batch_id="batch_0", - input_file_path="/x/inputs/batch_0000.jsonl", - submitted_at=datetime.now(timezone.utc), - num_requests=1, - batch_index=0, - id_map={"p0": "p0"}, - ) - failed = BatchSubmission( - provider="openai", - model="gpt-5.1", - batch_id="", - input_file_path="/x/inputs/batch_0001.jsonl", - submitted_at=datetime.now(timezone.utc), - num_requests=1, - batch_index=1, - id_map={"p1": "p1"}, - error="RuntimeError: boom", - ) - err = BatchSubmitError("submit failed", successes=[ok], failed=failed) - assert err.successes == [ok] - assert err.failed is failed - assert isinstance(err, RuntimeError) - - class TestSurrogateIds: def test_custom_id_fn_surrogates_recovered_via_id_map(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( + run_dir = submit_batch_physics_reasoning( client, _dataset(2), output_dir=tmp_path, - batch_size=2, + minibatch_size=2, custom_id_fn=lambda p: f"sur-{p.problem_id}", ) - (s,) = subs - assert s.id_map == {"sur-p0": "p0", "sur-p1": "p1"} + sub = BatchSubmission.load(run_dir) + (mb,) = sub.minibatches + assert mb["id_map"] == {"sur-p0": "p0", "sur-p1": "p1"} # The wire custom_id is the surrogate. line = json.loads( - (tmp_path / s.input_file_path.split("/")[-3] / "inputs") - .joinpath("batch_0000.jsonl") + (Path(run_dir) / "inputs" / "minibatch_0000.jsonl") .read_text() .splitlines()[0] ) @@ -236,26 +215,26 @@ def test_custom_id_fn_surrogates_recovered_via_id_map(self, tmp_path): class TestPartialFailure: - def test_middle_batch_failure_recorded_others_succeed(self, tmp_path): + def test_middle_minibatch_failure_recorded_others_succeed(self, tmp_path): client = _openai_client() client.client.batches.create.side_effect = [ MagicMock(id="batch_0"), RuntimeError("boom"), MagicMock(id="batch_2"), ] - subs = submit_batch_physics_reasoning( - client, _dataset(5), output_dir=tmp_path, batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(5), output_dir=tmp_path, minibatch_size=2 ) - assert len(subs) == 3 - assert subs[0].error is None and subs[0].batch_id == "batch_0" - assert subs[1].error is not None and "boom" in subs[1].error - assert subs[1].batch_id == "" - assert subs[2].error is None and subs[2].batch_id == "batch_2" - # Failed batch's input file still exists for resume. - for s in subs: - assert ( - tmp_path / s.input_file_path.split(str(tmp_path) + "/")[-1] - ).exists() + mbs = BatchSubmission.load(run_dir).minibatches + assert len(mbs) == 3 + assert mbs[0]["error"] is None and mbs[0]["batch_id"] == "batch_0" + assert mbs[0]["status"] == SUBMITTED + assert mbs[1]["error"] is not None and "boom" in mbs[1]["error"] + assert mbs[1]["batch_id"] == "" and mbs[1]["status"] == SUBMIT_ERROR + assert mbs[2]["error"] is None and mbs[2]["batch_id"] == "batch_2" + # Every minibatch's input file still exists for resume. + for mb in mbs: + assert Path(mb["input_file_path"]).exists() class TestFacade: @@ -263,7 +242,7 @@ def test_client_facade_delegates_to_module(self, tmp_path): client = _openai_client() ds = _dataset(1) with patch("prkit.batch.submit_batch_physics_reasoning") as mock_submit: - mock_submit.return_value = [] + mock_submit.return_value = str(tmp_path / "run") client.submit_batch_physics_reasoning(ds, output_dir=tmp_path) mock_submit.assert_called_once() args, kwargs = mock_submit.call_args @@ -274,14 +253,11 @@ def test_client_facade_delegates_to_module(self, tmp_path): class TestAnthropicInline: def test_artifact_written_and_inline_list_submitted(self, tmp_path): client = _anthropic_client() - subs = submit_batch_physics_reasoning( - client, _dataset(2), output_dir=tmp_path, batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, minibatch_size=2 ) - (s,) = subs # The .jsonl artifact is still written for Anthropic (audit/resume). - artifact = ( - tmp_path / s.input_file_path.split("/")[-3] / "inputs" / "batch_0000.jsonl" - ) + artifact = Path(run_dir) / "inputs" / "minibatch_0000.jsonl" assert artifact.exists() # submit_batch received the inline list of request dicts (no file upload). _, kwargs = client.client.messages.batches.create.call_args @@ -293,48 +269,48 @@ def test_artifact_written_and_inline_list_submitted(self, tmp_path): class TestRunFolderLayout: def test_layout_and_paths(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(3), output_dir=tmp_path, run_name="run-x", batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, run_name="run-x", minibatch_size=2 ) - run_dir = tmp_path / "run-x" - assert (run_dir / "metadata.json").is_file() - assert (run_dir / "inputs" / "batch_0000.jsonl").is_file() - for s in subs: - assert s.input_file_path.startswith(str(run_dir / "inputs")) - assert s.metadata_path == str(run_dir / "metadata.json") + rd = Path(run_dir) + assert rd == tmp_path / "run-x" + assert (rd / "metadata.json").is_file() + assert (rd / "inputs" / "minibatch_0000.jsonl").is_file() + sub = BatchSubmission.load(run_dir) + assert sub.run_dir == str(rd) + for mb in sub.minibatches: + assert mb["input_file_path"].startswith(str(rd / "inputs")) class TestMetadataContent: - def test_metadata_written_before_return_with_header_and_submissions(self, tmp_path): + def test_metadata_written_before_return_with_header_and_minibatches(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(3), output_dir=tmp_path, run_name="run-x", batch_size=2 + run_dir = submit_batch_physics_reasoning( + client, _dataset(3), output_dir=tmp_path, run_name="run-x", minibatch_size=2 ) - record = json.loads((tmp_path / "run-x" / "metadata.json").read_text()) - assert record["run_id"] == "run-x" + record = json.loads((Path(run_dir) / "metadata.json").read_text()) + assert record["run_name"] == "run-x" assert record["provider"] == "openai" assert record["model"] == "gpt-5.1" - assert record["batch_size"] == 2 + assert record["minibatch_size"] == 2 assert record["total_problems"] == 3 - assert record["num_batches"] == 2 + assert record["minibatch_count"] == 2 assert record["request_kind"] == "free_text" assert record["prkit_api_version"] assert record["dataset"] == {"name": "physreason", "version": "1.0"} - assert record["submissions"] == [s.to_dict() for s in subs] + assert len(record["minibatches"]) == 2 + # metadata.json == BatchSubmission.to_dict() (the consolidated ledger). + assert record == BatchSubmission.load(run_dir).to_dict() def test_bare_list_omits_dataset_field_and_segment(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _problems(2), output_dir=tmp_path, batch_size=2 - ) - record = json.loads( - ( - tmp_path / subs[0].metadata_path.split("/")[-2] / "metadata.json" - ).read_text() + run_dir = submit_batch_physics_reasoning( + client, _problems(2), output_dir=tmp_path, minibatch_size=2 ) + record = json.loads((Path(run_dir) / "metadata.json").read_text()) assert record["dataset"] is None # run name omits the segment -> starts with the slugified model. - assert record["run_id"].startswith("gpt-5.1-") + assert record["run_name"].startswith("gpt-5.1-") class TestIdNormalizationRoundTrip: @@ -348,23 +324,24 @@ class TestIdNormalizationRoundTrip: ) def test_wire_id_lands_verbatim(self, builder, expected, tmp_path): client = builder() - subs = submit_batch_physics_reasoning( - client, _dataset(1), output_dir=tmp_path, batch_size=1 + run_dir = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, minibatch_size=1 ) - assert subs[0].batch_id == expected + assert BatchSubmission.load(run_dir).minibatches[0]["batch_id"] == expected def test_openai_id_accepted_by_poll_batch(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( - client, _dataset(1), output_dir=tmp_path, batch_size=1 + run_dir = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, minibatch_size=1 ) + batch_id = BatchSubmission.load(run_dir).minibatches[0]["batch_id"] client.client.batches.retrieve.return_value = MagicMock( status="completed", output_file_id=None, error_file_id=None, request_counts=MagicMock(total=1, completed=1, failed=0), ) - status = client.poll_batch(subs[0].batch_id) + status = client.poll_batch(batch_id) assert status.batch_id == "batch_abc" @@ -372,18 +349,20 @@ class TestDefaultsAndOverrides: def test_default_output_dir_and_display_name(self, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) client = _openai_client() - subs = submit_batch_physics_reasoning(client, _dataset(1)) - run_id = subs[0].metadata_path.split("/")[-2] + run_dir = submit_batch_physics_reasoning(client, _dataset(1)) + sub = BatchSubmission.load(run_dir) + run_id = sub.run_name assert (tmp_path / "batch_runs" / run_id / "metadata.json").is_file() - assert subs[0].metadata["display_name"] == run_id + assert sub.display_name == run_id def test_explicit_run_name_is_slugified(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning( + run_dir = submit_batch_physics_reasoning( client, _dataset(1), output_dir=tmp_path, run_name="My Eval 001!" ) assert (tmp_path / "my-eval-001").is_dir() - assert subs[0].metadata_path == str(tmp_path / "my-eval-001" / "metadata.json") + assert run_dir == str(tmp_path / "my-eval-001") + assert BatchSubmission.load(run_dir).run_dir == str(tmp_path / "my-eval-001") class TestOverwrite: @@ -402,15 +381,18 @@ def test_overwrite_true_succeeds_and_clears_stale_inputs(self, tmp_path): submit_batch_physics_reasoning( client, _dataset(1), output_dir=tmp_path, run_name="run-x" ) - stale = tmp_path / "run-x" / "inputs" / "batch_0005.jsonl" + stale = tmp_path / "run-x" / "inputs" / "minibatch_0005.jsonl" stale.write_text("stale\n") - subs = submit_batch_physics_reasoning( + run_dir = submit_batch_physics_reasoning( client, _dataset(1), output_dir=tmp_path, run_name="run-x", overwrite=True ) - assert len(subs) == 1 + assert len(BatchSubmission.load(run_dir).minibatches) == 1 assert not stale.exists() - def test_returns_list_of_batch_submission(self, tmp_path): + def test_returns_run_dir_str_holding_metadata(self, tmp_path): client = _openai_client() - subs = submit_batch_physics_reasoning(client, _dataset(1), output_dir=tmp_path) - assert all(isinstance(s, BatchSubmission) for s in subs) + run_dir = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path + ) + assert isinstance(run_dir, str) + assert (Path(run_dir) / "metadata.json").is_file() From 3f89dc88e2a6939b4d48b63d81552cdabd36fb93 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 15:31:52 -0400 Subject: [PATCH 06/14] Add resumable batch fetch (fetch_batch, iter_batch_results) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drive the fetch side of the batch lane over the existing client primitives: poll each non-final minibatch, download terminal-and-retrievable ones, correlate each result back to its problem, and advance + persist the ledger. A bounded, resumable helper — no scoring, no pricing, no end-to-end runner. - `fetch_batch(client, run_dir|submission, *, wait=False, poll_interval=10.0, timeout=None, outputs_dirname="outputs", progress=True) -> BatchSubmission`: poll-once-and-persist by default (resumes across processes via the ledger; skips fetched / terminal-failed / submit_error minibatches), `wait=True` loops to completion with `time.sleep` backoff honouring `timeout`. Writes normalized BatchResult JSONL to `/outputs/minibatch_XXXX.jsonl`; EXPIRED minibatches are still retrieved (partial subset -> fetched), only FAILED/CANCELLED have nothing to fetch; a raised retrieve marks `fetch_error` and is retried next pass. Logs a one-line INFO progress summary (completion line once terminal) on `prkit.batch`. - `iter_batch_results(run_dir|submission) -> Iterator[(problem_id, BatchResult)]`: pure offline reader, correlates custom_id -> problem_id via each minibatch's id_map, emits one result per submitted problem (synthetic ERRORED for ids the provider never returned), drops + counts uncorrelated extras. No network. - Capability gate: `batch_fetch_supported` + `_FETCH_CAPABLE_PROVIDERS` {openai, anthropic, google} (Gemini is "google"; xAI has no batch surface). `fetch_batch` raises `BatchFetchUnsupportedError` up front, never a raw NotImplementedError mid-sweep. - Thin `BaseModelClient.fetch_batch_physics_reasoning` facade (lazy import, mirrors the submit facade). Reuses poll_batch / retrieve_batch_results / batch_types as-is; the wait/backoff loop lives here, not in the client. Leaf discipline: batch_types is imported lazily inside the fetch functions, so `import prkit.batch` stays light at load. test_import_isolation extended with the new symbols (load-time only); test_fetch_batch.py covers the gate, happy path, resume/skip, non-terminal, EXPIRED partials, terminal failures, submit-error skip, cross-provider correlation, completeness, the handle contract, progress reporting, the facade, and the wait loop + timeout. Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 385 ++++++++++++++- src/prkit/core/model_clients/base.py | 22 + tests/prkit/batch/test_fetch_batch.py | 549 +++++++++++++++++++++ tests/prkit/batch/test_import_isolation.py | 14 +- 4 files changed, 945 insertions(+), 25 deletions(-) create mode 100644 tests/prkit/batch/test_fetch_batch.py diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py index dec6b00..8f773d1 100644 --- a/src/prkit/batch/__init__.py +++ b/src/prkit/batch/__init__.py @@ -1,47 +1,67 @@ -"""Batch-mode *submit* for physics-reasoning runs (N4). +"""Batch-mode *submit* and *fetch* for physics-reasoning runs (N4). -This leaf turns a set of :class:`~prkit.core.domain.PhysicsProblem`\\s into -submitted provider batch jobs and one whole-batch :class:`BatchSubmission` ledger. +This leaf drives the discounted provider **batch lane** end to end while staying +a bounded helper — never an orchestrating runner. **Vocabulary** (owner-set): a **batch** is the whole thing a user triggers over a dataset (one :func:`submit_batch_physics_reasoning` call → one *run folder* → one :class:`BatchSubmission` ledger). A **minibatch** is one ``minibatch_size``-problem group = one provider batch job = one ``minibatch_XXXX.jsonl`` = one element of -:attr:`BatchSubmission.minibatches`. ``prkit.batch`` / ``submit_batch_*`` keep -"batch" because it names the *batch lane*, not a unit. +:attr:`BatchSubmission.minibatches`. ``prkit.batch`` / ``submit_batch_*`` / +:func:`fetch_batch` keep "batch" because it names the *batch lane*, not a unit. -:func:`submit_batch_physics_reasoning` preprocesses each problem exactly like the -synchronous ``solve_physics_problem`` path (via the client's -``build_problem_batch_request``), splits the dataset into minibatches of +**Submit half** (Stage 1): :func:`submit_batch_physics_reasoning` preprocesses each +problem exactly like the synchronous ``solve_physics_problem`` path (via the +client's ``build_problem_batch_request``), splits the dataset into minibatches of ``minibatch_size`` problems, writes each minibatch's requests as provider-correct JSONL under one run folder, submits each minibatch, saves the consolidated -``metadata.json`` ledger, and **returns the run-folder path** (a ``str``). The -ledger is mutable on purpose — it is the resume state store (each minibatch carries -a ``status``); disk is the source of truth across the ~24h provider window, so the -object is reconstructed on demand via :meth:`BatchSubmission.load`. Polling / -downloading / correlating the results is a later milestone that consumes the ledger. - -Import discipline: at module load this imports only :mod:`prkit.core.domain` -and the standard library — never ``prkit.api``, the dataset hub, a scorer, the -cost meter, or a provider SDK. The model client is duck-typed. +``metadata.json`` ledger, and **returns the run-folder path** (a ``str``). + +**Fetch half** (Stage 2): :func:`fetch_batch` reconstructs the ledger from that +path, polls each non-final minibatch once (or loops until terminal when +``wait=True``), downloads terminal-and-retrievable minibatches to +``outputs/minibatch_XXXX.jsonl``, and persists the advanced ledger. It is +idempotent/resumable: already-fetched and terminal-failed minibatches are skipped, +so re-runs never re-hit the network. :func:`iter_batch_results` is a pure offline +reader that correlates each persisted result back to its ``problem_id`` via the +minibatch's ``id_map``, emitting one ``(problem_id, BatchResult)`` per submitted +problem. Scoring and pricing are **not** done here — the consumer calls +``prkit.api.Scorer`` / ``Verdict`` directly (cost metering is N6's job). + +Import discipline: at module load this imports only :mod:`prkit.core.domain` and +the standard library — never ``prkit.api``, the dataset hub, a scorer, the cost +meter, or a provider SDK. The model client is duck-typed. The batch-lifecycle +types (:class:`~prkit.core.model_clients.batch_types.BatchResult` etc.) live under +``prkit.core.model_clients`` (an import-isolation forbidden module), so they are +imported **lazily inside** :func:`fetch_batch` / :func:`iter_batch_results`, never +at module load. See ``tests/prkit/batch/test_import_isolation.py``. """ from __future__ import annotations import json +import logging import re -from collections.abc import Callable, Sequence +import time +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from prkit.core.domain import PhysicsDataset, PhysicsProblem +if TYPE_CHECKING: + from prkit.core.model_clients.batch_types import BatchResult, BatchState + __all__ = [ "BatchSubmission", "BatchInputError", + "BatchFetchUnsupportedError", "submit_batch_physics_reasoning", + "fetch_batch", + "iter_batch_results", + "batch_fetch_supported", "dumps_batch_jsonl", "write_batch_jsonl", "validate_batch_requests", @@ -78,22 +98,32 @@ _SKIP_FETCH_STATUSES = frozenset({FETCHED, SUBMIT_ERROR, FAILED, CANCELLED}) # Minibatches in these statuses are terminal for ``is_complete()`` — fetched, or -# terminal-failed with nothing left to retrieve. +# terminal-failed with nothing left to retrieve. ``wait=True`` stops once every +# minibatch is one of these. _COMPLETE_STATUSES = frozenset({FETCHED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED}) +# Providers with a full batch fetch lifecycle (poll + retrieve). Gemini's +# provider string is "google" (not "gemini"); xAI / DeepSeek / Dashscope / Ollama +# have no batch fetch surface and are intentionally absent. +_FETCH_CAPABLE_PROVIDERS = frozenset({"openai", "anthropic", "google"}) + # Providers whose batch line correlates by ``key`` rather than ``custom_id``. _KEY_ID_PROVIDERS = frozenset({"google", "gemini"}) # Anthropic restricts custom ids to this charset/length. _ANTHROPIC_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") _MAX_ID_LEN = 64 +# Leaf-light logger: flows through PRKitLogger handlers when the app configured +# them, plain stdlib logging otherwise. Used for the per-pass progress summary. +_logger = logging.getLogger("prkit.batch") + @dataclass class BatchSubmission: """The whole-batch ledger = one run folder = one ``metadata.json``. - Mutable on purpose: the fetch step advances each minibatch's ``status`` and - persists the ledger, which is what makes fetch idempotent/resumable across + Mutable on purpose: :func:`fetch_batch` advances each minibatch's ``status`` + and persists the ledger, which is what makes fetch idempotent/resumable across processes. Batch-level facts are stored once; the per-minibatch fetch state lives in :attr:`minibatches`, a list of plain dicts (one per provider job). @@ -247,6 +277,15 @@ class BatchInputError(ValueError): pre-existing non-empty run folder (when ``overwrite`` is False).""" +class BatchFetchUnsupportedError(BatchInputError): + """The client's provider has no batch fetch lifecycle (poll + retrieve). + + Raised **up front** by :func:`fetch_batch` for providers outside + :data:`_FETCH_CAPABLE_PROVIDERS` — never a raw ``NotImplementedError`` from + partway through a sweep. + """ + + def _slug(text: str) -> str: """Lowercase, collapse non-``[a-z0-9._-]`` runs to ``-``, trim edges. @@ -361,7 +400,8 @@ def submit_batch_physics_reasoning( Disk is the source of truth across the ~24h provider window, so the ledger is persisted rather than handed back as an object: reconstruct it on demand at - fetch time via :meth:`BatchSubmission.load`. + fetch time via :meth:`BatchSubmission.load` (or just pass the ``run_dir`` to + :func:`fetch_batch` / :func:`iter_batch_results`). A minibatch whose submit raises is recorded with ``status=SUBMIT_ERROR``, ``error`` set, and ``batch_id=""`` (its input file remains for re-submission via @@ -498,3 +538,302 @@ def submit_batch_physics_reasoning( submission.save() return str(run_dir) + + +# --------------------------------------------------------------------------- # +# Fetch half (Stage 2) # +# --------------------------------------------------------------------------- # +def batch_fetch_supported(client: Any) -> bool: + """True only for providers with a full fetch lifecycle (in the allow-list).""" + return getattr(client, "provider", None) in _FETCH_CAPABLE_PROVIDERS + + +def fetch_batch( + client: Any, + submission: BatchSubmission | str | Path, + *, + wait: bool = False, + poll_interval: float = 10.0, + timeout: float | None = None, + outputs_dirname: str = "outputs", + progress: bool = True, +) -> BatchSubmission: + """Poll, download, and persist a submitted batch run; return the ledger. + + Loads the ledger if given a ``run_dir`` (else uses *submission* as-is), then + polls each non-final minibatch once — downloading terminal-and-retrievable ones + to ``//minibatch_XXXX.jsonl`` and advancing + saving + the ledger after each change. With ``wait=True`` the pass repeats every + ``poll_interval`` seconds until every minibatch is terminal (or ``timeout`` + seconds elapse). Idempotent: ``FETCHED`` / terminal-failed / ``SUBMIT_ERROR`` + minibatches are skipped, so re-runs never re-hit the network for finished work. + + ``EXPIRED`` minibatches are still retrieved (they can carry a completed subset); + only ``FAILED`` / ``CANCELLED`` have nothing to fetch. When ``progress=True`` + (the default), a one-line status summary is logged at INFO after each pass. + + Raises: + BatchFetchUnsupportedError: up front, if the client's provider has no fetch + lifecycle (never a raw ``NotImplementedError`` mid-sweep). + """ + if not batch_fetch_supported(client): + provider = getattr(client, "provider", None) or "unknown" + raise BatchFetchUnsupportedError( + f"Provider {provider!r} has no batch fetch lifecycle; fetch-capable " + f"providers are {sorted(_FETCH_CAPABLE_PROVIDERS)}." + ) + + sub = ( + submission + if isinstance(submission, BatchSubmission) + else BatchSubmission.load(submission) + ) + # Lazy import: batch_types sits under prkit.core.model_clients (forbidden at + # leaf load time). The caller already holds a live client, so the package is + # loaded by now anyway. + from prkit.core.model_clients.batch_types import BatchState + + outputs_dir = Path(sub.run_dir) / outputs_dirname + start = time.monotonic() + while True: + newly_fetched = _run_fetch_pass(client, sub, outputs_dir, BatchState) + if progress: + _log_progress(sub, newly_fetched) + if not wait or sub.is_complete(): + break + if timeout is not None and (time.monotonic() - start) >= timeout: + break + time.sleep(poll_interval) + return sub + + +def _run_fetch_pass( + client: Any, + sub: BatchSubmission, + outputs_dir: Path, + batch_state: type[BatchState], +) -> int: + """One poll-and-download pass over the non-final minibatches; return Δ fetched.""" + newly_fetched = 0 + for mb in sub.minibatches_to_fetch(): + index = mb["index"] + old_status = mb["status"] + st = client.poll_batch(mb["batch_id"]) + counts = dict(st.counts) + + if st.state in (batch_state.COMPLETED, batch_state.EXPIRED): + try: + results = list(client.retrieve_batch_results(mb["batch_id"])) + except Exception as exc: # noqa: BLE001 - retried on the next pass + sub.set_status( + index, + FETCH_ERROR, + counts=counts, + error=f"{type(exc).__name__}: {exc}", + ) + _log_transition(index, old_status, FETCH_ERROR) + sub.save() + continue + # COMPLETED always persists (an empty file still marks the minibatch + # final; iter_batch_results synthesizes failures for the missing ids). + # EXPIRED persists only when it carried a partial subset. + if results or st.state == batch_state.COMPLETED: + output_path = outputs_dir / f"minibatch_{index:04d}.jsonl" + _write_results(output_path, results) + sub.set_status( + index, + FETCHED, + counts=counts, + output_path=str(output_path), + fetched_at=_now_iso(), + ) + newly_fetched += 1 + _log_transition(index, old_status, FETCHED, len(results)) + else: + sub.set_status(index, EXPIRED, counts=counts) + _log_transition(index, old_status, EXPIRED) + else: + new_status = _status_for_state(st.state, batch_state) + if new_status is None: # UNKNOWN: keep prior status, refresh counts + sub.set_status(index, old_status, counts=counts) + else: + sub.set_status(index, new_status, counts=counts) + if new_status != old_status: + _log_transition(index, old_status, new_status) + sub.save() + return newly_fetched + + +def _status_for_state(state: BatchState, batch_state: type[BatchState]) -> str | None: + """Map a poll's ``BatchState`` to a minibatch status. + + Returns ``None`` for ``UNKNOWN`` (keep the prior status and keep polling). + ``COMPLETED`` / ``EXPIRED`` are handled by the retrieve path, not here. + """ + return { + batch_state.PENDING: RUNNING, + batch_state.IN_PROGRESS: RUNNING, + batch_state.FAILED: FAILED, + batch_state.CANCELLED: CANCELLED, + }.get(state) + + +def _write_results(path: Path, results: Sequence[BatchResult]) -> None: + """Write normalized ``BatchResult`` lines as JSONL to *path* (creating parents).""" + path.parent.mkdir(parents=True, exist_ok=True) + lines = [ + json.dumps( + { + "custom_id": r.custom_id, + "status": str(r.status), + "text": r.text, + "error": r.error, + }, + ensure_ascii=False, + ) + for r in results + ] + path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") + + +def iter_batch_results( + submission: BatchSubmission | str | Path, +) -> Iterator[tuple[str, BatchResult]]: + """Pure offline reader: yield ``(problem_id, BatchResult)`` in input order. + + Loads the ledger if given a ``run_dir`` (else uses *submission* as-is). For + every ``FETCHED`` minibatch, reads its persisted ``outputs/`` file and + correlates each line's ``custom_id`` back to its ``problem_id`` via that + minibatch's ``id_map``, emitting one result per ``id_map`` entry (input order). + A synthetic ERRORED :class:`BatchResult` is emitted for any submitted id the + provider never returned (completeness); extra/uncorrelated ids are dropped and + counted on the minibatch (``uncorrelated_count``). Reads no network — safe to + re-run for re-scoring. + """ + from prkit.core.model_clients.batch_types import BatchItemStatus, BatchResult + + sub = ( + submission + if isinstance(submission, BatchSubmission) + else BatchSubmission.load(submission) + ) + for mb in sub.minibatches: + if mb["status"] != FETCHED: + continue + id_map: dict[str, str] = mb.get("id_map") or {} + output_path = mb.get("output_path") + + results_by_cid: dict[str, BatchResult] = {} + uncorrelated = 0 + if output_path and Path(output_path).exists(): + for line in _read_jsonl(output_path): + obj = json.loads(line) + cid = str(obj.get("custom_id", "")) + result = BatchResult( + custom_id=cid, + status=_coerce_item_status(obj.get("status"), BatchItemStatus), + text=obj.get("text"), + error=obj.get("error"), + ) + if cid in id_map: + results_by_cid[cid] = result + else: + uncorrelated += 1 + if uncorrelated: + mb["uncorrelated_count"] = uncorrelated + + for cid, problem_id in id_map.items(): + correlated = results_by_cid.get(cid) + if correlated is None: + correlated = BatchResult( + custom_id=cid, + status=BatchItemStatus.ERRORED, + error="No result returned by the provider for this request.", + ) + yield problem_id, correlated + + +def _coerce_item_status(value: Any, batch_item_status: type[Any]) -> Any: + """Map a persisted status string back to ``BatchItemStatus`` (ERRORED on miss).""" + try: + return batch_item_status(value) + except (ValueError, KeyError): + return batch_item_status.ERRORED + + +def _read_jsonl(path: str | Path) -> Iterator[str]: + """Yield non-empty stripped lines from a JSONL file.""" + for line in Path(path).read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped: + yield stripped + + +def _now_iso() -> str: + """Current UTC time as an ISO 8601 string.""" + return datetime.now(timezone.utc).isoformat() + + +# --------------------------------------------------------------------------- # +# Progress reporting (per fetch pass) # +# --------------------------------------------------------------------------- # +def _log_progress(sub: BatchSubmission, newly_fetched: int) -> None: + """Emit the one-line INFO status summary (or completion line) for a pass.""" + counts = sub.status_counts() + n = sub.minibatch_count + fetched = counts.get(FETCHED, 0) + running = ( + counts.get(SUBMITTED, 0) + + counts.get(RUNNING, 0) + + counts.get(COMPLETED, 0) + + counts.get(FETCH_ERROR, 0) + ) + failed = counts.get(FAILED, 0) + counts.get(CANCELLED, 0) + counts.get(EXPIRED, 0) + not_submitted = counts.get(SUBMIT_ERROR, 0) + + if sub.is_complete(): + _logger.info( + "✓ Batch %r complete — %d/%d minibatches fetched, %d failed, " + "%d not submitted.", + sub.run_name, + fetched, + n, + failed, + not_submitted, + ) + return + + problems_done = sum( + mb["num_requests"] for mb in sub.minibatches if mb["status"] == FETCHED + ) + _logger.info( + "Batch %r [%s/%s] — fetched %d/%d minibatches (+%d this pass) · " + "running %d · failed %d · not-submitted %d | results %d/%d", + sub.run_name, + sub.provider, + sub.model, + fetched, + n, + newly_fetched, + running, + failed, + not_submitted, + problems_done, + sub.total_problems, + ) + + +def _log_transition( + index: int, old_status: str, new_status: str, num_results: int | None = None +) -> None: + """Emit a per-minibatch DEBUG transition line (off unless the logger is DEBUG).""" + if num_results is None: + _logger.debug("minibatch %d: %s → %s", index, old_status, new_status) + else: + _logger.debug( + "minibatch %d: %s → %s (%d results)", + index, + old_status, + new_status, + num_results, + ) diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index 8727de9..aee2be0 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -28,6 +28,9 @@ ) if TYPE_CHECKING: + from pathlib import Path + + from prkit.batch import BatchSubmission from prkit.semantics.schema import PhysicsQuestionSemantics T = TypeVar("T", bound=BaseModel) @@ -394,6 +397,25 @@ def submit_batch_physics_reasoning( return _submit_batch(self, problems, **kwargs) + def fetch_batch_physics_reasoning( + self, + run_dir_or_submission: BatchSubmission | str | Path, + **kwargs: Any, + ) -> BatchSubmission: + """Poll + download a submitted batch run; return the updated ledger. + + Thin one-line facade mirroring :meth:`submit_batch_physics_reasoning`: lazily + imports :mod:`prkit.batch` and delegates to :func:`prkit.batch.fetch_batch`. + Typically called with the ``run_dir`` string that submit returned; + reconstructs the ledger via :meth:`BatchSubmission.load` and returns the + freshly-updated :class:`~prkit.batch.BatchSubmission`. See + :func:`prkit.batch.fetch_batch` for the keyword arguments (``wait`` / + ``poll_interval`` / ``timeout`` / ``outputs_dirname`` / ``progress``). + """ + from prkit.batch import fetch_batch + + return fetch_batch(self, run_dir_or_submission, **kwargs) + def _build_batch_request( self, *, diff --git a/tests/prkit/batch/test_fetch_batch.py b/tests/prkit/batch/test_fetch_batch.py new file mode 100644 index 0000000..a63c5e0 --- /dev/null +++ b/tests/prkit/batch/test_fetch_batch.py @@ -0,0 +1,549 @@ +"""Tests for the fetch half: ``fetch_batch`` / ``iter_batch_results`` / capability. + +Fully offline. Submission runs are built with SDK-patched real clients (so the +``id_map``s + run folder are realistic), then polled/downloaded with a small +duck-typed ``_FetchClient`` that maps each ``batch_id`` to a ``BatchState`` (or a +sequence of them, for ``wait=True``) and a list of ``BatchResult``s. Covers the +capability gate, the poll→download happy path, resume/skip, non-terminal, +EXPIRED partials, FAILED/CANCELLED, submit-error skip, cross-provider correlation, +completeness, the run_dir/object handle contract, the progress summary, the +facade, and the ``wait=True`` loop + timeout. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.batch import ( + CANCELLED, + EXPIRED, + FETCHED, + RUNNING, + SUBMIT_ERROR, + BatchFetchUnsupportedError, + BatchSubmission, + batch_fetch_supported, + fetch_batch, + iter_batch_results, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem +from prkit.core.model_clients.batch_types import ( + BatchItemStatus, + BatchResult, + BatchState, + BatchStatus, +) + + +# --------------------------------------------------------------------------- # +# Offline SUBMIT client builders (SDK constructor patched) + a run helper # +# --------------------------------------------------------------------------- # +def _openai_submit_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _anthropic_submit_client(model: str = "claude-opus-4-8"): + with patch("prkit.core.model_clients.anthropic.Anthropic") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.anthropic import AnthropicModel + + return AnthropicModel(model) + + +def _gemini_submit_client(model: str = "gemini-3.5-flash"): + with patch("prkit.core.model_clients.gemini.genai.Client") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.gemini import GeminiModel + + return GeminiModel(model) + + +_SUBMIT_BUILDERS = { + "openai": _openai_submit_client, + "anthropic": _anthropic_submit_client, + "google": _gemini_submit_client, +} + + +def _dataset(n: int): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +def _submit_run( + tmp_path, + *, + n: int, + minibatch_size: int, + batch_ids, + provider: str = "openai", + custom_id_fn=None, + run_name: str = "run", +) -> str: + """Submit a run offline; ``submit_batch`` returns the given ``batch_ids``.""" + client = _SUBMIT_BUILDERS[provider]() + client.submit_batch = MagicMock(side_effect=list(batch_ids)) + return submit_batch_physics_reasoning( + client, + _dataset(n), + output_dir=tmp_path, + run_name=run_name, + minibatch_size=minibatch_size, + custom_id_fn=custom_id_fn, + ) + + +# --------------------------------------------------------------------------- # +# Duck-typed FETCH client # +# --------------------------------------------------------------------------- # +class _FetchClient: + """A minimal poll/retrieve client driven by per-``batch_id`` scripted state.""" + + def __init__(self, provider: str = "openai", *, states=None, results=None): + self.provider = provider + self.model = "model-x" + self._states: dict[str, list[BatchState]] = states or {} + self._results: dict[str, list[BatchResult]] = results or {} + self.poll_calls: list[str] = [] + self.retrieve_calls: list[str] = [] + + def poll_batch(self, batch_id: str) -> BatchStatus: + self.poll_calls.append(batch_id) + seq = self._states[batch_id] + state = seq.pop(0) if len(seq) > 1 else seq[0] + return BatchStatus( + batch_id=batch_id, + state=state, + provider=self.provider, + raw_status=str(state), + counts={"total": 1}, + ) + + def retrieve_batch_results(self, batch_id: str): + self.retrieve_calls.append(batch_id) + return iter(self._results.get(batch_id, [])) + + +def _cids(sub: BatchSubmission, index: int = 0) -> list[str]: + return list(sub.minibatches[index]["id_map"]) + + +# --------------------------------------------------------------------------- # +class TestCapability: + @pytest.mark.parametrize( + "provider,ok", + [ + ("openai", True), + ("anthropic", True), + ("google", True), + ("xai", False), + ("deepseek", False), + ("dashscope", False), + ("ollama", False), + ], + ) + def test_batch_fetch_supported(self, provider, ok): + assert batch_fetch_supported(_FetchClient(provider)) is ok + + def test_unsupported_raises_up_front_without_polling(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + client = _FetchClient("xai", states={"b0": [BatchState.COMPLETED]}) + with pytest.raises(BatchFetchUnsupportedError, match="xai"): + fetch_batch(client, run_dir) + assert client.poll_calls == [] + + +class TestHappyPath: + def test_poll_download_persists_and_correlates(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A0"), + BatchResult(cids[1], BatchItemStatus.SUCCEEDED, text="A1"), + ] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + assert mb["status"] == FETCHED + assert mb["output_path"] and Path(mb["output_path"]).exists() + assert mb["fetched_at"] + assert mb["output_path"] == str( + Path(run_dir) / "outputs" / "minibatch_0000.jsonl" + ) + assert len(Path(mb["output_path"]).read_text().splitlines()) == 2 + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p1"} + assert all(r.succeeded for r in pairs.values()) + assert sub.is_complete() + + +class TestResume: + def test_fetched_minibatches_skipped_on_second_pass(self, tmp_path): + bids = [f"b{i}" for i in range(10)] + run_dir = _submit_run(tmp_path, n=10, minibatch_size=1, batch_ids=bids) + sub0 = BatchSubmission.load(run_dir) + cid_of = {mb["batch_id"]: next(iter(mb["id_map"])) for mb in sub0.minibatches} + states = { + bid: [BatchState.COMPLETED if i < 3 else BatchState.IN_PROGRESS] + for i, bid in enumerate(bids) + } + results = { + bid: [BatchResult(cid_of[bid], BatchItemStatus.SUCCEEDED, text="x")] + for bid in bids[:3] + } + client = _FetchClient("openai", states=states, results=results) + + sub = fetch_batch(client, run_dir, progress=False) + assert sum(mb["status"] == FETCHED for mb in sub.minibatches) == 3 + assert set(client.poll_calls) == set(bids) # first pass polled all 10 + + client.poll_calls.clear() + client.retrieve_calls.clear() + fetch_batch(client, run_dir, progress=False) + # The 3 fetched are skipped (never re-polled); only the 7 running are. + assert not any(b in client.poll_calls for b in bids[:3]) + assert sorted(client.poll_calls) == sorted(bids[3:]) + assert client.retrieve_calls == [] + + +class TestNonTerminal: + def test_in_progress_leaves_running_writes_no_output(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + client = _FetchClient("openai", states={"b0": [BatchState.IN_PROGRESS]}) + sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + assert mb["status"] == RUNNING + assert mb["output_path"] is None + assert mb["counts"] == {"total": 1} # poll counts recorded + assert client.retrieve_calls == [] + # Ledger persisted to disk even on a no-download pass. + assert BatchSubmission.load(run_dir).minibatches[0]["status"] == RUNNING + + +class TestExpired: + def test_expired_partial_subset_persisted_with_synthetic_missing(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.EXPIRED]}, + results={ + "b0": [BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A0")] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + assert sub.minibatches[0]["status"] == FETCHED # partial subset -> fetched + pairs = dict(iter_batch_results(run_dir)) + assert pairs["p0"].succeeded and pairs["p0"].text == "A0" + assert pairs["p1"].status == BatchItemStatus.ERRORED # synthetic completeness + + def test_expired_with_nothing_marks_expired_terminal(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + client = _FetchClient( + "openai", states={"b0": [BatchState.EXPIRED]}, results={"b0": []} + ) + sub = fetch_batch(client, run_dir, progress=False) + assert sub.minibatches[0]["status"] == EXPIRED + assert sub.minibatches[0]["output_path"] is None + assert sub.is_complete() # expired-empty counts terminal for wait/completion + + +class TestTerminalFailures: + def test_failed_and_cancelled_not_retried(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=1, batch_ids=["b0", "b1"]) + client = _FetchClient( + "openai", + states={"b0": [BatchState.FAILED], "b1": [BatchState.CANCELLED]}, + ) + sub = fetch_batch(client, run_dir, progress=False) + assert sub.minibatches[0]["status"] == "failed" + assert sub.minibatches[1]["status"] == CANCELLED + client.poll_calls.clear() + fetch_batch(client, run_dir, progress=False) + assert client.poll_calls == [] # both terminal-failed -> skipped + assert client.retrieve_calls == [] + + +class TestSubmitError: + def test_submit_error_minibatch_never_polled(self, tmp_path): + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=["b0", RuntimeError("boom")]) + run_dir = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, run_name="run", minibatch_size=1 + ) + sub0 = BatchSubmission.load(run_dir) + assert sub0.minibatches[1]["status"] == SUBMIT_ERROR + assert sub0.minibatches[1]["batch_id"] == "" + cid0 = next(iter(sub0.minibatches[0]["id_map"])) + fc = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid0, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + sub = fetch_batch(fc, run_dir, progress=False) + assert fc.poll_calls == ["b0"] # the "" submit-error id is never polled + assert sub.minibatches[0]["status"] == FETCHED + assert sub.minibatches[1]["status"] == SUBMIT_ERROR + + +class TestCorrelation: + @pytest.mark.parametrize("provider", ["openai", "anthropic", "google"]) + def test_custom_id_to_problem_id_each_provider(self, tmp_path, provider): + run_dir = _submit_run( + tmp_path, n=2, minibatch_size=2, batch_ids=["b0"], provider=provider + ) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + provider, + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(c, BatchItemStatus.SUCCEEDED, text=f"ans-{c}") + for c in cids + ] + }, + ) + fetch_batch(client, run_dir, progress=False) + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p1"} + assert all(r.succeeded for r in pairs.values()) + + def test_surrogate_ids_recover_problem_id(self, tmp_path): + run_dir = _submit_run( + tmp_path, + n=2, + minibatch_size=2, + batch_ids=["b0"], + custom_id_fn=lambda p: f"sur-{p.problem_id}", + ) + sub0 = BatchSubmission.load(run_dir) + assert set(sub0.minibatches[0]["id_map"]) == {"sur-p0", "sur-p1"} + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(c, BatchItemStatus.SUCCEEDED, text="x") + for c in _cids(sub0) + ] + }, + ) + fetch_batch(client, run_dir, progress=False) + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p1"} # recovered problem_ids, not surrogates + + +class TestCompleteness: + def test_one_result_per_id_in_order_extra_dropped_and_counted(self, tmp_path): + run_dir = _submit_run(tmp_path, n=3, minibatch_size=3, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) # [p0, p1, p2] in input order + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A0"), + # cids[1] omitted -> synthetic ERRORED + BatchResult(cids[2], BatchItemStatus.SUCCEEDED, text="A2"), + BatchResult("ghost", BatchItemStatus.SUCCEEDED, text="X"), # extra + ] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + pairs = list( + iter_batch_results(sub) + ) # pass the object to read uncorrelated_count + assert [pid for pid, _ in pairs] == [ + "p0", + "p1", + "p2", + ] # exactly one each, in order + by = dict(pairs) + assert by["p0"].succeeded and by["p2"].succeeded + assert ( + by["p1"].status == BatchItemStatus.ERRORED + ) # synthetic for the missing id + assert sub.minibatches[0]["uncorrelated_count"] == 1 # ghost dropped + counted + + +class TestHandleContract: + def test_submit_returns_run_dir_and_fetch_accepts_path_or_object(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + assert isinstance(run_dir, str) + assert (Path(run_dir) / "metadata.json").is_file() + cid = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + + # (a) accepts the run_dir string (loads the ledger from disk) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + loaded = fetch_batch(client, run_dir, progress=False) + assert isinstance(loaded, BatchSubmission) + assert loaded.minibatches[0]["status"] == FETCHED + + # (b) accepts a BatchSubmission directly and threads the same object back + obj = BatchSubmission.load(run_dir) + obj.minibatches[0]["status"] = SUBMIT_ERROR # force a skip to prove no reload + returned = fetch_batch(client, obj, progress=False) + assert returned is obj # the in-memory object is used as-is (no disk reload) + + def test_iter_batch_results_is_offline(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(c, BatchItemStatus.SUCCEEDED, text="A") for c in cids + ] + }, + ) + fetch_batch(client, run_dir, progress=False) + # Reads purely from disk via the run_dir handle — no client passed. + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p1"} + + +class TestProgressReporting: + def test_summary_line_emitted_with_correct_tallies(self, tmp_path, caplog): + bids = [f"b{i}" for i in range(10)] + run_dir = _submit_run(tmp_path, n=10, minibatch_size=1, batch_ids=bids) + cid_of = { + mb["batch_id"]: next(iter(mb["id_map"])) + for mb in BatchSubmission.load(run_dir).minibatches + } + states = { + bid: [BatchState.COMPLETED if i < 3 else BatchState.IN_PROGRESS] + for i, bid in enumerate(bids) + } + results = { + bid: [BatchResult(cid_of[bid], BatchItemStatus.SUCCEEDED, text="x")] + for bid in bids[:3] + } + client = _FetchClient("openai", states=states, results=results) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + fetch_batch(client, run_dir, progress=True) + summaries = [ + r.getMessage() + for r in caplog.records + if r.name == "prkit.batch" and r.getMessage().startswith("Batch ") + ] + assert len(summaries) == 1 + line = summaries[0] + assert "fetched 3/10 minibatches" in line + assert "(+3 this pass)" in line + assert "running 7" in line + assert "failed 0" in line + assert "not-submitted 0" in line + assert "results 3/10" in line + + def test_progress_false_suppresses_summary(self, tmp_path, caplog): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + cid = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + fetch_batch(client, run_dir, progress=False) + assert [r for r in caplog.records if r.name == "prkit.batch"] == [] + + def test_completion_line_only_when_complete(self, tmp_path, caplog): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + cid = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + fetch_batch(client, run_dir, progress=True) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any(m.startswith("✓ Batch ") and "complete" in m for m in msgs) + # The non-complete summary line is not emitted when the pass completes. + assert not any(m.startswith("Batch ") for m in msgs) + + def test_status_counts_tally(self, tmp_path): + run_dir = _submit_run( + tmp_path, n=3, minibatch_size=1, batch_ids=["b0", "b1", "b2"] + ) + cid0 = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + "openai", + states={ + "b0": [BatchState.COMPLETED], + "b1": [BatchState.IN_PROGRESS], + "b2": [BatchState.FAILED], + }, + results={"b0": [BatchResult(cid0, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + sub = fetch_batch(client, run_dir, progress=False) + assert sub.status_counts() == {FETCHED: 1, RUNNING: 1, "failed": 1} + + +class TestFacade: + def test_fetch_facade_delegates_to_module(self): + client = _openai_submit_client() + with patch("prkit.batch.fetch_batch") as mock_fetch: + mock_fetch.return_value = "LEDGER" + out = client.fetch_batch_physics_reasoning("run-dir", wait=True) + mock_fetch.assert_called_once() + args, kwargs = mock_fetch.call_args + assert args[0] is client and args[1] == "run-dir" + assert kwargs["wait"] is True + assert out == "LEDGER" + + +class TestWaitLoop: + def test_wait_loops_until_all_terminal(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + cid = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + "openai", + states={"b0": [BatchState.IN_PROGRESS, BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="A")]}, + ) + with patch("prkit.batch.time.sleep") as slept: + sub = fetch_batch( + client, run_dir, wait=True, poll_interval=0.01, progress=False + ) + assert sub.is_complete() + assert client.poll_calls == ["b0", "b0"] # running, then completed + slept.assert_called() # slept between the two passes + + def test_wait_timeout_stops_before_completion(self, tmp_path): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + client = _FetchClient( + "openai", + states={"b0": [BatchState.IN_PROGRESS]}, # never completes + ) + with patch("prkit.batch.time.sleep") as slept: + sub = fetch_batch( + client, + run_dir, + wait=True, + poll_interval=0.01, + timeout=0, + progress=False, + ) + assert not sub.is_complete() + assert client.poll_calls == ["b0"] # timeout=0 -> exactly one pass + slept.assert_not_called() diff --git a/tests/prkit/batch/test_import_isolation.py b/tests/prkit/batch/test_import_isolation.py index 9f6f690..b2af9e4 100644 --- a/tests/prkit/batch/test_import_isolation.py +++ b/tests/prkit/batch/test_import_isolation.py @@ -5,7 +5,11 @@ SDK, the dataset hub, the semantics/api packages, ``pint``, ``pandas``, or even ``prkit.core.model_clients``. The submit orchestrator depends only on ``prkit.core.domain`` + stdlib at load; the client is duck-typed and the -``prkit.api`` version read is lazy (call-time, not import-time). +``prkit.api`` version read is lazy (call-time, not import-time). The fetch half +(``fetch_batch`` / ``iter_batch_results``) imports ``batch_types`` lazily *inside* +its functions, so the new symbols are checked here for load-time cleanliness only — +a call-time no-leak assertion would fail by design (the live client already loaded +``prkit.core.model_clients``). """ from __future__ import annotations @@ -35,7 +39,13 @@ def test_batch_import_does_not_pull_heavy_deps(): code = textwrap.dedent(f""" import sys import prkit.batch - from prkit.batch import submit_batch_physics_reasoning, BatchSubmission + from prkit.batch import ( + submit_batch_physics_reasoning, + BatchSubmission, + fetch_batch, + iter_batch_results, + batch_fetch_supported, + ) forbidden = {_FORBIDDEN!r} leaked = [name for name in forbidden if name in sys.modules] From 0694092f59312276e75605b8229a966303b98e24 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 19:24:55 -0400 Subject: [PATCH 07/14] Add CONSOLIDATED status, finalize status sets, and shared correlation helper Stage-3 (finalize) structural groundwork on the prkit.batch leaf, ahead of the two new finalize verbs: - New minibatch status CONSOLIDATED (a FETCHED minibatch whose per-problem results/ files are written); added to __all__, _SKIP_FETCH_STATUSES (fetch never re-polls it), and _COMPLETE_STATUSES (terminal for is_complete()). - New status sets _HAS_OUTPUT_STATUSES {FETCHED, CONSOLIDATED} and _RESUBMIT_STATUSES {FAILED, SUBMIT_ERROR, EXPIRED} (excludes CANCELLED). - iter_batch_results' per-minibatch correlation body extracted into a private _iter_minibatch_results(mb) (reused by consolidation), and its status gate widened from == FETCHED to in _HAS_OUTPUT_STATUSES so re-scoring still reads a minibatch's outputs/ file after consolidation. The batch_types import stays lazy inside the helper (leaf import discipline preserved). - New BatchNotTerminalError(BatchInputError) for the resubmit precondition. - New stdlib infra _atomic_write_text (tempfile + os.replace) and _safe_results_filename; add import os / import tempfile. - submit_batch_physics_reasoning(overwrite=True) cleanup also clears results/ and results_manifest.json (the only auto-clear of results/). Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 209 +++++++++++++++++++++++++++--------- 1 file changed, 159 insertions(+), 50 deletions(-) diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py index 8f773d1..5a5a786 100644 --- a/src/prkit/batch/__init__.py +++ b/src/prkit/batch/__init__.py @@ -41,7 +41,9 @@ import json import logging +import os import re +import tempfile import time from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass, field @@ -58,6 +60,7 @@ "BatchSubmission", "BatchInputError", "BatchFetchUnsupportedError", + "BatchNotTerminalError", "submit_batch_physics_reasoning", "fetch_batch", "iter_batch_results", @@ -74,6 +77,7 @@ "FAILED", "CANCELLED", "FETCHED", + "CONSOLIDATED", "SUBMIT_ERROR", "FETCH_ERROR", ] @@ -88,19 +92,36 @@ FAILED = "failed" # batch-level provider failure CANCELLED = "cancelled" # cancelled at the provider FETCHED = "fetched" # results downloaded + persisted to outputs/ (terminal-good) -SUBMIT_ERROR = "submit_error" # never submitted (batch_id == ""); re-submit via Stage 1 +CONSOLIDATED = "consolidated" # results/ files written (post-FETCHED, terminal-good) +SUBMIT_ERROR = "submit_error" # never submitted (batch_id == ""); resubmit target FETCH_ERROR = "fetch_error" # retrieve raised; non-final, retried on the next pass # Minibatches in these statuses are done with the fetch loop and skipped on the # next pass (idempotent resume). EXPIRED is deliberately NOT here: an expired job # is re-polled (cheap snapshot) until its results window truly closes — matches -# the design's skip set exactly. -_SKIP_FETCH_STATUSES = frozenset({FETCHED, SUBMIT_ERROR, FAILED, CANCELLED}) - -# Minibatches in these statuses are terminal for ``is_complete()`` — fetched, or -# terminal-failed with nothing left to retrieve. ``wait=True`` stops once every -# minibatch is one of these. -_COMPLETE_STATUSES = frozenset({FETCHED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED}) +# the design's skip set exactly. CONSOLIDATED is FETCHED-and-finalized, so it is +# skipped too (fetch never re-polls a consolidated minibatch). +_SKIP_FETCH_STATUSES = frozenset( + {FETCHED, CONSOLIDATED, SUBMIT_ERROR, FAILED, CANCELLED} +) + +# Minibatches in these statuses are terminal for ``is_complete()`` — fetched (or +# consolidated), or terminal-failed with nothing left to retrieve. ``wait=True`` +# stops once every minibatch is one of these. +_COMPLETE_STATUSES = frozenset( + {FETCHED, CONSOLIDATED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED} +) + +# Minibatches that have a readable ``outputs/`` file on disk: FETCHED, and +# FETCHED-then-CONSOLIDATED (consolidation never deletes the outputs/ file). The +# offline readers (:func:`iter_batch_results` / :func:`consolidate_batch_results`) +# gate on this so re-scoring still works after finalize. +_HAS_OUTPUT_STATUSES = frozenset({FETCHED, CONSOLIDATED}) + +# Minibatches that :func:`resubmit_failed_minibatches` re-submits: terminal, not +# consolidatable, and NOT a deliberate CANCELLED (excluded by owner decision). This +# is exactly ``_COMPLETE_STATUSES − {FETCHED, CONSOLIDATED, CANCELLED}``. +_RESUBMIT_STATUSES = frozenset({FAILED, SUBMIT_ERROR, EXPIRED}) # Providers with a full batch fetch lifecycle (poll + retrieve). Gemini's # provider string is "google" (not "gemini"); xAI / DeepSeek / Dashscope / Ollama @@ -113,6 +134,14 @@ _ANTHROPIC_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") _MAX_ID_LEN = 64 +# ``problem_id`` is not constrained to a filesystem-safe charset (only wire ids are +# validated), so consolidation renders it to a safe ``.json`` filename: runs +# of non-``[A-Za-z0-9._-]`` collapse to ``_`` and the stem is capped well under the +# 255-byte POSIX name limit. The mapping is lossy, so a collision is possible and +# guarded loudly (never a silent overwrite) by :func:`consolidate_batch_results`. +_UNSAFE_FILENAME_RE = re.compile(r"[^A-Za-z0-9._-]+") +_MAX_RESULTS_STEM_LEN = 200 + # Leaf-light logger: flows through PRKitLogger handlers when the app configured # them, plain stdlib logging otherwise. Used for the per-pass progress summary. _logger = logging.getLogger("prkit.batch") @@ -286,6 +315,16 @@ class BatchFetchUnsupportedError(BatchInputError): """ +class BatchNotTerminalError(BatchInputError): + """The ledger is not terminal, so a terminal-gated finalize step refuses. + + Raised **up front** by :func:`resubmit_failed_minibatches` when some minibatch + is still ``SUBMITTED`` / ``RUNNING`` / ``COMPLETED`` / ``FETCH_ERROR``. Run + :func:`fetch_batch` first to drive every minibatch terminal (which also resolves + any transient ``FETCH_ERROR`` by re-downloading), then resubmit the failures. + """ + + def _slug(text: str) -> str: """Lowercase, collapse non-``[a-z0-9._-]`` runs to ``-``, trim edges. @@ -448,6 +487,17 @@ def submit_batch_physics_reasoning( stale.unlink() for stale in (run_dir / "outputs").glob("minibatch_*.jsonl"): stale.unlink() + # Also clear a stale Stage-3 results set, so reusing a run folder for a fresh + # submit cannot leave old per-problem answers next to a new ledger. This is + # the *only* auto-clear of results/ (consolidate/resubmit never clear it). + results_dir = run_dir / "results" + if results_dir.is_dir(): + for stale in results_dir.iterdir(): + if stale.is_file(): + stale.unlink() + stale_manifest = run_dir / "results_manifest.json" + if stale_manifest.exists(): + stale_manifest.unlink() inputs_dir = run_dir / "inputs" inputs_dir.mkdir(parents=True, exist_ok=True) @@ -702,55 +752,71 @@ def iter_batch_results( """Pure offline reader: yield ``(problem_id, BatchResult)`` in input order. Loads the ledger if given a ``run_dir`` (else uses *submission* as-is). For - every ``FETCHED`` minibatch, reads its persisted ``outputs/`` file and - correlates each line's ``custom_id`` back to its ``problem_id`` via that - minibatch's ``id_map``, emitting one result per ``id_map`` entry (input order). - A synthetic ERRORED :class:`BatchResult` is emitted for any submitted id the - provider never returned (completeness); extra/uncorrelated ids are dropped and - counted on the minibatch (``uncorrelated_count``). Reads no network — safe to - re-run for re-scoring. + every minibatch with a persisted ``outputs/`` file (``FETCHED`` *or* + ``CONSOLIDATED`` — consolidation keeps the file, so re-scoring still works + post-finalize), reads it and correlates each line's ``custom_id`` back to its + ``problem_id`` via that minibatch's ``id_map``, emitting one result per + ``id_map`` entry (input order). A synthetic ERRORED :class:`BatchResult` is + emitted for any submitted id the provider never returned (completeness); + extra/uncorrelated ids are dropped and counted on the minibatch + (``uncorrelated_count``). Reads no network — safe to re-run for re-scoring. """ - from prkit.core.model_clients.batch_types import BatchItemStatus, BatchResult - sub = ( submission if isinstance(submission, BatchSubmission) else BatchSubmission.load(submission) ) for mb in sub.minibatches: - if mb["status"] != FETCHED: - continue - id_map: dict[str, str] = mb.get("id_map") or {} - output_path = mb.get("output_path") - - results_by_cid: dict[str, BatchResult] = {} - uncorrelated = 0 - if output_path and Path(output_path).exists(): - for line in _read_jsonl(output_path): - obj = json.loads(line) - cid = str(obj.get("custom_id", "")) - result = BatchResult( - custom_id=cid, - status=_coerce_item_status(obj.get("status"), BatchItemStatus), - text=obj.get("text"), - error=obj.get("error"), - ) - if cid in id_map: - results_by_cid[cid] = result - else: - uncorrelated += 1 - if uncorrelated: - mb["uncorrelated_count"] = uncorrelated - - for cid, problem_id in id_map.items(): - correlated = results_by_cid.get(cid) - if correlated is None: - correlated = BatchResult( - custom_id=cid, - status=BatchItemStatus.ERRORED, - error="No result returned by the provider for this request.", - ) - yield problem_id, correlated + if mb["status"] in _HAS_OUTPUT_STATUSES: + yield from _iter_minibatch_results(mb) + + +def _iter_minibatch_results( + mb: dict[str, Any], +) -> Iterator[tuple[str, BatchResult]]: + """Correlate one minibatch's persisted ``outputs/`` file back to ``problem_id``. + + Reads ``mb["output_path"]``, maps each line's ``custom_id`` to its ``problem_id`` + via ``mb["id_map"]``, and yields one ``(problem_id, BatchResult)`` per ``id_map`` + entry in input order — synthesizing an ERRORED result for any id the provider + never returned, and dropping/counting extras on ``mb["uncorrelated_count"]``. + Shared by :func:`iter_batch_results` and :func:`consolidate_batch_results`; the + caller gates on status (this helper assumes a readable ``outputs/`` file). Reads + no network. ``batch_types`` is imported lazily here to keep the leaf light. + """ + from prkit.core.model_clients.batch_types import BatchItemStatus, BatchResult + + id_map: dict[str, str] = mb.get("id_map") or {} + output_path = mb.get("output_path") + + results_by_cid: dict[str, BatchResult] = {} + uncorrelated = 0 + if output_path and Path(output_path).exists(): + for line in _read_jsonl(output_path): + obj = json.loads(line) + cid = str(obj.get("custom_id", "")) + result = BatchResult( + custom_id=cid, + status=_coerce_item_status(obj.get("status"), BatchItemStatus), + text=obj.get("text"), + error=obj.get("error"), + ) + if cid in id_map: + results_by_cid[cid] = result + else: + uncorrelated += 1 + if uncorrelated: + mb["uncorrelated_count"] = uncorrelated + + for cid, problem_id in id_map.items(): + correlated = results_by_cid.get(cid) + if correlated is None: + correlated = BatchResult( + custom_id=cid, + status=BatchItemStatus.ERRORED, + error="No result returned by the provider for this request.", + ) + yield problem_id, correlated def _coerce_item_status(value: Any, batch_item_status: type[Any]) -> Any: @@ -774,6 +840,49 @@ def _now_iso() -> str: return datetime.now(timezone.utc).isoformat() +def _atomic_write_text(path: str | Path, text: str) -> None: + """Write *text* to *path* atomically (temp file in the same dir + ``os.replace``). + + The temp file is created in the destination directory so ``os.replace`` is a + same-filesystem rename (atomic on POSIX and Windows); a crash mid-write therefore + never leaves a truncated or corrupt file at *path*. Used for every + ``results/.json`` and ``results_manifest.json`` write so finalize is + crash-safe. + """ + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + dir=str(target.parent), prefix=f".{target.name}.", suffix=".tmp" + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(text) + os.replace(tmp_name, target) + except BaseException: + # Best-effort cleanup if the replace never happened (temp file orphaned). + try: + os.unlink(tmp_name) + except OSError: # pragma: no cover - defensive + pass + raise + + +def _safe_results_filename(problem_id: str) -> str: + """Render *problem_id* as a filesystem-safe ``.json`` results filename. + + Collapses runs of non-``[A-Za-z0-9._-]`` characters to ``_``, strips leading + ``._-`` (so the result is never a dotfile or empty), caps the stem length, and + falls back to ``"problem"`` when nothing usable remains. The mapping is lossy, so + two distinct ids can collide on one name — :func:`consolidate_batch_results` + guards that by refusing to overwrite a file that holds a *different* + ``problem_id`` (it never silently clobbers). + """ + stem = _UNSAFE_FILENAME_RE.sub("_", problem_id).strip("._-") or "problem" + if len(stem) > _MAX_RESULTS_STEM_LEN: + stem = stem[:_MAX_RESULTS_STEM_LEN] + return f"{stem}.json" + + # --------------------------------------------------------------------------- # # Progress reporting (per fetch pass) # # --------------------------------------------------------------------------- # From 47942bf4190c155266fc294706890eb9e045fe29 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 19:32:35 -0400 Subject: [PATCH 08/14] Add consolidate_batch_results + resubmit_failed_minibatches finalize verbs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two Stage-3 (finalize) verbs over a fetched ledger, plus the resubmit client facade and the owner's next-command guidance prompts. Bounded helpers — no dataset loading, no auto-chaining; the consumer still drives the outer loop. - consolidate_batch_results(submission, *, results_dirname="results"): offline, lenient, incremental. Streams each FETCHED-not-yet-CONSOLIDATED minibatch's results to per-problem results/.json (atomic write, one record in RAM at a time — never a whole-run aggregate), marks each minibatch CONSOLIDATED and saves after each (crash-safe/resumable), warns about minibatches not yet succeeded, and refreshes results_manifest.json (atomic, written last). Raises BatchInputError on an empty ledger or a filename collision (two problems sanitizing to one file — never a silent overwrite). - resubmit_failed_minibatches(client, submission): re-submits each FAILED / SUBMIT_ERROR / EXPIRED minibatch (not CANCELLED) by re-reading its persisted inputs/ file. Submit-first / mutate-second / save-third per minibatch; a per item failure becomes SUBMIT_ERROR and the loop continues. Requires a terminal ledger (BatchNotTerminalError otherwise) and a batch-capable provider (BatchFetchUnsupportedError up front). - BaseModelClient.resubmit_failed_minibatches facade (thin, lazy-imports prkit.batch), mirroring fetch_batch_physics_reasoning. - Next-command prompts at the end of submit / fetch (3-way) / resubmit / consolidate via the prkit.batch logger (guidance text, not orchestration). Folds in the one forced existing-test adjustment (the fetch progress test now isolates the per-pass summary from the new end-of-fetch next-command line). Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 369 ++++++++++++++++++++++++++ src/prkit/core/model_clients/base.py | 19 ++ tests/prkit/batch/test_fetch_batch.py | 5 +- 3 files changed, 392 insertions(+), 1 deletion(-) diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py index 5a5a786..5d540c5 100644 --- a/src/prkit/batch/__init__.py +++ b/src/prkit/batch/__init__.py @@ -64,6 +64,8 @@ "submit_batch_physics_reasoning", "fetch_batch", "iter_batch_results", + "consolidate_batch_results", + "resubmit_failed_minibatches", "batch_fetch_supported", "dumps_batch_jsonl", "write_batch_jsonl", @@ -586,6 +588,7 @@ def submit_batch_physics_reasoning( minibatches=minibatches, ) submission.save() + _log_next_after_submit(submission) return str(run_dir) @@ -654,6 +657,8 @@ def fetch_batch( if timeout is not None and (time.monotonic() - start) >= timeout: break time.sleep(poll_interval) + if progress: + _log_next_after_fetch(sub) return sub @@ -946,3 +951,367 @@ def _log_transition( new_status, num_results, ) + + +# --------------------------------------------------------------------------- # +# Finalize half (Stage 3): consolidate + resubmit # +# --------------------------------------------------------------------------- # +def consolidate_batch_results( + submission: BatchSubmission | str | Path, + *, + results_dirname: str = "results", +) -> BatchSubmission: + """Stream every FETCHED-not-yet-CONSOLIDATED minibatch's results to disk. + + For each such minibatch, correlate each output line's ``custom_id`` back to its + ``problem_id`` (via the minibatch ``id_map``) and write + ``//.json`` = + ``{problem_id, custom_id, status, text, error}`` **atomically**, then mark the + minibatch ``CONSOLIDATED`` and persist the ledger. This is incremental and + resumable: a re-run skips already-``CONSOLIDATED`` minibatches and a crash + resumes from the first non-consolidated FETCHED one. Results are **streamed** — + one record is held in memory at a time, never the whole run. + + Offline (no client, like :func:`iter_batch_results`) and **lenient**: it WARNS + (does not raise) when some minibatches are not yet succeeded, consolidating the + succeeded subset; re-run after resubmit + fetch to complete. A refreshed + ``/results_manifest.json`` summary is written last (atomic). + + Raises: + BatchInputError: an empty ledger (no minibatches), or a filename collision + (two problems sanitizing to the same file — never a silent overwrite). + """ + sub = ( + submission + if isinstance(submission, BatchSubmission) + else BatchSubmission.load(submission) + ) + if not sub.minibatches: + raise BatchInputError( + "Ledger has no minibatches to consolidate (empty/degenerate run)." + ) + + pending = [mb for mb in sub.minibatches if mb["status"] not in _HAS_OUTPUT_STATUSES] + if pending: + _logger.warning( + "Consolidating the succeeded subset: %d/%d minibatches not yet succeeded " + "(resubmit + fetch them, then re-run consolidate).", + len(pending), + len(sub.minibatches), + ) + + results_dir = Path(sub.run_dir) / results_dirname + results_dir.mkdir(parents=True, exist_ok=True) # NEVER cleared here (incremental) + + status_counts: dict[str, int] = {} + results_written = 0 + uncorrelated_total = 0 + for mb in sub.minibatches: + if mb["status"] != FETCHED: # CONSOLIDATED already done; the rest are skipped + continue + for problem_id, result in _iter_minibatch_results(mb): + path = results_dir / _safe_results_filename(problem_id) + existing_pid = _existing_problem_id(path) + if existing_pid is not None and existing_pid != problem_id: + raise BatchInputError( + f"Results filename collision: {path.name!r} already holds " + f"problem_id {existing_pid!r}, cannot also write {problem_id!r}. " + "Two problems sanitize to the same file; disambiguate their ids." + ) + status_str = str(result.status) + _atomic_write_text( + path, + json.dumps( + { + "problem_id": problem_id, + "custom_id": result.custom_id, + "status": status_str, + "text": result.text, + "error": result.error, + }, + ensure_ascii=False, + indent=2, + ), + ) + status_counts[status_str] = status_counts.get(status_str, 0) + 1 + results_written += 1 + uncorrelated_total += int(mb.get("uncorrelated_count", 0) or 0) + sub.set_status(mb["index"], CONSOLIDATED) + sub.save() # crash-safe per minibatch (ledger owns resumability) + + _write_results_manifest(sub, status_counts, results_written, uncorrelated_total) + _log_next_after_consolidate(sub, results_dirname) + return sub + + +def _existing_problem_id(path: Path) -> str | None: + """Return the ``problem_id`` recorded in an existing results file, else ``None``. + + Lets :func:`consolidate_batch_results` tell a genuine filename collision (the + file holds a *different* problem) from an idempotent re-write of the same problem + on a resume. A missing/malformed/unreadable file returns ``None`` (a safe + overwrite target — the atomic re-write replaces it cleanly). + """ + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + pid = data.get("problem_id") + return pid if isinstance(pid, str) else None + + +def _write_results_manifest( + sub: BatchSubmission, + status_counts: dict[str, int], + results_written: int, + uncorrelated_total: int, +) -> None: + """Refresh ``/results_manifest.json`` (atomic, written last each call). + + A human/consumer-facing at-a-glance summary at the run-dir root (``results/`` + holds only per-problem files); crash-safety/resumability is owned by the ledger + (``metadata.json``), not this marker. ``minibatches_consolidated`` / + ``fully_consolidated`` are ledger-derived (cumulative); ``results_written`` / + ``status_counts`` / ``uncorrelated_total`` reflect *this call's* newly + consolidated minibatches (per the design's streaming counters). A failed or + CANCELLED minibatch keeps ``fully_consolidated`` false (see Known limitations). + """ + consolidated = sub.status_counts().get(CONSOLIDATED, 0) + manifest = { + "run_name": sub.run_name, + "provider": sub.provider, + "model": sub.model, + "total_problems": sub.total_problems, + "results_written": results_written, + "status_counts": status_counts, + "minibatches_consolidated": consolidated, + "minibatches_total": sub.minibatch_count, + "fully_consolidated": consolidated == sub.minibatch_count, + "uncorrelated_total": uncorrelated_total, + "consolidated_at": _now_iso(), + } + _atomic_write_text( + Path(sub.run_dir) / "results_manifest.json", + json.dumps(manifest, indent=2, ensure_ascii=False), + ) + + +def resubmit_failed_minibatches( + client: Any, + submission: BatchSubmission | str | Path, +) -> BatchSubmission: + """Re-submit each FAILED / SUBMIT_ERROR / EXPIRED minibatch (NOT CANCELLED). + + Re-reads each target minibatch's persisted ``inputs/minibatch_XXXX.jsonl`` and + calls ``client.submit_batch`` with the run's merged metadata, then — **only after + the submit returns a new ``batch_id``** — resets that ledger entry (status → + ``SUBMITTED``, new ``batch_id``, cleared ``output_path`` / ``fetched_at`` / + ``counts`` / ``error``) in memory and on disk. A single minibatch's submit + failure is recorded as ``SUBMIT_ERROR`` (``batch_id=""``, ``error`` set) and the + loop continues. CANCELLED minibatches are left untouched (a deliberate cancel is + not auto-resubmitted — a known dead-end). Returns the updated ledger; the + consumer then re-runs :func:`fetch_batch` on the new jobs. + + Raises: + BatchFetchUnsupportedError: up front, for a provider with no batch lifecycle. + BatchNotTerminalError: if the ledger is not terminal (run ``fetch_batch`` + first to resolve any RUNNING / FETCH_ERROR minibatch). + """ + if not batch_fetch_supported(client): + provider = getattr(client, "provider", None) or "unknown" + raise BatchFetchUnsupportedError( + f"Provider {provider!r} has no batch lifecycle; batch-capable providers " + f"are {sorted(_FETCH_CAPABLE_PROVIDERS)}." + ) + + sub = ( + submission + if isinstance(submission, BatchSubmission) + else BatchSubmission.load(submission) + ) + if not sub.is_complete(): + raise BatchNotTerminalError( + f"Batch {sub.run_name!r} is not terminal; run fetch_batch first to drive " + "every minibatch terminal (resolving any FETCH_ERROR), then resubmit." + ) + + targets = [mb for mb in sub.minibatches if mb["status"] in _RESUBMIT_STATUSES] + if not targets: + _log_next_after_resubmit(sub, resubmitted=0, still_failed=0) + return sub + + # The ledger stores the un-merged user metadata + display_name separately, so + # reconstruct the provider-side merged metadata exactly as submit did (:486). + merged_metadata = {**sub.metadata, "display_name": sub.display_name} + outputs_dir = Path(sub.run_dir) / "outputs" + resubmitted = 0 + still_failed = 0 + for mb in targets: + index = mb["index"] + old_status = mb["status"] + try: + requests = [json.loads(line) for line in _read_jsonl(mb["input_file_path"])] + new_id = client.submit_batch(requests, metadata=merged_metadata) + except Exception as exc: # noqa: BLE001 - recorded on the ledger; loop goes on + sub.set_status( + index, + SUBMIT_ERROR, + batch_id="", + error=f"{type(exc).__name__}: {exc}", + output_path=None, + fetched_at=None, + counts={}, + ) + still_failed += 1 + _log_transition(index, old_status, SUBMIT_ERROR) + sub.save() + continue + # Submit succeeded: drop a stale downloaded artifact (defensive — targets + # normally have output_path=None), then reset the entry to SUBMITTED. + stale_output = outputs_dir / f"minibatch_{index:04d}.jsonl" + if stale_output.exists(): + stale_output.unlink() + sub.set_status( + index, + SUBMITTED, + batch_id=new_id, + submitted_at=_now_iso(), + error=None, + output_path=None, + fetched_at=None, + counts={}, + ) + resubmitted += 1 + _log_transition(index, old_status, SUBMITTED) + sub.save() # crash-safe per minibatch + + _log_next_after_resubmit(sub, resubmitted=resubmitted, still_failed=still_failed) + return sub + + +# --------------------------------------------------------------------------- # +# Next-command guidance (one INFO line telling the human what to run next) # +# --------------------------------------------------------------------------- # +def _log_next_after_submit(sub: BatchSubmission) -> None: + """After submit: point at fetch (a fetch routes any SUBMIT_ERROR afterward).""" + _logger.info( + 'Submitted %d minibatches. Next: client.fetch_batch_physics_reasoning("%s").', + sub.minibatch_count, + sub.run_dir, + ) + + +def _log_next_after_fetch(sub: BatchSubmission) -> None: + """After a fetch pass: the owner's 3-way next-command prompt (§5.6).""" + counts = sub.status_counts() + n = sub.minibatch_count + succeeded = counts.get(FETCHED, 0) + counts.get(CONSOLIDATED, 0) + + if not sub.is_complete(): + running = ( + counts.get(SUBMITTED, 0) + + counts.get(RUNNING, 0) + + counts.get(COMPLETED, 0) + + counts.get(FETCH_ERROR, 0) + ) + _logger.info( + "Batch in progress (%d running). Re-run " + 'fetch_batch_physics_reasoning("%s") later to continue.', + running, + sub.run_dir, + ) + return + + if succeeded == n: + _logger.info( + 'All %d minibatches succeeded. Next: consolidate_batch_results("%s") ' + "to write results/.", + n, + sub.run_dir, + ) + return + + # Terminal, some failed. + failed_total = n - succeeded + resubmittable = sum(counts.get(s, 0) for s in _RESUBMIT_STATUSES) + if resubmittable: + if succeeded: + _logger.info( + '%d minibatches failed. Next: client.resubmit_failed_minibatches("%s"),' + " then fetch again. (%d already succeeded — consolidate_batch_results" + " can capture them now.)", + failed_total, + sub.run_dir, + succeeded, + ) + else: + _logger.info( + '%d minibatches failed. Next: client.resubmit_failed_minibatches("%s"),' + " then fetch again.", + failed_total, + sub.run_dir, + ) + return + + # All failures are CANCELLED dead-ends (none resubmittable; see Known limitations). + if succeeded: + _logger.info( + "%d minibatches were CANCELLED — a dead-end (not auto-resubmitted; see " + 'Known limitations). consolidate_batch_results("%s") can still capture ' + "the %d succeeded.", + failed_total, + sub.run_dir, + succeeded, + ) + else: + _logger.info( + "%d minibatches were CANCELLED — a dead-end (not auto-resubmitted; see " + "Known limitations); nothing left to fetch or consolidate.", + failed_total, + ) + + +def _log_next_after_resubmit( + sub: BatchSubmission, *, resubmitted: int, still_failed: int +) -> None: + """After resubmit: the K-resubmitted / M-failed summary + the fetch-next hint.""" + if resubmitted == 0 and still_failed == 0: + cancelled = sub.status_counts().get(CANCELLED, 0) + if cancelled: + _logger.info( + "Nothing to resubmit: %d CANCELLED minibatch(es) are a dead-end (not " + "auto-resubmitted) and the rest already succeeded. Run " + 'consolidate_batch_results("%s") to capture the succeeded subset.', + cancelled, + sub.run_dir, + ) + else: + _logger.info( + "Nothing to resubmit — all minibatches already succeeded. " + 'Next: consolidate_batch_results("%s").', + sub.run_dir, + ) + return + _logger.info( + "Resubmitted %d minibatches (%d still failed). " + 'Next: client.fetch_batch_physics_reasoning("%s").', + resubmitted, + still_failed, + sub.run_dir, + ) + + +def _log_next_after_consolidate(sub: BatchSubmission, results_dirname: str) -> None: + """After consolidate: 'done' when fully consolidated, else the resume hint.""" + consolidated = sub.status_counts().get(CONSOLIDATED, 0) + n = sub.minibatch_count + if consolidated == n: + _logger.info("Done — results in %s/%s/.", sub.run_dir, results_dirname) + else: + _logger.info( + "Consolidated %d/%d minibatches; the rest are not yet succeeded — " + 'resubmit + fetch them, then re-run consolidate_batch_results("%s").', + consolidated, + n, + sub.run_dir, + ) diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index aee2be0..ce7806a 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -416,6 +416,25 @@ def fetch_batch_physics_reasoning( return fetch_batch(self, run_dir_or_submission, **kwargs) + def resubmit_failed_minibatches( + self, + run_dir_or_submission: BatchSubmission | str | Path, + **kwargs: Any, + ) -> BatchSubmission: + """Re-submit a terminal batch run's failed minibatches; return the ledger. + + Thin one-line facade mirroring :meth:`fetch_batch_physics_reasoning`: lazily + imports :mod:`prkit.batch` and delegates to + :func:`prkit.batch.resubmit_failed_minibatches`. Re-submits each + FAILED / SUBMIT_ERROR / EXPIRED minibatch (not CANCELLED) by re-reading its + persisted ``inputs/`` file; the run must be terminal first (run + :meth:`fetch_batch_physics_reasoning` until it is). Returns the updated + :class:`~prkit.batch.BatchSubmission`; fetch the new jobs next. + """ + from prkit.batch import resubmit_failed_minibatches + + return resubmit_failed_minibatches(self, run_dir_or_submission, **kwargs) + def _build_batch_request( self, *, diff --git a/tests/prkit/batch/test_fetch_batch.py b/tests/prkit/batch/test_fetch_batch.py index a63c5e0..7977035 100644 --- a/tests/prkit/batch/test_fetch_batch.py +++ b/tests/prkit/batch/test_fetch_batch.py @@ -440,10 +440,12 @@ def test_summary_line_emitted_with_correct_tallies(self, tmp_path, caplog): client = _FetchClient("openai", states=states, results=results) with caplog.at_level(logging.INFO, logger="prkit.batch"): fetch_batch(client, run_dir, progress=True) + # The per-pass summary is the line carrying "(+N this pass)"; the Stage-3 + # end-of-fetch next-command line ("Batch in progress …") is excluded. summaries = [ r.getMessage() for r in caplog.records - if r.name == "prkit.batch" and r.getMessage().startswith("Batch ") + if r.name == "prkit.batch" and "this pass" in r.getMessage() ] assert len(summaries) == 1 line = summaries[0] @@ -463,6 +465,7 @@ def test_progress_false_suppresses_summary(self, tmp_path, caplog): results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="A")]}, ) with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() # drop submit's own next-command line; isolate the fetch fetch_batch(client, run_dir, progress=False) assert [r for r in caplog.records if r.name == "prkit.batch"] == [] From 13719ed3e65d8d30fa7b203259c06ab1b690464a Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 19:41:13 -0400 Subject: [PATCH 09/14] Add Stage-3 finalize tests + extend import-isolation and next-command coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New offline test modules and extensions for the finalize half: - test_consolidate_results.py: happy path (per-problem files + manifest at the run-dir root, results/ holds only per-problem files), lenient succeeded-subset (+WARNING), incremental resume (already-CONSOLIDATED not rewritten), streaming (one record per file, manifest written last — guards the no-aggregate rule), filename sanitization + collision (BatchInputError, no silent overwrite), crash-safety before the manifest and mid-minibatch (atomic, no partial file), iter_batch_results after consolidate (the _HAS_OUTPUT_STATUSES widening), empty-ledger guard, and the consolidate next-command prompt. - test_resubmit_minibatches.py: happy path (requests re-read from the persisted input file + merged metadata, entries reset to SUBMITTED with cleared fetch fields), CANCELLED exclusion + only-CANCELLED no-op, terminal + capability preconditions (no submit on failure), per-item failure + continue, submit-first ordering, the client facade, and the resubmit next-command prompt. - test_import_isolation.py: import consolidate_batch_results / resubmit_failed_minibatches in the subprocess load-time cleanliness check. - test_fetch_batch.py: the 3-way end-of-fetch next-command prompt. - test_submit_physics_reasoning.py: submit's next-command prompt and overwrite=True clearing a stale results/ + results_manifest.json. Co-Authored-By: Claude Opus 4.8 --- tests/prkit/batch/test_consolidate_results.py | 414 ++++++++++++++++++ tests/prkit/batch/test_fetch_batch.py | 56 +++ tests/prkit/batch/test_import_isolation.py | 2 + .../prkit/batch/test_resubmit_minibatches.py | 268 ++++++++++++ .../batch/test_submit_physics_reasoning.py | 33 ++ 5 files changed, 773 insertions(+) create mode 100644 tests/prkit/batch/test_consolidate_results.py create mode 100644 tests/prkit/batch/test_resubmit_minibatches.py diff --git a/tests/prkit/batch/test_consolidate_results.py b/tests/prkit/batch/test_consolidate_results.py new file mode 100644 index 0000000..4dd071a --- /dev/null +++ b/tests/prkit/batch/test_consolidate_results.py @@ -0,0 +1,414 @@ +"""Tests for the finalize half: ``consolidate_batch_results`` (Stage 3). + +Fully offline. Runs are built with an SDK-patched submit client (realistic +``id_map`` + input files) then driven to terminal statuses by the duck-typed +``_FetchClient`` from the fetch tests' pattern, so each minibatch lands FETCHED +(carries a SUCCEEDED result) or terminal-failed (none). Covers the happy path, +the lenient succeeded-subset, incremental resume, streaming (one record at a +time), filename sanitization + collision, crash-safety (before the manifest and +mid-minibatch), the manifest layout, ``iter_batch_results`` after consolidate, +the empty-ledger guard, and the next-command prompt. +""" + +from __future__ import annotations + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +import prkit.batch as batch_module +from prkit.batch import ( + CONSOLIDATED, + FETCHED, + BatchInputError, + BatchSubmission, + consolidate_batch_results, + fetch_batch, + iter_batch_results, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem +from prkit.core.model_clients.batch_types import ( + BatchItemStatus, + BatchResult, + BatchState, +) + + +# --------------------------------------------------------------------------- # +# Offline builders (mirroring tests/prkit/batch/test_fetch_batch.py) # +# --------------------------------------------------------------------------- # +def _openai_submit_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _dataset(n: int): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +def _dataset_ids(ids): + problems = [PhysicsProblem(problem_id=pid, question=f"Q-{pid}") for pid in ids] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +def _submit_dataset(tmp_path, dataset, *, minibatch_size, batch_ids, run_name="run"): + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=list(batch_ids)) + return submit_batch_physics_reasoning( + client, + dataset, + output_dir=tmp_path, + run_name=run_name, + minibatch_size=minibatch_size, + ) + + +class _FetchClient: + """Minimal poll/retrieve client driven by per-``batch_id`` scripted state.""" + + def __init__(self, provider="openai", *, states=None, results=None): + self.provider = provider + self.model = "model-x" + self._states = states or {} + self._results = results or {} + self.poll_calls: list[str] = [] + + def poll_batch(self, batch_id: str): + self.poll_calls.append(batch_id) + from prkit.core.model_clients.batch_types import BatchStatus + + seq = self._states[batch_id] + state = seq.pop(0) if len(seq) > 1 else seq[0] + return BatchStatus( + batch_id=batch_id, + state=state, + provider=self.provider, + raw_status=str(state), + counts={"total": 1}, + ) + + def retrieve_batch_results(self, batch_id: str): + return iter(self._results.get(batch_id, [])) + + +_STATE_MAP = { + "fetched": BatchState.COMPLETED, + "failed": BatchState.FAILED, + "expired": BatchState.EXPIRED, + "cancelled": BatchState.CANCELLED, +} + + +def _build_run(tmp_path, *, statuses, run_name="run"): + """One minibatch per status; ``"fetched"`` carries one SUCCEEDED result. + + After this returns, the ledger is terminal and minibatch ``i`` is FETCHED + (text ``ans-i``) or terminal-failed per ``statuses[i]``. + """ + n = len(statuses) + bids = [f"b{i}" for i in range(n)] + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=bids) + run_dir = submit_batch_physics_reasoning( + client, _dataset(n), output_dir=tmp_path, run_name=run_name, minibatch_size=1 + ) + sub0 = BatchSubmission.load(run_dir) + cid_of = {mb["batch_id"]: next(iter(mb["id_map"])) for mb in sub0.minibatches} + states = {bids[i]: [_STATE_MAP[s]] for i, s in enumerate(statuses)} + results = { + bids[i]: [ + BatchResult(cid_of[bids[i]], BatchItemStatus.SUCCEEDED, text=f"ans-{i}") + ] + for i, s in enumerate(statuses) + if s == "fetched" + } + fetch_batch(_FetchClient(states=states, results=results), run_dir, progress=False) + return run_dir + + +# --------------------------------------------------------------------------- # +class TestHappyPath: + def test_all_fetched_writes_per_problem_files_and_manifest(self, tmp_path): + run_dir = _build_run(tmp_path, statuses=["fetched", "fetched", "fetched"]) + sub = consolidate_batch_results(run_dir) + + results_dir = Path(run_dir) / "results" + assert sorted(p.name for p in results_dir.iterdir()) == [ + "p0.json", + "p1.json", + "p2.json", + ] # results/ holds ONLY per-problem files + for i in range(3): + rec = json.loads((results_dir / f"p{i}.json").read_text()) + assert rec == { + "problem_id": f"p{i}", + "custom_id": f"p{i}", + "status": "succeeded", + "text": f"ans-{i}", + "error": None, + } + assert all(mb["status"] == CONSOLIDATED for mb in sub.minibatches) + + # Manifest lives at the run-dir ROOT, not inside results/. + assert (Path(run_dir) / "results_manifest.json").is_file() + assert not (results_dir / "results_manifest.json").exists() + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + assert manifest["fully_consolidated"] is True + assert manifest["minibatches_consolidated"] == 3 + assert manifest["minibatches_total"] == 3 + assert manifest["results_written"] == 3 + assert manifest["status_counts"] == {"succeeded": 3} + assert manifest["uncorrelated_total"] == 0 + assert manifest["run_name"] == sub.run_name + + def test_custom_results_dirname(self, tmp_path): + run_dir = _build_run(tmp_path, statuses=["fetched"]) + consolidate_batch_results(run_dir, results_dirname="answers") + assert (Path(run_dir) / "answers" / "p0.json").is_file() + assert not (Path(run_dir) / "results").exists() + + +class TestLenientSubset: + def test_mixed_ledger_consolidates_only_fetched_and_warns(self, tmp_path, caplog): + run_dir = _build_run(tmp_path, statuses=["fetched", "failed", "fetched"]) + with caplog.at_level(logging.WARNING, logger="prkit.batch"): + sub = consolidate_batch_results(run_dir) + + assert sub.minibatches[0]["status"] == CONSOLIDATED + assert sub.minibatches[1]["status"] == "failed" # untouched + assert sub.minibatches[2]["status"] == CONSOLIDATED + assert sorted(p.name for p in (Path(run_dir) / "results").iterdir()) == [ + "p0.json", + "p2.json", + ] # the failed minibatch's problem is not written + + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + assert manifest["fully_consolidated"] is False + assert manifest["minibatches_consolidated"] == 2 + + warnings = [ + r.getMessage() + for r in caplog.records + if r.name == "prkit.batch" and r.levelno == logging.WARNING + ] + assert any("not yet succeeded" in w for w in warnings) + + +class TestIncrementalResume: + def test_rerun_skips_consolidated_and_adds_new(self, tmp_path): + run_dir = _build_run(tmp_path, statuses=["fetched", "failed"]) + consolidate_batch_results(run_dir) + p0 = Path(run_dir) / "results" / "p0.json" + assert p0.is_file() + mtime0 = p0.stat().st_mtime_ns + + # Simulate the failed minibatch later landing FETCHED (resubmit + fetch). + sub = BatchSubmission.load(run_dir) + mb1 = sub.minibatches[1] + cid1 = next(iter(mb1["id_map"])) + out_path = Path(run_dir) / "outputs" / "minibatch_0001.jsonl" + out_path.write_text( + json.dumps( + { + "custom_id": cid1, + "status": "succeeded", + "text": "late", + "error": None, + } + ) + + "\n", + encoding="utf-8", + ) + sub.set_status( + 1, + FETCHED, + output_path=str(out_path), + fetched_at="2026-01-01T00:00:00+00:00", + ) + sub.save() + + sub2 = consolidate_batch_results(run_dir) + assert all(mb["status"] == CONSOLIDATED for mb in sub2.minibatches) + assert (Path(run_dir) / "results" / "p1.json").is_file() + # The already-consolidated p0.json was NOT rewritten (mtime stable). + assert p0.stat().st_mtime_ns == mtime0 + + +class TestStreaming: + def test_writes_one_record_at_a_time_manifest_last(self, tmp_path, monkeypatch): + run_dir = _build_run(tmp_path, statuses=["fetched", "fetched", "fetched"]) + real = batch_module._atomic_write_text + writes: list[str] = [] + + def spy(path, text): + name = Path(path).name + if name != "results_manifest.json": + obj = json.loads(text) + # One record per file — never a list/aggregate of the whole run. + assert isinstance(obj, dict) and "problem_id" in obj + writes.append(name) + real(path, text) + + monkeypatch.setattr(batch_module, "_atomic_write_text", spy) + consolidate_batch_results(run_dir) + + per_problem = [w for w in writes if w != "results_manifest.json"] + assert per_problem == ["p0.json", "p1.json", "p2.json"] # streamed in order + assert writes[-1] == "results_manifest.json" # marker written last + + +class TestFilenameSafety: + def test_unsafe_problem_id_is_sanitized(self, tmp_path): + run_dir = _submit_dataset( + tmp_path, _dataset_ids(["a/b#1"]), minibatch_size=1, batch_ids=["b0"] + ) + cid = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cid, BatchItemStatus.SUCCEEDED, text="x")]}, + ) + fetch_batch(client, run_dir, progress=False) + consolidate_batch_results(run_dir) + files = [p.name for p in (Path(run_dir) / "results").iterdir()] + assert files == ["a_b_1.json"] + rec = json.loads((Path(run_dir) / "results" / "a_b_1.json").read_text()) + assert rec["problem_id"] == "a/b#1" # original id preserved inside + + def test_filename_collision_raises_no_silent_overwrite(self, tmp_path): + run_dir = _submit_dataset( + tmp_path, _dataset_ids(["a/b", "a:b"]), minibatch_size=2, batch_ids=["b0"] + ) + cids = list(BatchSubmission.load(run_dir).minibatches[0]["id_map"]) + client = _FetchClient( + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(c, BatchItemStatus.SUCCEEDED, text="x") for c in cids + ] + }, + ) + fetch_batch(client, run_dir, progress=False) + with pytest.raises(BatchInputError, match="collision"): + consolidate_batch_results(run_dir) + # The first problem's file was written; the second was refused (not clobbered). + rec = json.loads((Path(run_dir) / "results" / "a_b.json").read_text()) + assert rec["problem_id"] == "a/b" + + +class TestCrashSafety: + def test_crash_before_manifest_leaves_valid_files_and_resumes( + self, tmp_path, monkeypatch + ): + run_dir = _build_run(tmp_path, statuses=["fetched", "fetched"]) + + def boom(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(batch_module, "_write_results_manifest", boom) + with pytest.raises(RuntimeError, match="boom"): + consolidate_batch_results(run_dir) + + # Per-problem files are complete + valid; the manifest never appeared. + assert ( + json.loads((Path(run_dir) / "results" / "p0.json").read_text())[ + "problem_id" + ] + == "p0" + ) + assert not (Path(run_dir) / "results_manifest.json").exists() + + monkeypatch.undo() + sub = consolidate_batch_results(run_dir) + assert all(mb["status"] == CONSOLIDATED for mb in sub.minibatches) + assert (Path(run_dir) / "results_manifest.json").is_file() + + def test_crash_mid_minibatch_resumes_without_partial_file( + self, tmp_path, monkeypatch + ): + run_dir = _build_run(tmp_path, statuses=["fetched", "fetched"]) + real = batch_module._atomic_write_text + + def flaky(path, text): + if Path(path).name == "p1.json": + raise RuntimeError("disk full") + real(path, text) + + monkeypatch.setattr(batch_module, "_atomic_write_text", flaky) + with pytest.raises(RuntimeError, match="disk full"): + consolidate_batch_results(run_dir) + + sub = BatchSubmission.load(run_dir) + assert sub.minibatches[0]["status"] == CONSOLIDATED # saved before the crash + assert sub.minibatches[1]["status"] == FETCHED # never reached CONSOLIDATED + assert (Path(run_dir) / "results" / "p0.json").is_file() + assert not ( + Path(run_dir) / "results" / "p1.json" + ).exists() # atomic: no partial + + monkeypatch.undo() + sub2 = consolidate_batch_results(run_dir) + assert all(mb["status"] == CONSOLIDATED for mb in sub2.minibatches) + assert (Path(run_dir) / "results" / "p1.json").is_file() + + +class TestPostConsolidateReader: + def test_iter_batch_results_still_reads_consolidated(self, tmp_path): + run_dir = _build_run(tmp_path, statuses=["fetched", "fetched"]) + consolidate_batch_results(run_dir) + assert all( + mb["status"] == CONSOLIDATED + for mb in BatchSubmission.load(run_dir).minibatches + ) + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p1"} + assert all(r.succeeded for r in pairs.values()) + + +class TestEmptyLedger: + def test_empty_ledger_raises(self, tmp_path): + sub = BatchSubmission( + run_name="r", + provider="openai", + model="m", + created_at=datetime.now(timezone.utc), + minibatch_size=1, + minibatch_count=0, + total_problems=0, + dataset=None, + request_kind="free_text", + prkit_api_version="1.0", + display_name="r", + run_dir=str(tmp_path / "r"), + metadata={}, + minibatches=[], + ) + sub.save() + with pytest.raises(BatchInputError, match="no minibatches"): + consolidate_batch_results(str(tmp_path / "r")) + + +class TestNextCommand: + def test_done_line_when_fully_consolidated(self, tmp_path, caplog): + run_dir = _build_run(tmp_path, statuses=["fetched"]) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + consolidate_batch_results(run_dir) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any(m.startswith("Done — results in") for m in msgs) + + def test_resume_line_when_partial(self, tmp_path, caplog): + run_dir = _build_run(tmp_path, statuses=["fetched", "failed"]) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + consolidate_batch_results(run_dir) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any("Consolidated 1/2 minibatches" in m for m in msgs) diff --git a/tests/prkit/batch/test_fetch_batch.py b/tests/prkit/batch/test_fetch_batch.py index 7977035..1dac98c 100644 --- a/tests/prkit/batch/test_fetch_batch.py +++ b/tests/prkit/batch/test_fetch_batch.py @@ -550,3 +550,59 @@ def test_wait_timeout_stops_before_completion(self, tmp_path): assert not sub.is_complete() assert client.poll_calls == ["b0"] # timeout=0 -> exactly one pass slept.assert_not_called() + + +class TestNextCommandAfterFetch: + """The §5.6 3-way end-of-fetch next-command prompt (gated by ``progress``).""" + + def _msgs(self, caplog): + return [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + + def test_in_progress_points_back_at_fetch(self, tmp_path, caplog): + run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) + client = _FetchClient("openai", states={"b0": [BatchState.IN_PROGRESS]}) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + fetch_batch(client, run_dir, progress=True) + assert any( + m.startswith("Batch in progress") and "Re-run" in m + for m in self._msgs(caplog) + ) + + def test_all_succeeded_points_at_consolidate(self, tmp_path, caplog): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=1, batch_ids=["b0", "b1"]) + cid_of = { + mb["batch_id"]: next(iter(mb["id_map"])) + for mb in BatchSubmission.load(run_dir).minibatches + } + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED], "b1": [BatchState.COMPLETED]}, + results={ + b: [BatchResult(c, BatchItemStatus.SUCCEEDED, text="x")] + for b, c in cid_of.items() + }, + ) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + fetch_batch(client, run_dir, progress=True) + assert any( + "All 2 minibatches succeeded" in m and "consolidate_batch_results" in m + for m in self._msgs(caplog) + ) + + def test_some_failed_points_at_resubmit(self, tmp_path, caplog): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=1, batch_ids=["b0", "b1"]) + cid0 = next(iter(BatchSubmission.load(run_dir).minibatches[0]["id_map"])) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED], "b1": [BatchState.FAILED]}, + results={"b0": [BatchResult(cid0, BatchItemStatus.SUCCEEDED, text="x")]}, + ) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + fetch_batch(client, run_dir, progress=True) + assert any( + "1 minibatches failed" in m and "resubmit_failed_minibatches" in m + for m in self._msgs(caplog) + ) diff --git a/tests/prkit/batch/test_import_isolation.py b/tests/prkit/batch/test_import_isolation.py index b2af9e4..4ffb2a3 100644 --- a/tests/prkit/batch/test_import_isolation.py +++ b/tests/prkit/batch/test_import_isolation.py @@ -44,6 +44,8 @@ def test_batch_import_does_not_pull_heavy_deps(): BatchSubmission, fetch_batch, iter_batch_results, + consolidate_batch_results, + resubmit_failed_minibatches, batch_fetch_supported, ) diff --git a/tests/prkit/batch/test_resubmit_minibatches.py b/tests/prkit/batch/test_resubmit_minibatches.py new file mode 100644 index 0000000..84b6675 --- /dev/null +++ b/tests/prkit/batch/test_resubmit_minibatches.py @@ -0,0 +1,268 @@ +"""Tests for the finalize half: ``resubmit_failed_minibatches`` (Stage 3). + +Fully offline. Runs are built with an SDK-patched submit client then driven to +terminal statuses by the duck-typed ``_FetchClient``; resubmit is then driven by +a tiny duck-typed ``_Client`` whose ``submit_batch`` is a ``MagicMock``. Covers +the happy path (requests re-read from the persisted input file + merged +metadata, entries reset to SUBMITTED), CANCELLED exclusion + no-op, the terminal +and capability preconditions, per-item failure + continue, submit-first ordering, +the client facade, and the next-command prompt. +""" + +from __future__ import annotations + +import logging +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.batch import ( + CANCELLED, + FETCHED, + SUBMIT_ERROR, + SUBMITTED, + BatchFetchUnsupportedError, + BatchNotTerminalError, + BatchSubmission, + fetch_batch, + resubmit_failed_minibatches, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem +from prkit.core.model_clients.batch_types import ( + BatchItemStatus, + BatchResult, + BatchState, +) + + +# --------------------------------------------------------------------------- # +# Offline builders # +# --------------------------------------------------------------------------- # +def _openai_submit_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _dataset(n: int): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +class _FetchClient: + def __init__(self, provider="openai", *, states=None, results=None): + self.provider = provider + self.model = "model-x" + self._states = states or {} + self._results = results or {} + + def poll_batch(self, batch_id: str): + from prkit.core.model_clients.batch_types import BatchStatus + + seq = self._states[batch_id] + state = seq.pop(0) if len(seq) > 1 else seq[0] + return BatchStatus( + batch_id=batch_id, + state=state, + provider=self.provider, + raw_status=str(state), + counts={"total": 1}, + ) + + def retrieve_batch_results(self, batch_id: str): + return iter(self._results.get(batch_id, [])) + + +class _Client: + """A duck-typed resubmit client: only ``provider`` + ``submit_batch`` are used.""" + + def __init__(self, provider="openai"): + self.provider = provider + self.model = "model-x" + + +_STATE_MAP = { + "fetched": BatchState.COMPLETED, + "failed": BatchState.FAILED, + "expired": BatchState.EXPIRED, + "cancelled": BatchState.CANCELLED, +} + + +def _build_terminal_run(tmp_path, *, statuses, run_name="run"): + """One minibatch per status, driven to a terminal ledger (1 problem each).""" + n = len(statuses) + bids = [f"b{i}" for i in range(n)] + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=bids) + run_dir = submit_batch_physics_reasoning( + client, _dataset(n), output_dir=tmp_path, run_name=run_name, minibatch_size=1 + ) + sub0 = BatchSubmission.load(run_dir) + cid_of = {mb["batch_id"]: next(iter(mb["id_map"])) for mb in sub0.minibatches} + states = {bids[i]: [_STATE_MAP[s]] for i, s in enumerate(statuses)} + results = { + bids[i]: [ + BatchResult(cid_of[bids[i]], BatchItemStatus.SUCCEEDED, text=f"ans-{i}") + ] + for i, s in enumerate(statuses) + if s == "fetched" + } + fetch_batch(_FetchClient(states=states, results=results), run_dir, progress=False) + return run_dir + + +def _build_running_run(tmp_path): + """A non-terminal run (one minibatch left RUNNING).""" + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=["b0"]) + run_dir = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run", minibatch_size=1 + ) + fetch_batch( + _FetchClient(states={"b0": [BatchState.IN_PROGRESS]}), + run_dir, + progress=False, + ) + return run_dir + + +# --------------------------------------------------------------------------- # +class TestHappyPath: + def test_resubmits_failed_and_expired_resetting_each_entry(self, tmp_path): + run_dir = _build_terminal_run( + tmp_path, statuses=["fetched", "failed", "expired"] + ) + sub0 = BatchSubmission.load(run_dir) + client = _Client() + client.submit_batch = MagicMock(side_effect=["new1", "new2"]) + + sub = resubmit_failed_minibatches(client, run_dir) + + assert client.submit_batch.call_count == 2 # FAILED + EXPIRED only + for call in client.submit_batch.call_args_list: + requests = call.args[0] + # Requests are parsed from the persisted input file (real request dicts). + assert requests and all("custom_id" in r for r in requests) + assert call.kwargs["metadata"] == {"display_name": sub0.display_name} + + mb1, mb2 = sub.minibatches[1], sub.minibatches[2] + assert mb1["status"] == SUBMITTED and mb1["batch_id"] == "new1" + assert mb2["status"] == SUBMITTED and mb2["batch_id"] == "new2" + for mb in (mb1, mb2): + assert mb["error"] is None + assert mb["output_path"] is None + assert mb["fetched_at"] is None + assert mb["counts"] == {} + assert sub.minibatches[0]["status"] == FETCHED # the succeeded one is untouched + # Persisted to disk. + assert BatchSubmission.load(run_dir).minibatches[1]["batch_id"] == "new1" + + +class TestCancelledExcluded: + def test_cancelled_minibatch_not_resubmitted(self, tmp_path): + run_dir = _build_terminal_run( + tmp_path, statuses=["fetched", "cancelled", "failed"] + ) + client = _Client() + client.submit_batch = MagicMock(side_effect=["new2"]) + sub = resubmit_failed_minibatches(client, run_dir) + assert client.submit_batch.call_count == 1 # only the FAILED one + assert sub.minibatches[1]["status"] == CANCELLED # untouched dead-end + assert sub.minibatches[2]["status"] == SUBMITTED + + def test_only_cancelled_is_logged_noop(self, tmp_path, caplog): + run_dir = _build_terminal_run(tmp_path, statuses=["fetched", "cancelled"]) + client = _Client() + client.submit_batch = MagicMock() + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + resubmit_failed_minibatches(client, run_dir) + client.submit_batch.assert_not_called() + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any("Nothing to resubmit" in m and "CANCELLED" in m for m in msgs) + + +class TestPreconditions: + def test_non_terminal_ledger_raises_before_any_submit(self, tmp_path): + run_dir = _build_running_run(tmp_path) + client = _Client() + client.submit_batch = MagicMock() + with pytest.raises(BatchNotTerminalError): + resubmit_failed_minibatches(client, run_dir) + client.submit_batch.assert_not_called() + + def test_non_batch_provider_raises_up_front(self, tmp_path): + run_dir = _build_terminal_run(tmp_path, statuses=["failed"]) + client = _Client("xai") + client.submit_batch = MagicMock() + with pytest.raises(BatchFetchUnsupportedError, match="xai"): + resubmit_failed_minibatches(client, run_dir) + client.submit_batch.assert_not_called() + + +class TestPerItemFailure: + def test_one_submit_raises_recorded_and_loop_continues(self, tmp_path, caplog): + run_dir = _build_terminal_run(tmp_path, statuses=["failed", "failed", "failed"]) + client = _Client() + client.submit_batch = MagicMock( + side_effect=["new0", RuntimeError("boom"), "new2"] + ) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + sub = resubmit_failed_minibatches(client, run_dir) + + assert ( + sub.minibatches[0]["status"] == SUBMITTED + and sub.minibatches[0]["batch_id"] == "new0" + ) + assert sub.minibatches[1]["status"] == SUBMIT_ERROR + assert sub.minibatches[1]["batch_id"] == "" + assert "boom" in sub.minibatches[1]["error"] + assert ( + sub.minibatches[2]["status"] == SUBMITTED + and sub.minibatches[2]["batch_id"] == "new2" + ) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any("Resubmitted 2 minibatches (1 still failed)" in m for m in msgs) + + def test_failed_submit_never_leaves_submitted_with_empty_id(self, tmp_path): + run_dir = _build_terminal_run(tmp_path, statuses=["failed"]) + client = _Client() + client.submit_batch = MagicMock(side_effect=RuntimeError("boom")) + sub = resubmit_failed_minibatches(client, run_dir) + mb = sub.minibatches[0] + # Submit-first / mutate-second: a raising submit must NOT yield SUBMITTED+"". + assert not (mb["status"] == SUBMITTED and mb["batch_id"] == "") + assert mb["status"] == SUBMIT_ERROR and mb["batch_id"] == "" + + +class TestFacade: + def test_resubmit_facade_delegates_to_module(self): + client = _openai_submit_client() + with patch("prkit.batch.resubmit_failed_minibatches") as mock_resub: + mock_resub.return_value = "LEDGER" + out = client.resubmit_failed_minibatches("run-dir") + mock_resub.assert_called_once() + args, _ = mock_resub.call_args + assert args[0] is client and args[1] == "run-dir" + assert out == "LEDGER" + + +class TestNextCommand: + def test_resubmit_logs_fetch_next(self, tmp_path, caplog): + run_dir = _build_terminal_run(tmp_path, statuses=["failed", "fetched"]) + client = _Client() + client.submit_batch = MagicMock(side_effect=["new0"]) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + resubmit_failed_minibatches(client, run_dir) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any( + m.startswith("Resubmitted 1 minibatches (0 still failed)") + and "fetch_batch_physics_reasoning" in m + for m in msgs + ) diff --git a/tests/prkit/batch/test_submit_physics_reasoning.py b/tests/prkit/batch/test_submit_physics_reasoning.py index b9f978a..5118a48 100644 --- a/tests/prkit/batch/test_submit_physics_reasoning.py +++ b/tests/prkit/batch/test_submit_physics_reasoning.py @@ -14,6 +14,7 @@ from __future__ import annotations import json +import logging from datetime import datetime, timezone from pathlib import Path from unittest.mock import MagicMock, patch @@ -389,6 +390,23 @@ def test_overwrite_true_succeeds_and_clears_stale_inputs(self, tmp_path): assert len(BatchSubmission.load(run_dir).minibatches) == 1 assert not stale.exists() + def test_overwrite_true_clears_stale_results_and_manifest(self, tmp_path): + client = _openai_client() + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x" + ) + rd = tmp_path / "run-x" + (rd / "results").mkdir() + stale_result = rd / "results" / "p0.json" + stale_result.write_text("{}\n") + stale_manifest = rd / "results_manifest.json" + stale_manifest.write_text("{}\n") + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x", overwrite=True + ) + assert not stale_result.exists() + assert not stale_manifest.exists() + def test_returns_run_dir_str_holding_metadata(self, tmp_path): client = _openai_client() run_dir = submit_batch_physics_reasoning( @@ -396,3 +414,18 @@ def test_returns_run_dir_str_holding_metadata(self, tmp_path): ) assert isinstance(run_dir, str) assert (Path(run_dir) / "metadata.json").is_file() + + +class TestNextCommand: + def test_submit_logs_fetch_next(self, tmp_path, caplog): + client = _openai_client() + with caplog.at_level(logging.INFO, logger="prkit.batch"): + submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run-x" + ) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any( + m.startswith("Submitted 1 minibatches") + and "fetch_batch_physics_reasoning" in m + for m in msgs + ) From 0d3dcab6a0ecea3a377c0b6ecbe66c6bbc86483f Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 19:57:27 -0400 Subject: [PATCH 10/14] Add BATCH_MODE.md: rolled-up batch-mode design & usage reference Single source of truth for the prkit batch lane, reconciling the three per-stage design notes (internal/N4_BATCH_DESIGN_STAGE_{1,2,3}.md, kept as historical rationale). Documents the live submit/fetch/finalize API with all three stages implemented, the ledger and run-folder layout, the status lifecycle, the scoring and cost-meter seams, and a Design-evolution table recording the terminology, return-type, and ledger changes across stages. Co-Authored-By: Claude Opus 4.8 --- BATCH_MODE.md | 482 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 482 insertions(+) create mode 100644 BATCH_MODE.md diff --git a/BATCH_MODE.md b/BATCH_MODE.md new file mode 100644 index 0000000..ed3c6dc --- /dev/null +++ b/BATCH_MODE.md @@ -0,0 +1,482 @@ +# Batch Mode (N4) — Design & Usage Reference + +**Single source of truth** for prkit's discounted provider **batch lane**. This doc is the rolled-up, +up-to-date design: it supersedes the per-stage design notes in `internal/N4_BATCH_DESIGN_STAGE_{1,2,3}.md` +(kept only as historical rationale). When a future stage is designed, **roll its committed decisions into this +file** rather than leaving them stranded in a stage doc — see [Maintaining this doc](#maintaining-this-doc). + +## Implementation status (at a glance) + +| Stage | Capability | Status | Public surface | +|---|---|---|---| +| **1 — Submit** | preprocess → split → write JSONL → submit → ledger | ✅ **Implemented & usable** | `submit_batch_physics_reasoning`, `build_problem_batch_request` | +| **2 — Fetch** | poll → download → correlate → resume | ✅ **Implemented & usable** | `fetch_batch`, `iter_batch_results`, `batch_fetch_supported` | +| **3 — Finalize** | consolidate per-problem results · resubmit failed minibatches | ✅ **Implemented & usable** | `consolidate_batch_results`, `resubmit_failed_minibatches` | + +> **All three stages are implemented and callable** as of 2026-06-22 (104 batch tests pass on +> `feat/n4-batch-mode`). Everything below describes the live API. + +--- + +## Concepts & vocabulary + +The vocabulary was fixed in Stage 2 (it renamed Stage 1's terms — see [Design evolution](#design-evolution)). + +- **batch** — the *whole thing* a user triggers over a dataset. One `submit_batch_physics_reasoning(...)` call + → one **run folder** → one `BatchSubmission` **ledger** (`metadata.json`). +- **minibatch** — one `minibatch_size`-problem group = one provider batch job = one `minibatch_XXXX.jsonl` = + one element of `BatchSubmission.minibatches`. **The minibatch is the unit of success**: it either succeeded + (`FETCHED`) or failed as a whole. (Record-level failures within a fetched minibatch surface as `ERRORED` + results, not as minibatch failures.) +- **run folder** — the consumer-owned directory holding the ledger, the input JSONL, and (after fetch) the + output JSONL. **Disk is the source of truth** across the provider's ~24h window. +- The names `prkit.batch` / `submit_batch_*` / `fetch_batch` keep the word "batch" because it names the + discounted **batch lane**, not a unit. + +### What batch mode is (and is not) + +`prkit.batch` is a **bounded helper**, not an end-to-end runner. It owns the tedious, provider-specific glue a +hub should own once — preprocessing parity, splitting, JSONL writing, submission, polling, downloading, +`problem_id` correlation, and a resumable ledger — and hands control back to the consumer at every seam. + +It deliberately does **not** own: + +- **An end-to-end runner / outer loop.** No dataset loading, no inference orchestration, no auto-chaining. The + consumer drives `submit → fetch → [resubmit → fetch]* → consolidate`. +- **Scoring.** No scorer adapter. Batch mode stops at correlated `BatchResult`s; the consumer calls + `prkit.api.Scorer` / `Verdict` directly (see [Scoring seam](#scoring-seam)). +- **Pricing.** That is N6's job (the cost meter). Batch mode forks no pricing and imports no `prkit.cost`. The + ~50% batch discount is realized at the provider regardless; N6 only makes it *visible* + (see [Cost-meter seam](#cost-meter-n6-seam)). +- **A hidden state store.** The only state is the consumer-owned ledger + output files inside the run folder. + No checkpoint DB, no response cache. +- **Concurrency orchestration.** Minibatches are submitted / polled / fetched sequentially. A consumer wanting + parallelism runs the sync call in its own threads. + +--- + +## Provider support + +| Provider | `client.provider` | Batch input transport | Fetch lifecycle | +|---|---|---|---| +| OpenAI | `"openai"` | JSONL **file** upload | ✅ poll + retrieve (output **and** error file) | +| Anthropic | `"anthropic"` | **inline list** (no file upload) | ✅ poll + retrieve (`.results()` stream) | +| Gemini | **`"google"`** | keyed JSONL **file** upload (`key` → `custom_id`) | ✅ poll + retrieve | +| xAI / DeepSeek / Dashscope / Ollama | — | — | ❌ no batch surface | + +- Gemini's provider string is **`"google"`**, not `"gemini"` — relevant if you key off `client.provider`. +- Fetch capability is gated by `batch_fetch_supported(client)` against `{"openai", "anthropic", "google"}`; + `fetch_batch` raises `BatchFetchUnsupportedError` **up front** for anything else (never a raw + `NotImplementedError` mid-sweep). +- A local `.jsonl` artifact is written under `inputs/` for **all three** providers (reproducibility / inspection + / resume), even Anthropic, whose `submit_batch` actually sends an inline list. + +--- + +## Quick start + +```python +from prkit.datasets import DatasetHub +from prkit.core.model_clients import create_model_client +from prkit.batch import iter_batch_results + +dataset = DatasetHub.load("physreason", variant="full", split="test") +client = create_model_client("gpt-5.1") + +# 1) SUBMIT — preprocess (matches solve_physics_problem), split, write JSONL, submit. +# Returns the run-folder PATH (a str). The ledger is saved to /metadata.json. +run_dir = client.submit_batch_physics_reasoning(dataset) +# Default run folder: ./batch_runs/--/ +# Override: client.submit_batch_physics_reasoning(dataset, output_dir="/data/runs", run_name="eval-001") + +# ... the provider holds results for ~24h. Fetch can happen later, in a different process — +# only the run_dir string is needed (disk is the source of truth). ... + +# 2) FETCH — poll, download ready minibatches, persist, return the fresh ledger. +# Idempotent/resumable: re-run to make progress; fetched minibatches are skipped. +sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) # wait=True loops until terminal +print(sub.is_complete(), sub.status_counts()) + +# 3) READ + SCORE — offline correlation back to problem_id (no network; re-runnable). +for problem_id, result in iter_batch_results(run_dir): # accepts run_dir or the ledger + if result.succeeded: + verdict = scorer.score(result.text, gold[problem_id]) # consumer's scoring seam (no adapter) + # non-success: result.status is ERRORED/EXPIRED/CANCELED, result.text is None, result.error set +``` + +`fetch_batch` without `wait=True` does **one** poll-and-download pass and returns — ideal for a cron / `/loop` +driver that re-invokes until `sub.is_complete()`. Each pass logs a one-line INFO summary on the `prkit.batch` +logger (suppress with `progress=False`). + +--- + +## Public API (implemented) + +All symbols live in the import-light leaf `prkit.batch` (`src/prkit/batch/__init__.py`). The `BaseModelClient` +facades are the one-line ergonomic entry points; the module functions are the same thing without the facade. + +### Submit + +```python +# Module function: +prkit.batch.submit_batch_physics_reasoning( + client, # duck-typed BaseModelClient (openai / anthropic / google) + problems, # PhysicsDataset | Sequence[PhysicsProblem] + *, + output_dir: str | Path = "batch_runs", # root dir that holds run folders + run_name: str | None = None, # default: slug("--") + minibatch_size: int = 500, # problems per provider job; safely under all provider ceilings + instructions: str | None = None, + max_output_tokens: int | None = None, + temperature: float | None = None, + custom_id_fn: Callable[[PhysicsProblem], str] | None = None, # default: problem.problem_id + display_name: str | None = None, # provider-facing label; defaults to run_name + metadata: dict[str, str] | None = None, + overwrite: bool = False, # reuse a non-empty run folder (clears stale minibatch_*.jsonl) +) -> str # returns the run-folder path + +# Facade (mirrors solve_physics_problem ergonomics): +client.submit_batch_physics_reasoning(problems, **kwargs) -> str +``` + +- Builds **one free-text request per problem** via `client.build_problem_batch_request`, which reuses the same + prompt builder + images as the synchronous `solve_physics_problem` path — so **batch prompts ≡ sync prompts**. + Free-text only today (mirrors the `ANSWER_TEXT`-only sync path). +- Writes all `inputs/minibatch_XXXX.jsonl` first (cheap, local), then submits sequentially. A minibatch whose + submit raises is recorded with `status=SUBMIT_ERROR`, `error` set, `batch_id=""`, and its input file present — + so the run is resumable. +- **Correlation ids** default to `problem_id`. Validated up front (fail-fast) for uniqueness and per-provider + limits: OpenAI ≤ 64 chars; Anthropic `^[a-zA-Z0-9_-]{1,64}$`. Pass `custom_id_fn` to map a problem → a + surrogate id; the ledger's per-minibatch `id_map` recovers `problem_id` at read time regardless. + +```python +# Lower-level client method submit uses internally (also useful directly): +client.build_problem_batch_request(problem, *, request_id=None, instructions=None, + max_output_tokens=None, temperature=None, **kwargs) -> dict +``` + +### Fetch + +```python +prkit.batch.fetch_batch( + client, + submission, # BatchSubmission | run_dir str | Path + *, + wait: bool = False, # False: one poll-and-download pass; True: loop to completion + poll_interval: float = 10.0, # seconds between polls when wait=True + timeout: float | None = None, # max seconds when wait=True; None = no cap + outputs_dirname: str = "outputs", + progress: bool = True, # emit the one-line INFO summary after each pass +) -> BatchSubmission # the freshly-loaded, updated ledger + +client.fetch_batch_physics_reasoning(run_dir_or_submission, **kwargs) -> BatchSubmission +``` + +- Polls each non-final minibatch once; on `COMPLETED`/`EXPIRED` it downloads results to + `outputs/minibatch_XXXX.jsonl` and marks the minibatch `FETCHED`. The ledger is saved after **every** change + (crash-safe). +- **Idempotent/resumable:** `FETCHED` / `SUBMIT_ERROR` / `FAILED` / `CANCELLED` (and `CONSOLIDATED`) minibatches + are skipped, so re-runs never re-hit the network for finished work. `EXPIRED` is still *re-polled* (cheap) and + still *retrieved* (it can carry a completed subset). +- A retrieve that raises is recorded as `FETCH_ERROR` (non-terminal) and retried on the next pass. + +### Read results (offline) + +```python +prkit.batch.iter_batch_results(submission) -> Iterator[tuple[str, BatchResult]] +``` + +- Pure offline reader (no network, re-runnable for re-scoring). Accepts a `BatchSubmission` or a `run_dir`. +- For every minibatch with a downloaded output file, correlates each line's `custom_id` → `problem_id` via the + minibatch's `id_map`, yielding **one `(problem_id, BatchResult)` per submitted problem in input order**. +- **Completeness:** any submitted id the provider never returned yields a synthetic `ERRORED` `BatchResult`; + extra/uncorrelated ids are dropped and counted on the minibatch (`uncorrelated_count`). + +### Capability check + +```python +prkit.batch.batch_fetch_supported(client) -> bool # True for openai / anthropic / google +``` + +### Result types (from `prkit.core.model_clients.batch_types`) + +```python +@dataclass(frozen=True) +class BatchResult: + custom_id: str + status: BatchItemStatus # SUCCEEDED | ERRORED | EXPIRED | CANCELED + text: str | None = None # the model's free-text output on success; None otherwise + error: str | None = None # failure description on any non-success outcome + @property + def succeeded(self) -> bool: ... +``` + +### Errors + +| Error | Base | Raised when | +|---|---|---| +| `BatchInputError` | `ValueError` | empty input, duplicate/illegal id, pre-existing non-empty run folder without `overwrite` | +| `BatchFetchUnsupportedError` | `BatchInputError` | `fetch_batch` (or resubmit) called on a non-batch provider | +| `BatchNotTerminalError` | `BatchInputError` | *(Stage 3)* `resubmit_failed_minibatches` called on a non-terminal ledger | + +--- + +## The ledger (`BatchSubmission` / `metadata.json`) + +One run folder = one `metadata.json` = `BatchSubmission.to_dict()`. The ledger is a **mutable status record** +(Stage 1's frozen per-batch receipt was reshaped into this in Stage 2). `submit` builds it, saves it, and +returns the `run_dir`; `fetch` loads, advances, and re-saves it. + +```python +@dataclass # mutable: it is a status ledger +class BatchSubmission: + # ---- batch-level (stored once) ---- + run_name: str + provider: str # "openai" | "anthropic" | "google" + model: str + created_at: datetime # UTC + minibatch_size: int + minibatch_count: int + total_problems: int + dataset: dict | None # {"name", "version"} or None for a bare problem list + request_kind: str # "free_text" + prkit_api_version: str + display_name: str + run_dir: str + metadata: dict[str, str] + # ---- per-minibatch ledger (mutable, fetch-updated) ---- + minibatches: list[dict] + + # pure state helpers (NO network): + def minibatches_to_fetch(self) -> list[dict]: ... # status not in the fetch-skip set + def set_status(self, index: int, status: str, **fields) -> None: ... + def is_complete(self) -> bool: ... # every minibatch terminal + def status_counts(self) -> dict[str, int]: ... # {status: count} + # serialization: + def to_dict(self) -> dict: ... # JSON-safe (datetimes -> ISO 8601) + @classmethod + def from_dict(cls, data: dict) -> "BatchSubmission": ... + def save(self, run_dir=None) -> Path: ... # write /metadata.json + @classmethod + def load(cls, run_dir_or_metadata_path) -> "BatchSubmission": ... +``` + +Each `minibatches[i]` dict: + +```python +{ + "index": int, # 0-based ordinal + "batch_id": str, # provider job id ("" => submit failed) + "status": str, # one of the status constants below + "input_file_path": str, # /inputs/minibatch_XXXX.jsonl + "num_requests": int, + "id_map": dict[str, str], # wire custom_id / key -> problem_id (per minibatch) + "submitted_at": str | None, # ISO 8601 + "error": str | None, # local submit error (with SUBMIT_ERROR) + "output_path": str | None, # /outputs/... (set on fetch) + "fetched_at": str | None, # ISO 8601 (set on fetch) + "counts": dict[str, int], # last poll's request_counts + "endpoint": str | None, # audit-only (provider parity) + "completion_window": str | None, # audit-only + # "uncorrelated_count": int # added by iter_batch_results when extra ids were dropped +} +``` + +### Run-folder layout + +``` +// + metadata.json # the BatchSubmission ledger (mutable; advanced by fetch) + inputs/ + minibatch_0000.jsonl # one provider-correct request line per problem + minibatch_0001.jsonl + ... + outputs/ # written by fetch: one normalized-BatchResult JSONL per fetched minibatch + minibatch_0000.jsonl # lines: {"custom_id", "status", "text", "error"} + ... + results/ # written by consolidate (Stage 3): one file per problem + .json + results_manifest.json # written by consolidate (Stage 3): consolidation summary +``` + +--- + +## Status lifecycle + +Minibatch status constants (`prkit.batch`, plain strings): +`SUBMITTED`, `RUNNING`, `COMPLETED`, `EXPIRED`, `FAILED`, `CANCELLED`, `FETCHED`, `CONSOLIDATED`, +`SUBMIT_ERROR`, `FETCH_ERROR`. + +``` +SUBMITTED ─poll→ RUNNING ─poll→ COMPLETED ─retrieve→ FETCHED ─consolidate→ CONSOLIDATED (terminal-good) + └poll→ EXPIRED(partial+results) → FETCHED (counts as success) + └poll→ EXPIRED(empty) ─────────┐ + └poll→ FAILED ─────────────────┤ resubmit (Stage 3) ─→ SUBMITTED (re-enters loop) +SUBMIT_ERROR ────────────────────────────────────────────┘ + └poll→ CANCELLED (terminal; NOT auto-resubmitted — see Stage 3 limitation) +FETCH_ERROR (non-terminal) ─re-run fetch_batch→ ... +``` + +**Poll state → minibatch status** (the fetch-pass mapping): + +| `BatchState` from poll | fetch action | resulting status | +|---|---|---| +| `PENDING` / `IN_PROGRESS` | none | `RUNNING` | +| `COMPLETED` | retrieve → write `outputs/` (empty file still marks it final) | `FETCHED` | +| `EXPIRED` | retrieve; persist if any results returned | `FETCHED` (partial) if any, else `EXPIRED` | +| `FAILED` | none | `FAILED` | +| `CANCELLED` | none | `CANCELLED` | +| `UNKNOWN` | none (keep polling) | unchanged | +| retrieve raised | record error | `FETCH_ERROR` (retried next pass) | +| local submit failure (`batch_id == ""`) | skip | `SUBMIT_ERROR` | + +Key status sets (internal, but they explain behavior): + +- `_SKIP_FETCH_STATUSES = {FETCHED, CONSOLIDATED, SUBMIT_ERROR, FAILED, CANCELLED}` — skipped by the fetch loop. + `EXPIRED` is intentionally **not** here (re-polled until its window truly closes). +- `_COMPLETE_STATUSES = {FETCHED, CONSOLIDATED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED}` — `is_complete()` / + `wait=True` stop once every minibatch is one of these. +- `_HAS_OUTPUT_STATUSES = {FETCHED, CONSOLIDATED}` — minibatches with a readable `outputs/` file. + +--- + +## Stage 3 — Finalize + +Stage 3 closes the loop with two **bounded** verbs over a fetched ledger. It is still **not** an end-to-end +runner — the consumer owns the outer loop; each verb does one step and returns the ledger. + +### `consolidate_batch_results` (offline, incremental) + +```python +prkit.batch.consolidate_batch_results( + submission, # BatchSubmission | run_dir str | Path + *, + results_dirname: str = "results", +) -> BatchSubmission +``` + +- Streams every `FETCHED`-but-not-yet-`CONSOLIDATED` minibatch's results into **per-problem files** + `results/.json` = `{problem_id, custom_id, status, text, error}`, then marks each minibatch + `CONSOLIDATED` and persists the ledger (so a re-run skips it). Writes a `results_manifest.json` summary at the + run-dir root. +- **Lenient + incremental:** consolidates the succeeded subset and *warns* (does not raise) about minibatches + not yet succeeded; safe to run mid-flight and re-run after the rest land. Streams one record at a time — no + in-memory aggregate. **Offline, no client** (mirrors `iter_batch_results`). +- **Filename safety:** `problem_id` is sanitized to a filesystem-safe name; a collision (two problems mapping to + the same file) raises `BatchInputError` rather than silently overwriting. + +### `resubmit_failed_minibatches` (needs client, terminal-gated) + +```python +prkit.batch.resubmit_failed_minibatches(client, submission) -> BatchSubmission +client.resubmit_failed_minibatches(run_dir_or_submission, **kwargs) -> BatchSubmission # facade +``` + +- Re-submits each `FAILED` / `SUBMIT_ERROR` / `EXPIRED` minibatch (`_RESUBMIT_STATUSES`) by re-reading its + persisted `inputs/minibatch_XXXX.jsonl` and calling `client.submit_batch`, resetting the ledger entry to + `SUBMITTED` with a new `batch_id` **only after submit succeeds** (a failed submit → `SUBMIT_ERROR`, loop + continues). The consumer then re-runs `fetch_batch` on the new jobs. +- **Excludes `CANCELLED`** (a cancel can be deliberate; auto-resubmit would fight intent). +- Requires a terminal ledger (`is_complete()`), else raises `BatchNotTerminalError` (run `fetch_batch` first); + raises `BatchFetchUnsupportedError` up front for a non-batch provider. + +### Intended end-to-end loop (consumer-driven) + +```python +run_dir = client.submit_batch_physics_reasoning(dataset) +sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) +while not all(mb["status"] in ("fetched", "consolidated") for mb in sub.minibatches): + sub = client.resubmit_failed_minibatches(run_dir) # Stage 3 + sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) +sub = consolidate_batch_results(run_dir) # Stage 3 -> results/.json +for problem_id, result in iter_batch_results(run_dir): + if result.succeeded: + verdict = scorer.score(result.text, gold[problem_id]) +``` + +### Known limitation (Stage 3) + +A `CANCELLED` minibatch is neither resubmitted nor consolidatable, so a batch containing one can never reach +"fully consolidated" and a naïve `while not all-done` loop would spin. Surfaced via the manifest's +`fully_consolidated` flag and the next-command guidance; a future stage may add an opt-in +`force_resubmit_cancelled` path. + +### Status → finalize action (single source of truth) + +| minibatch `status` | consolidate | resubmit | +|---|---|---| +| `FETCHED` | write `results/.json` → `CONSOLIDATED` | skip | +| `CONSOLIDATED` | skip (done) | skip | +| `FAILED` / `SUBMIT_ERROR` / `EXPIRED` | skip + warn | re-read input → `submit_batch` → `SUBMITTED` (or `SUBMIT_ERROR`) | +| `CANCELLED` | skip + warn (dead-end) | skip (known limitation) | +| `RUNNING` / `SUBMITTED` / `COMPLETED` / `FETCH_ERROR` | skip + warn | precondition `is_complete()` fails → raise | + +--- + +## Scoring seam + +Batch mode stops at correlated `BatchResult`s. The consumer scores by reading `outputs/` via +`iter_batch_results` (or the per-problem `results/.json` files after consolidation) and calling +`prkit.api.Scorer.score(prediction, reference) -> Verdict` directly. The reference/gold answers live on the +consumer's `PhysicsProblem`s, not in the ledger. `prkit.batch` imports no `prkit.api`, accepts no `Scorer`, and +owns no references — scoring is genuinely the consumer's. + +## Cost-meter (N6) seam + +Pricing is N6's job; batch mode forks none and imports no `prkit.cost`. **Forward dependency:** `BatchResult` +carries only `custom_id` / `status` / `text` / `error` — **no token usage** — so batch-path pricing is not +possible until N6 ships and adds a usage channel to the batch path. The ~50% discount is realized at the +provider regardless; N6 only makes it *visible*. + +## Import discipline (leaf-light) + +`prkit.batch` imports only the standard library + `prkit.core.domain` at module load — never `prkit.api`, the +dataset hub, a scorer, the cost meter, or a provider SDK. Batch-lifecycle types (`BatchResult`, +`BatchItemStatus`, `BatchState`) live under `prkit.core.model_clients` (an import-isolation forbidden module), +so they are imported **lazily inside** `fetch_batch` / `iter_batch_results`. The `BaseModelClient` facades +likewise lazily import `prkit.batch`. The client is duck-typed throughout. Enforced by +`tests/prkit/batch/test_import_isolation.py` (load-time cleanliness only). + +--- + +## Design evolution + +Decisions that changed across stages — recorded so the renames and contract shifts are unambiguous. + +| Stage 1 (original) | Current (Stage 2+) | Why | +|---|---|---| +| "run" = whole submit; "batch" = unit | **"batch"** = whole submit; **"minibatch"** = unit | clearer naming; "batch" names the lane | +| `batch_size`, `num_batches`, `batch_index`, `batch_*.jsonl` | `minibatch_size`, `minibatch_count`, `index`, `minibatch_*.jsonl` | follows the rename | +| `BatchSubmission` = frozen per-batch receipt | **mutable whole-batch ledger** (+ `from_dict`/`save`/`load`/status helpers) | fetch needs a resumable status record | +| submit returns `list[BatchSubmission]` | submit returns **`run_dir` str** | disk is the source of truth across the ~24h gap; an in-memory object would be stale by fetch time | +| `metadata.json` = immutable audit record, "never read back" | **mutable resume ledger** prkit reads & updates | this is what makes fetch idempotent (boundary preserved: still consumer-owned, no *hidden* store) | +| xAI = "build-only" batch | xAI = **no batch surface** | Stage-1 cleanup deleted its only (unused, structured) batch method | +| `iter_batch_results` gates on `== FETCHED` | *(Stage 3)* widens to `_HAS_OUTPUT_STATUSES` so re-scoring still works after consolidation | consolidate keeps `outputs/` files | +| `submit(overwrite=True)` clears `inputs/`+`outputs/` | *(Stage 3)* also clears `results/`+`results_manifest.json` | a reused folder must not leave a stale results set | + +Also removed during Stage 1 (no longer part of the contract): the reserved `api.Runner` Protocol, and the +unused structured-batch request/response wrappers. The structured-output *engine* (`StructuredOutputPlan` / +`.parse()`) was untouched. + +--- + +## Maintaining this doc + +This file is the rollup. When a new stage is designed and its decisions are owner-approved: + +1. Write/keep the detailed rationale in `internal/N4_BATCH_DESIGN_STAGE_.md` (the working design + Q&A). +2. **Fold the committed surface into this file**: update the status table, [Public API](#public-api-implemented), + the ledger/layout, the lifecycle, and add a row to [Design evolution](#design-evolution) for anything that + *changes* an earlier decision. +3. Flip a capability from 🟡 designed to ✅ implemented **only** when it is actually callable in + `src/prkit/batch/__init__.py` (and its facade, if any, in `base.py`) — re-verify against the code, not the + stage doc. + +### Sources (historical rationale) + +- `internal/N4_BATCH_DESIGN_STAGE_1.md` — submit half; Q1–Q6; `Runner` + structured-batch cleanup. +- `internal/N4_BATCH_DESIGN_STAGE_2.md` — fetch half; ledger reshape; Stage-1 amendments. +- `internal/N4_BATCH_DESIGN_STAGE_3.md` — finalize (consolidate + resubmit); Stage-1/2 amendments. +- Code: `src/prkit/batch/__init__.py`, `src/prkit/core/model_clients/base.py` (facades), + `src/prkit/core/model_clients/batch_types.py` (`BatchResult` / `BatchItemStatus` / `BatchState`). +- Roadmap: `internal/DEVELOPMENT_ROADMAP.md` §N4, §N6. From 77189663ff4ac4ca39b7f94203c099cca1ab7c74 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 22:29:11 -0400 Subject: [PATCH 11/14] Add BatchItemStatus.MAX_ATTEMPTED per-record terminal status The first prkit-synthesized per-record status beyond the provider set: a record that has exhausted MAX_ATTEMPTS total submissions (whole-minibatch + record-level retries) is given up on as MAX_ATTEMPTED, distinct from a transient ERRORED so a downstream scorer can tell "we gave up" from "errored". Mirrors the synthetic-ERRORED precedent in the batch leaf's correlation reader. Co-Authored-By: Claude Opus 4.8 --- src/prkit/core/model_clients/batch_types.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/prkit/core/model_clients/batch_types.py b/src/prkit/core/model_clients/batch_types.py index f95bbf0..f76d2cc 100644 --- a/src/prkit/core/model_clients/batch_types.py +++ b/src/prkit/core/model_clients/batch_types.py @@ -45,6 +45,12 @@ class BatchItemStatus(_StrEnum): ERRORED = "errored" EXPIRED = "expired" CANCELED = "canceled" + # prkit-synthesized (not a provider status): a record that has exhausted + # ``prkit.batch.MAX_ATTEMPTS`` total submissions (whole-minibatch + record-level + # retries) and is given up on as a terminal failure. Distinct from a transient + # ERRORED so a downstream scorer can tell "we gave up" from "errored". Mirrors + # the synthetic-ERRORED precedent in ``prkit.batch._iter_minibatch_results``. + MAX_ATTEMPTED = "max_attempted" @dataclass(frozen=True) From 4f8b4fbdc3de7a1996700b409a36f0d29ec6525d Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 22:31:05 -0400 Subject: [PATCH 12/14] Add Stage-4 record recovery: siphon-at-fetch + resubmit_failures drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fetch_batch now partitions each retrieved minibatch and siphons recoverable per-record failures (real ERRORED/EXPIRED/CANCELED + synthetic-missing, under MAX_ATTEMPTS=3) onto the source minibatch's ledger entry — pruning its id_map and decrementing num_requests in lockstep — writing succeeded-only to outputs/ and refreshing a derived failed-records-batch-input.jsonl; a record that has exhausted its submissions is instead rewritten as MAX_ATTEMPTED and kept in id_map so it consolidates terminally. resubmit_failed_minibatches is renamed resubmit_failures (module + BaseModelClient facade): it now also resubmits CANCELLED minibatches (reversing the Stage-3 exclusion) and bumps each whole-minibatch attempt, and additionally drains the failed-records accumulator into fresh retry minibatches (monotonic index, is_retry/retry_sources/record_attempts, minibatch_count++). The results manifest gains pending_failed_records and gates fully_consolidated on it being zero; the next-command prompts drop the CANCELLED dead-end and point at resubmit_failures, and consolidate emits the record-recovery hint when records are still pending. The leaf stays import-light: batch_types (BatchResult/BatchItemStatus.MAX_ATTEMPTED) is referenced only inside the already-lazy helpers, and _SIPHON_RECORD_STATUSES holds string values so the partition needs no enum at module load. Amends the Stage-1/2/3 surface and updates the affected tests. Co-Authored-By: Claude Opus 4.8 --- src/prkit/batch/__init__.py | 417 +++++++++++++++--- src/prkit/core/model_clients/base.py | 15 +- tests/prkit/batch/test_consolidate_results.py | 44 ++ tests/prkit/batch/test_fetch_batch.py | 36 +- tests/prkit/batch/test_import_isolation.py | 3 +- ...nibatches.py => test_resubmit_failures.py} | 58 +-- 6 files changed, 458 insertions(+), 115 deletions(-) rename tests/prkit/batch/{test_resubmit_minibatches.py => test_resubmit_failures.py} (82%) diff --git a/src/prkit/batch/__init__.py b/src/prkit/batch/__init__.py index 5d540c5..c671548 100644 --- a/src/prkit/batch/__init__.py +++ b/src/prkit/batch/__init__.py @@ -46,7 +46,7 @@ import tempfile import time from collections.abc import Callable, Iterator, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any @@ -65,7 +65,7 @@ "fetch_batch", "iter_batch_results", "consolidate_batch_results", - "resubmit_failed_minibatches", + "resubmit_failures", "batch_fetch_supported", "dumps_batch_jsonl", "write_batch_jsonl", @@ -120,10 +120,25 @@ # gate on this so re-scoring still works after finalize. _HAS_OUTPUT_STATUSES = frozenset({FETCHED, CONSOLIDATED}) -# Minibatches that :func:`resubmit_failed_minibatches` re-submits: terminal, not -# consolidatable, and NOT a deliberate CANCELLED (excluded by owner decision). This -# is exactly ``_COMPLETE_STATUSES − {FETCHED, CONSOLIDATED, CANCELLED}``. -_RESUBMIT_STATUSES = frozenset({FAILED, SUBMIT_ERROR, EXPIRED}) +# Minibatches that :func:`resubmit_failures` re-submits: every terminal, +# non-consolidatable status — including CANCELLED (Stage 4 D1 reverses the Stage-3 +# CANCELLED exclusion, so a cancelled job is re-driven into the loop instead of +# being a dead-end). This is exactly ``_COMPLETE_STATUSES − {FETCHED, CONSOLIDATED}``. +_RESUBMIT_STATUSES = frozenset({FAILED, SUBMIT_ERROR, EXPIRED, CANCELLED}) + +# Total submissions a single record gets — spanning BOTH whole-minibatch retries +# (each bumps the minibatch ``attempt``) and record-level retries (carried in +# ``record_attempts``) — before it is given up on. A record's submission count is +# ``record_attempts.get(cid, 0) + minibatch_attempt``; it is siphoned for retry +# while that is ``< MAX_ATTEMPTS`` and rewritten as MAX_ATTEMPTED once it reaches it. +# A fixed module constant by design (Stage 4 D3), never a per-call argument. +MAX_ATTEMPTS = 3 + +# Per-record ``BatchItemStatus`` *string values* (not enum members — so the siphon +# partition needs no ``batch_types`` import at module load, preserving leaf import +# discipline) that mark a recoverable per-record failure. ``MAX_ATTEMPTED`` is +# deliberately absent: an exhausted record is terminal and never re-siphoned. +_SIPHON_RECORD_STATUSES = frozenset({"errored", "expired", "canceled"}) # Providers with a full batch fetch lifecycle (poll + retrieve). Gemini's # provider string is "google" (not "gemini"); xAI / DeepSeek / Dashscope / Ollama @@ -174,7 +189,22 @@ class BatchSubmission: "counts": dict[str, int], # last poll's request_counts "endpoint": str | None, # audit-only (provider parity) "completion_window": str | None, # audit-only (provider parity) + # ---- Stage-4 additive keys (round-trip for free via dict(mb)) ---- + "attempt": int, # whole-minibatch submission count + # (submit => 1; whole-minibatch resubmit => +1) + "failed_records": list[dict], # pending siphoned per-record failures, each + # {custom_id, problem_id, error, attempt}; + # present only when records have been siphoned + # ---- on RETRY minibatches only (built by resubmit_failures' drain) ---- + "is_retry": bool, # True for a record-drain minibatch + "retry_sources": list[int], # source minibatch indices pooled into it + "record_attempts": dict[str, int], # per-record prior submission count + # (custom_id -> count), carried forward } + + Invariants: ``num_requests == len(id_map)`` on every minibatch (the siphon + decrements both in lockstep); a record's submission count is + ``record_attempts.get(custom_id, 0) + attempt``. """ # ---- batch-level (stored once; never repeated per minibatch) ---- @@ -320,7 +350,7 @@ class BatchFetchUnsupportedError(BatchInputError): class BatchNotTerminalError(BatchInputError): """The ledger is not terminal, so a terminal-gated finalize step refuses. - Raised **up front** by :func:`resubmit_failed_minibatches` when some minibatch + Raised **up front** by :func:`resubmit_failures` when some minibatch is still ``SUBMITTED`` / ``RUNNING`` / ``COMPLETED`` / ``FETCH_ERROR``. Run :func:`fetch_batch` first to drive every minibatch terminal (which also resolves any transient ``FETCH_ERROR`` by re-downloading), then resubmit the failures. @@ -562,6 +592,7 @@ def submit_batch_physics_reasoning( "counts": {}, "endpoint": None, "completion_window": None, + "attempt": 1, } ) @@ -694,7 +725,11 @@ def _run_fetch_pass( # EXPIRED persists only when it carried a partial subset. if results or st.state == batch_state.COMPLETED: output_path = outputs_dir / f"minibatch_{index:04d}.jsonl" - _write_results(output_path, results) + # Stage 4: partition + siphon rides this one download pass — write + # succeeded-only (+ exhausted MAX_ATTEMPTED) to outputs/, siphon the + # recoverable failures onto the ledger (pruning id_map/num_requests + # in lockstep) and refresh the derived accumulator. + _siphon_minibatch(sub, mb, results, output_path) sub.set_status( index, FETCHED, @@ -751,6 +786,134 @@ def _write_results(path: Path, results: Sequence[BatchResult]) -> None: path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8") +def _partition_results( + results: Sequence[BatchResult], + id_map: dict[str, str], +) -> tuple[list[BatchResult], list[BatchResult], int]: + """Partition one retrieved minibatch into ``(succeeded, failed, uncorrelated)``. + + Pure, in-memory (no I/O), and a structural mirror of the correlation in + :func:`_iter_minibatch_results` (synthetic-ERRORED completeness + uncorrelated + counting). For each ``custom_id`` in *id_map* (input order): the correlated + result is "succeeded" unless its status is a recoverable per-record failure + (:data:`_SIPHON_RECORD_STATUSES`); an id the provider never returned is + synthesized as an ERRORED "failed" record. Any returned id absent from *id_map* + is an uncorrelated extra — counted, not partitioned. ``batch_types`` is imported + lazily here to keep the leaf light (import discipline, §10). + """ + from prkit.core.model_clients.batch_types import BatchItemStatus, BatchResult + + results_by_cid: dict[str, BatchResult] = {} + uncorrelated = 0 + for r in results: + if r.custom_id in id_map: + results_by_cid[r.custom_id] = r + else: + uncorrelated += 1 + + succeeded: list[BatchResult] = [] + failed: list[BatchResult] = [] + for cid in id_map: + correlated = results_by_cid.get(cid) + if correlated is None: + correlated = BatchResult( + custom_id=cid, + status=BatchItemStatus.ERRORED, + error="No result returned by the provider for this request.", + ) + if str(correlated.status) in _SIPHON_RECORD_STATUSES: + failed.append(correlated) + else: + succeeded.append(correlated) + return succeeded, failed, uncorrelated + + +def _siphon_minibatch( + sub: BatchSubmission, + mb: dict[str, Any], + results: Sequence[BatchResult], + output_path: Path, +) -> None: + """Partition a retrieved minibatch, siphon recoverable failures, write outputs. + + The Stage-4 record-recovery step that rides :func:`fetch_batch`'s single + poll-and-download pass (it adds no new sweep). Partition *results* in memory + (:func:`_partition_results`); write the succeeded records — plus any *exhausted* + failure rewritten as :class:`BatchItemStatus.MAX_ATTEMPTED` — to *output_path* + (succeeded-only on the wire). Each failure under the :data:`MAX_ATTEMPTS` bound + is **siphoned**: appended to ``mb["failed_records"]`` and pruned from + ``mb["id_map"]`` with ``mb["num_requests"]`` decremented in lockstep (so + ``num_requests == len(id_map)`` holds). A record whose submission count + (``record_attempts.get(cid, 0) + attempt``) has reached ``MAX_ATTEMPTS`` is NOT + re-siphoned — it is rewritten as MAX_ATTEMPTED and kept in ``id_map`` so it + consolidates terminally (distinct from a transient ERRORED). Finally refresh the + derived run-level accumulator (:func:`_rewrite_failed_records_input`); the ledger + ``failed_records`` key is the crash-safe source of truth. + """ + from prkit.core.model_clients.batch_types import BatchItemStatus + + succeeded, failed, uncorrelated = _partition_results(results, mb["id_map"]) + if uncorrelated: + mb["uncorrelated_count"] = uncorrelated + + output_records = list(succeeded) + attempt = mb.get("attempt", 1) + record_attempts: dict[str, int] = mb.get("record_attempts") or {} + for r in failed: + cid = r.custom_id + submissions = record_attempts.get(cid, 0) + attempt + if submissions < MAX_ATTEMPTS: + mb.setdefault("failed_records", []).append( + { + "custom_id": cid, + "problem_id": mb["id_map"][cid], + "error": r.error, + "attempt": submissions, + } + ) + del mb["id_map"][cid] + mb["num_requests"] -= 1 + else: + output_records.append(replace(r, status=BatchItemStatus.MAX_ATTEMPTED)) + + _write_results(output_path, output_records) + _rewrite_failed_records_input(sub) + + +def _rewrite_failed_records_input(sub: BatchSubmission) -> None: + """Rewrite ``/failed-records-batch-input.jsonl`` from the ledger (derived). + + The authoritative pending set is each minibatch's ``failed_records`` ledger key; + this file is a **derived** atomic full-rewrite of the provider request lines for + every pending failed record, re-read from each source minibatch's never-pruned + ``inputs/`` file and correlated by the provider id field (``key`` for Gemini, + else ``custom_id``) — never a positional zip, which would mis-correlate on order + drift. A full-rewrite (not append) keeps it crash-safe and idempotent: a re-run + can never double-count, and resubmit's correctness depends only on the ledger + plus the persisted ``inputs/`` files. When nothing is pending the stale file is + removed, so the artifact appears exactly when record recovery is in play. + """ + id_field = "key" if sub.provider in _KEY_ID_PROVIDERS else "custom_id" + path = Path(sub.run_dir) / "failed-records-batch-input.jsonl" + lines: list[str] = [] + for mb in sub.minibatches: + pending = mb.get("failed_records") + if not pending: + continue + lines_by_cid = { + json.loads(line)[id_field]: line + for line in _read_jsonl(mb["input_file_path"]) + } + for entry in pending: + line = lines_by_cid.get(entry["custom_id"]) + if line is not None: + lines.append(line) + if lines: + _atomic_write_text(path, "\n".join(lines) + "\n") + elif path.exists(): + path.unlink() + + def iter_batch_results( submission: BatchSubmission | str | Path, ) -> Iterator[tuple[str, BatchResult]]: @@ -1074,9 +1237,19 @@ def _write_results_manifest( ``fully_consolidated`` are ledger-derived (cumulative); ``results_written`` / ``status_counts`` / ``uncorrelated_total`` reflect *this call's* newly consolidated minibatches (per the design's streaming counters). A failed or - CANCELLED minibatch keeps ``fully_consolidated`` false (see Known limitations). + CANCELLED minibatch keeps ``fully_consolidated`` false (it is not yet consolidated). + + ``pending_failed_records`` (Stage 4 D4) is the run-wide count of siphoned records + still awaiting recovery (``Σ len(mb["failed_records"])``); ``fully_consolidated`` + is gated on it being zero, so a ``while not fully_consolidated`` consumer keeps + looping until every record is recovered or exhausted. An exhausted MAX_ATTEMPTED + record is *not* in ``failed_records``, so a run with only permanent failures is + legitimately "done." """ consolidated = sub.status_counts().get(CONSOLIDATED, 0) + pending_failed_records = sum( + len(mb.get("failed_records") or []) for mb in sub.minibatches + ) manifest = { "run_name": sub.run_name, "provider": sub.provider, @@ -1086,7 +1259,10 @@ def _write_results_manifest( "status_counts": status_counts, "minibatches_consolidated": consolidated, "minibatches_total": sub.minibatch_count, - "fully_consolidated": consolidated == sub.minibatch_count, + "pending_failed_records": pending_failed_records, + "fully_consolidated": ( + consolidated == sub.minibatch_count and pending_failed_records == 0 + ), "uncorrelated_total": uncorrelated_total, "consolidated_at": _now_iso(), } @@ -1096,21 +1272,40 @@ def _write_results_manifest( ) -def resubmit_failed_minibatches( +def resubmit_failures( client: Any, submission: BatchSubmission | str | Path, ) -> BatchSubmission: - """Re-submit each FAILED / SUBMIT_ERROR / EXPIRED minibatch (NOT CANCELLED). - - Re-reads each target minibatch's persisted ``inputs/minibatch_XXXX.jsonl`` and - calls ``client.submit_batch`` with the run's merged metadata, then — **only after - the submit returns a new ``batch_id``** — resets that ledger entry (status → - ``SUBMITTED``, new ``batch_id``, cleared ``output_path`` / ``fetched_at`` / - ``counts`` / ``error``) in memory and on disk. A single minibatch's submit - failure is recorded as ``SUBMIT_ERROR`` (``batch_id=""``, ``error`` set) and the - loop continues. CANCELLED minibatches are left untouched (a deliberate cancel is - not auto-resubmitted — a known dead-end). Returns the updated ledger; the - consumer then re-runs :func:`fetch_batch` on the new jobs. + """Re-drive a terminal batch run's failures: whole minibatches **and** records. + + Does both bounded-recovery jobs in one call (Stage 4 D5): + + (a) **Whole-minibatch retries.** Re-submit each FAILED / SUBMIT_ERROR / EXPIRED / + **CANCELLED** minibatch in place — re-reading its persisted + ``inputs/minibatch_XXXX.jsonl`` and calling ``client.submit_batch`` with the run's + merged metadata, then — **only after** the submit returns a new ``batch_id`` — + resetting that ledger entry (status → ``SUBMITTED``, new ``batch_id``, cleared + ``output_path`` / ``fetched_at`` / ``counts`` / ``error``, ``attempt`` bumped) in + memory and on disk. A single minibatch's submit failure is recorded as + ``SUBMIT_ERROR`` (``batch_id=""``, ``error`` set) and the loop continues. CANCELLED + is now resubmitted like any other failure (Stage 4 D1 reverses the Stage-3 + exclusion), so a cancelled job is no longer a dead-end. + + (b) **Record drain.** Drain the run-level failed-records accumulator into fresh + minibatch(es): each chunk of ``<= minibatch_size`` pending siphoned records is + re-read from its source ``inputs/`` file (correlated by the provider id field, not + a positional zip), validated, written to a fresh ``inputs/minibatch_XXXX.jsonl``, + submitted, and appended to the ledger as a SUBMITTED retry minibatch (fresh + monotonic ``index``, ``is_retry=True``, ``attempt=1``, ``retry_sources``, and a + ``record_attempts`` carrying each record's prior submission count). ``minibatch_count`` + is bumped per appended retry minibatch; the consumed entries are removed from their + source ``failed_records`` in the same atomic save; the derived accumulator file is + refreshed last. + + Returns the updated ledger; the consumer then re-runs :func:`fetch_batch` on the + new jobs (and re-consolidates). The recovery loop terminates because every record + either succeeds or reaches :data:`MAX_ATTEMPTS` and is consolidated as a terminal + MAX_ATTEMPTED result, so the pending count strictly decreases to zero. Raises: BatchFetchUnsupportedError: up front, for a provider with no batch lifecycle. @@ -1136,16 +1331,19 @@ def resubmit_failed_minibatches( ) targets = [mb for mb in sub.minibatches if mb["status"] in _RESUBMIT_STATUSES] - if not targets: + has_pending = any(mb.get("failed_records") for mb in sub.minibatches) + if not targets and not has_pending: _log_next_after_resubmit(sub, resubmitted=0, still_failed=0) return sub # The ledger stores the un-merged user metadata + display_name separately, so - # reconstruct the provider-side merged metadata exactly as submit did (:486). + # reconstruct the provider-side merged metadata exactly as submit did. merged_metadata = {**sub.metadata, "display_name": sub.display_name} outputs_dir = Path(sub.run_dir) / "outputs" resubmitted = 0 still_failed = 0 + + # (a) Whole-minibatch failures — now incl. CANCELLED (D1); each bumps ``attempt``. for mb in targets: index = mb["index"] old_status = mb["status"] @@ -1180,15 +1378,116 @@ def resubmit_failed_minibatches( output_path=None, fetched_at=None, counts={}, + attempt=mb.get("attempt", 1) + 1, ) resubmitted += 1 _log_transition(index, old_status, SUBMITTED) sub.save() # crash-safe per minibatch + # (b) Drain the run-level failed-records accumulator into fresh minibatch(es). + resub_drain, drain_failed = _drain_failed_records(client, sub, merged_metadata) + resubmitted += resub_drain + still_failed += drain_failed + _log_next_after_resubmit(sub, resubmitted=resubmitted, still_failed=still_failed) return sub +def _drain_failed_records( + client: Any, + sub: BatchSubmission, + merged_metadata: dict[str, str], +) -> tuple[int, int]: + """Drain pending siphoned records into fresh retry minibatches; return Δ counts. + + Reuse, not re-derive: each pending record's provider request line is re-read from + its source minibatch's never-pruned ``inputs/`` file (correlated by the id field, + not a positional zip — blocker #1). Each chunk of ``<= minibatch_size`` records + becomes a fresh SUBMITTED retry minibatch appended to the ledger with a monotonic + ``index``, rebuilt ``id_map`` / ``record_attempts``, and ``minibatch_count`` bumped; + the drained entries are consumed from their source ``failed_records`` in the same + crash-safe save. Returns ``(submitted_count, submit_error_count)``. + """ + pending = [ + (mb, entry) + for mb in sub.minibatches + for entry in list(mb.get("failed_records") or []) + ] + if not pending: + return 0, 0 + + id_field = "key" if sub.provider in _KEY_ID_PROVIDERS else "custom_id" + inputs_dir = Path(sub.run_dir) / "inputs" + # Cache each source minibatch's never-pruned input lines, keyed by wire id. + source_lines: dict[int, dict[str, str]] = {} + for mb in sub.minibatches: + if mb.get("failed_records"): + source_lines[mb["index"]] = { + json.loads(line)[id_field]: line + for line in _read_jsonl(mb["input_file_path"]) + } + + resubmitted = 0 + still_failed = 0 + for chunk in _chunked(pending, sub.minibatch_size): + new_index = max(m["index"] for m in sub.minibatches) + 1 + requests: list[dict[str, Any]] = [] + id_map: dict[str, str] = {} + record_attempts: dict[str, int] = {} + retry_sources: set[int] = set() + for src_mb, entry in chunk: + cid = entry["custom_id"] + requests.append(json.loads(source_lines[src_mb["index"]][cid])) + id_map[cid] = entry["problem_id"] + record_attempts[cid] = entry["attempt"] + retry_sources.add(src_mb["index"]) + + validate_batch_requests(requests, provider=sub.provider) + retry_input_path = inputs_dir / f"minibatch_{new_index:04d}.jsonl" + write_batch_jsonl(requests, retry_input_path) + + batch_id = "" + error: str | None = None + try: + batch_id = client.submit_batch(requests, metadata=merged_metadata) + except Exception as exc: # noqa: BLE001 - recorded on the ledger; loop goes on + error = f"{type(exc).__name__}: {exc}" + + sub.minibatches.append( + { + "index": new_index, + "batch_id": batch_id, + "status": SUBMITTED if batch_id else SUBMIT_ERROR, + "input_file_path": str(retry_input_path), + "num_requests": len(requests), + "id_map": id_map, + "submitted_at": _now_iso(), + "error": error, + "output_path": None, + "fetched_at": None, + "counts": {}, + "endpoint": None, + "completion_window": None, + "attempt": 1, + "is_retry": True, + "retry_sources": sorted(retry_sources), + "record_attempts": record_attempts, + } + ) + sub.minibatch_count += 1 # counts submit + retry minibatch jobs + if batch_id: + resubmitted += 1 + else: + still_failed += 1 + # Consume the drained entries from their source ledger lists in the same save. + for src_mb, entry in chunk: + src_mb["failed_records"].remove(entry) + sub.save() # crash-safe per chunk + + _rewrite_failed_records_input(sub) # now empty / shrunk (derived) + return resubmitted, still_failed + + # --------------------------------------------------------------------------- # # Next-command guidance (one INFO line telling the human what to run next) # # --------------------------------------------------------------------------- # @@ -1231,43 +1530,24 @@ def _log_next_after_fetch(sub: BatchSubmission) -> None: ) return - # Terminal, some failed. + # Terminal, some failed. Every terminal-not-good status (incl. CANCELLED, D1) is + # resubmittable, so there is no dead-end branch. failed_total = n - succeeded - resubmittable = sum(counts.get(s, 0) for s in _RESUBMIT_STATUSES) - if resubmittable: - if succeeded: - _logger.info( - '%d minibatches failed. Next: client.resubmit_failed_minibatches("%s"),' - " then fetch again. (%d already succeeded — consolidate_batch_results" - " can capture them now.)", - failed_total, - sub.run_dir, - succeeded, - ) - else: - _logger.info( - '%d minibatches failed. Next: client.resubmit_failed_minibatches("%s"),' - " then fetch again.", - failed_total, - sub.run_dir, - ) - return - - # All failures are CANCELLED dead-ends (none resubmittable; see Known limitations). if succeeded: _logger.info( - "%d minibatches were CANCELLED — a dead-end (not auto-resubmitted; see " - 'Known limitations). consolidate_batch_results("%s") can still capture ' - "the %d succeeded.", + '%d minibatches failed. Next: client.resubmit_failures("%s"),' + " then fetch again. (%d already succeeded — consolidate_batch_results" + " can capture them now.)", failed_total, sub.run_dir, succeeded, ) else: _logger.info( - "%d minibatches were CANCELLED — a dead-end (not auto-resubmitted; see " - "Known limitations); nothing left to fetch or consolidate.", + '%d minibatches failed. Next: client.resubmit_failures("%s"),' + " then fetch again.", failed_total, + sub.run_dir, ) @@ -1276,21 +1556,11 @@ def _log_next_after_resubmit( ) -> None: """After resubmit: the K-resubmitted / M-failed summary + the fetch-next hint.""" if resubmitted == 0 and still_failed == 0: - cancelled = sub.status_counts().get(CANCELLED, 0) - if cancelled: - _logger.info( - "Nothing to resubmit: %d CANCELLED minibatch(es) are a dead-end (not " - "auto-resubmitted) and the rest already succeeded. Run " - 'consolidate_batch_results("%s") to capture the succeeded subset.', - cancelled, - sub.run_dir, - ) - else: - _logger.info( - "Nothing to resubmit — all minibatches already succeeded. " - 'Next: consolidate_batch_results("%s").', - sub.run_dir, - ) + _logger.info( + "Nothing to resubmit — all minibatches already succeeded and no records " + 'are pending. Next: consolidate_batch_results("%s").', + sub.run_dir, + ) return _logger.info( "Resubmitted %d minibatches (%d still failed). " @@ -1302,11 +1572,22 @@ def _log_next_after_resubmit( def _log_next_after_consolidate(sub: BatchSubmission, results_dirname: str) -> None: - """After consolidate: 'done' when fully consolidated, else the resume hint.""" + """After consolidate: 'done', the record-recovery hint (D4), or the resume hint.""" consolidated = sub.status_counts().get(CONSOLIDATED, 0) n = sub.minibatch_count - if consolidated == n: + pending_failed = sum(len(mb.get("failed_records") or []) for mb in sub.minibatches) + if consolidated == n and pending_failed == 0: _logger.info("Done — results in %s/%s/.", sub.run_dir, results_dirname) + elif pending_failed: + # D4: minibatch-completeness alone is blind to pending siphoned records, so a + # `while not fully_consolidated` consumer would stop without record-retry. + _logger.info( + "Done consolidating %d minibatches, but %d records still failed — run " + 'client.resubmit_failures("%s"), then fetch + consolidate again.', + consolidated, + pending_failed, + sub.run_dir, + ) else: _logger.info( "Consolidated %d/%d minibatches; the rest are not yet succeeded — " diff --git a/src/prkit/core/model_clients/base.py b/src/prkit/core/model_clients/base.py index ce7806a..e714844 100644 --- a/src/prkit/core/model_clients/base.py +++ b/src/prkit/core/model_clients/base.py @@ -416,24 +416,25 @@ def fetch_batch_physics_reasoning( return fetch_batch(self, run_dir_or_submission, **kwargs) - def resubmit_failed_minibatches( + def resubmit_failures( self, run_dir_or_submission: BatchSubmission | str | Path, **kwargs: Any, ) -> BatchSubmission: - """Re-submit a terminal batch run's failed minibatches; return the ledger. + """Re-drive a terminal batch run's failures (minibatches + records); return it. Thin one-line facade mirroring :meth:`fetch_batch_physics_reasoning`: lazily imports :mod:`prkit.batch` and delegates to - :func:`prkit.batch.resubmit_failed_minibatches`. Re-submits each - FAILED / SUBMIT_ERROR / EXPIRED minibatch (not CANCELLED) by re-reading its - persisted ``inputs/`` file; the run must be terminal first (run + :func:`prkit.batch.resubmit_failures`. Re-submits each + FAILED / SUBMIT_ERROR / EXPIRED / CANCELLED minibatch in place (re-reading its + persisted ``inputs/`` file) *and* drains the run's siphoned failed-records + accumulator into fresh minibatches; the run must be terminal first (run :meth:`fetch_batch_physics_reasoning` until it is). Returns the updated :class:`~prkit.batch.BatchSubmission`; fetch the new jobs next. """ - from prkit.batch import resubmit_failed_minibatches + from prkit.batch import resubmit_failures - return resubmit_failed_minibatches(self, run_dir_or_submission, **kwargs) + return resubmit_failures(self, run_dir_or_submission, **kwargs) def _build_batch_request( self, diff --git a/tests/prkit/batch/test_consolidate_results.py b/tests/prkit/batch/test_consolidate_results.py index 4dd071a..302587f 100644 --- a/tests/prkit/batch/test_consolidate_results.py +++ b/tests/prkit/batch/test_consolidate_results.py @@ -396,6 +396,50 @@ def test_empty_ledger_raises(self, tmp_path): consolidate_batch_results(str(tmp_path / "r")) +class TestPendingFailedRecords: + """Stage 4 D4: a siphoned record keeps the manifest from reading 'done'.""" + + def _run_with_siphoned_record(self, tmp_path): + run_dir = _submit_dataset( + tmp_path, _dataset(2), minibatch_size=2, batch_ids=["b0"] + ) + cids = list(BatchSubmission.load(run_dir).minibatches[0]["id_map"]) + client = _FetchClient( + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.ERRORED, error="boom"), + BatchResult(cids[1], BatchItemStatus.SUCCEEDED, text="A1"), + ] + }, + ) + fetch_batch(client, run_dir, progress=False) + return run_dir + + def test_manifest_gates_fully_consolidated_on_pending_records(self, tmp_path): + run_dir = self._run_with_siphoned_record(tmp_path) + consolidate_batch_results(run_dir) + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + # All minibatches consolidated, but a record is still pending recovery. + assert manifest["minibatches_consolidated"] == manifest["minibatches_total"] + assert manifest["pending_failed_records"] == 1 + assert manifest["fully_consolidated"] is False + # The siphoned problem has no result file yet (siphoned-not-resubmitted). + assert not (Path(run_dir) / "results" / "p0.json").exists() + assert (Path(run_dir) / "results" / "p1.json").is_file() + + def test_post_consolidate_hint_points_at_resubmit_failures(self, tmp_path, caplog): + run_dir = self._run_with_siphoned_record(tmp_path) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + consolidate_batch_results(run_dir) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any( + "records still failed" in m and "resubmit_failures" in m for m in msgs + ) + assert not any(m.startswith("Done — results in") for m in msgs) + + class TestNextCommand: def test_done_line_when_fully_consolidated(self, tmp_path, caplog): run_dir = _build_run(tmp_path, statuses=["fetched"]) diff --git a/tests/prkit/batch/test_fetch_batch.py b/tests/prkit/batch/test_fetch_batch.py index 1dac98c..e2a38e2 100644 --- a/tests/prkit/batch/test_fetch_batch.py +++ b/tests/prkit/batch/test_fetch_batch.py @@ -235,7 +235,9 @@ def test_in_progress_leaves_running_writes_no_output(self, tmp_path): class TestExpired: - def test_expired_partial_subset_persisted_with_synthetic_missing(self, tmp_path): + def test_expired_partial_subset_siphons_missing_record(self, tmp_path): + # Stage 4: the missing record of an EXPIRED partial is now SIPHONED for retry + # (not left as a synthetic ERRORED read), so the minibatch reads all-success. run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) cids = _cids(BatchSubmission.load(run_dir)) client = _FetchClient( @@ -246,10 +248,17 @@ def test_expired_partial_subset_persisted_with_synthetic_missing(self, tmp_path) }, ) sub = fetch_batch(client, run_dir, progress=False) - assert sub.minibatches[0]["status"] == FETCHED # partial subset -> fetched + mb = sub.minibatches[0] + assert mb["status"] == FETCHED # partial subset -> fetched + # The missing record was siphoned: pruned from id_map, num_requests in lockstep. + assert mb["num_requests"] == len(mb["id_map"]) == 1 + assert [e["custom_id"] for e in mb["failed_records"]] == [cids[1]] + assert mb["failed_records"][0]["problem_id"] == "p1" + assert mb["failed_records"][0]["attempt"] == 1 + # outputs/ holds the succeeded record only; the reader no longer yields p1. pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0"} assert pairs["p0"].succeeded and pairs["p0"].text == "A0" - assert pairs["p1"].status == BatchItemStatus.ERRORED # synthetic completeness def test_expired_with_nothing_marks_expired_terminal(self, tmp_path): run_dir = _submit_run(tmp_path, n=1, minibatch_size=1, batch_ids=["b0"]) @@ -348,7 +357,9 @@ def test_surrogate_ids_recover_problem_id(self, tmp_path): class TestCompleteness: - def test_one_result_per_id_in_order_extra_dropped_and_counted(self, tmp_path): + def test_missing_siphoned_extra_dropped_and_counted(self, tmp_path): + # Stage 4: a missing id is siphoned (not synthesized on read); an extra + # (uncorrelated) id is still dropped and counted. run_dir = _submit_run(tmp_path, n=3, minibatch_size=3, batch_ids=["b0"]) cids = _cids(BatchSubmission.load(run_dir)) # [p0, p1, p2] in input order client = _FetchClient( @@ -357,26 +368,23 @@ def test_one_result_per_id_in_order_extra_dropped_and_counted(self, tmp_path): results={ "b0": [ BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A0"), - # cids[1] omitted -> synthetic ERRORED + # cids[1] omitted -> siphoned for retry (no longer in id_map) BatchResult(cids[2], BatchItemStatus.SUCCEEDED, text="A2"), BatchResult("ghost", BatchItemStatus.SUCCEEDED, text="X"), # extra ] }, ) sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + # The missing id was siphoned; the two succeeded stay; lockstep invariant. + assert mb["num_requests"] == len(mb["id_map"]) == 2 + assert [e["problem_id"] for e in mb["failed_records"]] == ["p1"] pairs = list( iter_batch_results(sub) ) # pass the object to read uncorrelated_count - assert [pid for pid, _ in pairs] == [ - "p0", - "p1", - "p2", - ] # exactly one each, in order + assert [pid for pid, _ in pairs] == ["p0", "p2"] # p1 siphoned out, in order by = dict(pairs) assert by["p0"].succeeded and by["p2"].succeeded - assert ( - by["p1"].status == BatchItemStatus.ERRORED - ) # synthetic for the missing id assert sub.minibatches[0]["uncorrelated_count"] == 1 # ghost dropped + counted @@ -603,6 +611,6 @@ def test_some_failed_points_at_resubmit(self, tmp_path, caplog): caplog.clear() fetch_batch(client, run_dir, progress=True) assert any( - "1 minibatches failed" in m and "resubmit_failed_minibatches" in m + "1 minibatches failed" in m and "resubmit_failures" in m for m in self._msgs(caplog) ) diff --git a/tests/prkit/batch/test_import_isolation.py b/tests/prkit/batch/test_import_isolation.py index 4ffb2a3..56a9255 100644 --- a/tests/prkit/batch/test_import_isolation.py +++ b/tests/prkit/batch/test_import_isolation.py @@ -45,8 +45,9 @@ def test_batch_import_does_not_pull_heavy_deps(): fetch_batch, iter_batch_results, consolidate_batch_results, - resubmit_failed_minibatches, + resubmit_failures, batch_fetch_supported, + MAX_ATTEMPTS, ) forbidden = {_FORBIDDEN!r} diff --git a/tests/prkit/batch/test_resubmit_minibatches.py b/tests/prkit/batch/test_resubmit_failures.py similarity index 82% rename from tests/prkit/batch/test_resubmit_minibatches.py rename to tests/prkit/batch/test_resubmit_failures.py index 84b6675..43bbb60 100644 --- a/tests/prkit/batch/test_resubmit_minibatches.py +++ b/tests/prkit/batch/test_resubmit_failures.py @@ -1,12 +1,14 @@ -"""Tests for the finalize half: ``resubmit_failed_minibatches`` (Stage 3). +"""Tests for the finalize half: ``resubmit_failures`` whole-minibatch path (Stage 3/4). Fully offline. Runs are built with an SDK-patched submit client then driven to terminal statuses by the duck-typed ``_FetchClient``; resubmit is then driven by a tiny duck-typed ``_Client`` whose ``submit_batch`` is a ``MagicMock``. Covers the happy path (requests re-read from the persisted input file + merged -metadata, entries reset to SUBMITTED), CANCELLED exclusion + no-op, the terminal -and capability preconditions, per-item failure + continue, submit-first ordering, -the client facade, and the next-command prompt. +metadata, entries reset to SUBMITTED, ``attempt`` bumped), CANCELLED now +resubmitted (Stage 4 D1) + no-op, the terminal and capability preconditions, +per-item failure + continue, submit-first ordering, the client facade, and the +next-command prompt. The record-drain half of ``resubmit_failures`` is covered in +``test_resubmit_records.py``. """ from __future__ import annotations @@ -17,7 +19,6 @@ import pytest from prkit.batch import ( - CANCELLED, FETCHED, SUBMIT_ERROR, SUBMITTED, @@ -25,7 +26,7 @@ BatchNotTerminalError, BatchSubmission, fetch_batch, - resubmit_failed_minibatches, + resubmit_failures, submit_batch_physics_reasoning, ) from prkit.core.domain import PhysicsDataset, PhysicsProblem @@ -140,7 +141,7 @@ def test_resubmits_failed_and_expired_resetting_each_entry(self, tmp_path): client = _Client() client.submit_batch = MagicMock(side_effect=["new1", "new2"]) - sub = resubmit_failed_minibatches(client, run_dir) + sub = resubmit_failures(client, run_dir) assert client.submit_batch.call_count == 2 # FAILED + EXPIRED only for call in client.submit_batch.call_args_list: @@ -158,32 +159,39 @@ def test_resubmits_failed_and_expired_resetting_each_entry(self, tmp_path): assert mb["fetched_at"] is None assert mb["counts"] == {} assert sub.minibatches[0]["status"] == FETCHED # the succeeded one is untouched + # Whole-minibatch resubmit bumps each entry's attempt (submit => 1, +1 here). + assert mb1["attempt"] == 2 and mb2["attempt"] == 2 # Persisted to disk. assert BatchSubmission.load(run_dir).minibatches[1]["batch_id"] == "new1" -class TestCancelledExcluded: - def test_cancelled_minibatch_not_resubmitted(self, tmp_path): +class TestCancelledResubmitted: + """Stage 4 D1: CANCELLED is now resubmitted like any other failure (no dead-end).""" + + def test_cancelled_minibatch_is_resubmitted(self, tmp_path): run_dir = _build_terminal_run( tmp_path, statuses=["fetched", "cancelled", "failed"] ) client = _Client() - client.submit_batch = MagicMock(side_effect=["new2"]) - sub = resubmit_failed_minibatches(client, run_dir) - assert client.submit_batch.call_count == 1 # only the FAILED one - assert sub.minibatches[1]["status"] == CANCELLED # untouched dead-end + client.submit_batch = MagicMock(side_effect=["new1", "new2"]) + sub = resubmit_failures(client, run_dir) + assert client.submit_batch.call_count == 2 # CANCELLED + FAILED both + assert sub.minibatches[1]["status"] == SUBMITTED # cancelled re-driven + assert sub.minibatches[1]["batch_id"] == "new1" + assert sub.minibatches[1]["attempt"] == 2 assert sub.minibatches[2]["status"] == SUBMITTED - def test_only_cancelled_is_logged_noop(self, tmp_path, caplog): + def test_only_cancelled_is_resubmitted_not_noop(self, tmp_path, caplog): run_dir = _build_terminal_run(tmp_path, statuses=["fetched", "cancelled"]) client = _Client() - client.submit_batch = MagicMock() + client.submit_batch = MagicMock(side_effect=["new1"]) with caplog.at_level(logging.INFO, logger="prkit.batch"): caplog.clear() - resubmit_failed_minibatches(client, run_dir) - client.submit_batch.assert_not_called() + sub = resubmit_failures(client, run_dir) + client.submit_batch.assert_called_once() + assert sub.minibatches[1]["status"] == SUBMITTED msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] - assert any("Nothing to resubmit" in m and "CANCELLED" in m for m in msgs) + assert any("Resubmitted 1 minibatches (0 still failed)" in m for m in msgs) class TestPreconditions: @@ -192,7 +200,7 @@ def test_non_terminal_ledger_raises_before_any_submit(self, tmp_path): client = _Client() client.submit_batch = MagicMock() with pytest.raises(BatchNotTerminalError): - resubmit_failed_minibatches(client, run_dir) + resubmit_failures(client, run_dir) client.submit_batch.assert_not_called() def test_non_batch_provider_raises_up_front(self, tmp_path): @@ -200,7 +208,7 @@ def test_non_batch_provider_raises_up_front(self, tmp_path): client = _Client("xai") client.submit_batch = MagicMock() with pytest.raises(BatchFetchUnsupportedError, match="xai"): - resubmit_failed_minibatches(client, run_dir) + resubmit_failures(client, run_dir) client.submit_batch.assert_not_called() @@ -213,7 +221,7 @@ def test_one_submit_raises_recorded_and_loop_continues(self, tmp_path, caplog): ) with caplog.at_level(logging.INFO, logger="prkit.batch"): caplog.clear() - sub = resubmit_failed_minibatches(client, run_dir) + sub = resubmit_failures(client, run_dir) assert ( sub.minibatches[0]["status"] == SUBMITTED @@ -233,7 +241,7 @@ def test_failed_submit_never_leaves_submitted_with_empty_id(self, tmp_path): run_dir = _build_terminal_run(tmp_path, statuses=["failed"]) client = _Client() client.submit_batch = MagicMock(side_effect=RuntimeError("boom")) - sub = resubmit_failed_minibatches(client, run_dir) + sub = resubmit_failures(client, run_dir) mb = sub.minibatches[0] # Submit-first / mutate-second: a raising submit must NOT yield SUBMITTED+"". assert not (mb["status"] == SUBMITTED and mb["batch_id"] == "") @@ -243,9 +251,9 @@ def test_failed_submit_never_leaves_submitted_with_empty_id(self, tmp_path): class TestFacade: def test_resubmit_facade_delegates_to_module(self): client = _openai_submit_client() - with patch("prkit.batch.resubmit_failed_minibatches") as mock_resub: + with patch("prkit.batch.resubmit_failures") as mock_resub: mock_resub.return_value = "LEDGER" - out = client.resubmit_failed_minibatches("run-dir") + out = client.resubmit_failures("run-dir") mock_resub.assert_called_once() args, _ = mock_resub.call_args assert args[0] is client and args[1] == "run-dir" @@ -259,7 +267,7 @@ def test_resubmit_logs_fetch_next(self, tmp_path, caplog): client.submit_batch = MagicMock(side_effect=["new0"]) with caplog.at_level(logging.INFO, logger="prkit.batch"): caplog.clear() - resubmit_failed_minibatches(client, run_dir) + resubmit_failures(client, run_dir) msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] assert any( m.startswith("Resubmitted 1 minibatches (0 still failed)") From fadc414665ab57b41c5b3e0713db3b6dd6c8ca56 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 22:32:47 -0400 Subject: [PATCH 13/14] Add Stage-4 record-recovery test suites (siphon + record drain) test_fetch_siphon.py: the siphon happy path (succeeded-only outputs; one failed_records entry; id_map pruned + num_requests in lockstep; derived accumulator), synthetic-missing siphon, re-fetch idempotency, no-failure no-artifact, and MAX_ATTEMPTS exhaustion (rewritten max_attempted, consolidated). test_resubmit_records.py: drain mechanics (fresh retry minibatch shape; minibatch_count bump; consumed accumulator; submit_batch fed the failed input line by wire id), a drain submit error keeping the record recoverable, the headline end-to-end recovery loop to fully_consolidated, the field-parse correlation regression (custom_id and Gemini key, with reordered input lines), the MAX_ATTEMPTS off-by-one spanning whole-minibatch + record retries, the terminal precondition, the ledger round-trip of the new minibatch keys, and the consolidate record-recovery hint. Co-Authored-By: Claude Opus 4.8 --- tests/prkit/batch/test_fetch_siphon.py | 238 +++++++++++++++ tests/prkit/batch/test_resubmit_records.py | 336 +++++++++++++++++++++ 2 files changed, 574 insertions(+) create mode 100644 tests/prkit/batch/test_fetch_siphon.py create mode 100644 tests/prkit/batch/test_resubmit_records.py diff --git a/tests/prkit/batch/test_fetch_siphon.py b/tests/prkit/batch/test_fetch_siphon.py new file mode 100644 index 0000000..7dc7983 --- /dev/null +++ b/tests/prkit/batch/test_fetch_siphon.py @@ -0,0 +1,238 @@ +"""Tests for the Stage-4 siphon inside ``fetch_batch`` (record-level recovery). + +Fully offline, mirroring ``tests/prkit/batch/test_fetch_batch.py``: runs are +submitted with an SDK-patched real client (realistic ``id_map`` + input files) +then polled/downloaded by a duck-typed ``_FetchClient`` whose +``retrieve_batch_results`` returns scripted ``BatchResult``s. A record-level +failure is injected either as a real ERRORED result for one cid or by omitting a +cid (the synthetic-missing path). Covers: the happy-path siphon (succeeded-only +on the wire; one ``failed_records`` entry; ``id_map`` pruned + ``num_requests`` +decremented in lockstep; the derived accumulator), synthetic-missing siphon, +siphon idempotency on a re-fetch, and the MAX_ATTEMPTS exhaustion path +(rewritten as ``max_attempted``, kept in ``id_map``, consolidated terminally). +""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +from prkit.batch import ( + FETCHED, + MAX_ATTEMPTS, + BatchSubmission, + consolidate_batch_results, + fetch_batch, + iter_batch_results, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem +from prkit.core.model_clients.batch_types import ( + BatchItemStatus, + BatchResult, + BatchState, + BatchStatus, +) + + +# --------------------------------------------------------------------------- # +# Offline builders (mirroring test_fetch_batch.py) # +# --------------------------------------------------------------------------- # +def _openai_submit_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _dataset(n: int): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +def _submit_run(tmp_path, *, n, minibatch_size, batch_ids, run_name="run"): + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=list(batch_ids)) + return submit_batch_physics_reasoning( + client, + _dataset(n), + output_dir=tmp_path, + run_name=run_name, + minibatch_size=minibatch_size, + ) + + +class _FetchClient: + def __init__(self, provider="openai", *, states=None, results=None): + self.provider = provider + self.model = "model-x" + self._states = states or {} + self._results = results or {} + self.poll_calls: list[str] = [] + self.retrieve_calls: list[str] = [] + + def poll_batch(self, batch_id: str) -> BatchStatus: + self.poll_calls.append(batch_id) + seq = self._states[batch_id] + state = seq.pop(0) if len(seq) > 1 else seq[0] + return BatchStatus( + batch_id=batch_id, + state=state, + provider=self.provider, + raw_status=str(state), + counts={"total": 1}, + ) + + def retrieve_batch_results(self, batch_id: str): + self.retrieve_calls.append(batch_id) + return iter(self._results.get(batch_id, [])) + + +def _cids(sub, index=0): + return list(sub.minibatches[index]["id_map"]) + + +# --------------------------------------------------------------------------- # +class TestSiphonHappyPath: + def test_errored_record_siphoned_outputs_succeeded_only(self, tmp_path): + run_dir = _submit_run(tmp_path, n=3, minibatch_size=3, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) # [p0, p1, p2] + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A0"), + BatchResult(cids[1], BatchItemStatus.ERRORED, error="rate limited"), + BatchResult(cids[2], BatchItemStatus.SUCCEEDED, text="A2"), + ] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + assert mb["status"] == FETCHED + + # One siphoned record, recorded on the ledger (attempt=1, the source error). + assert [e["custom_id"] for e in mb["failed_records"]] == [cids[1]] + entry = mb["failed_records"][0] + assert entry == { + "custom_id": cids[1], + "problem_id": "p1", + "error": "rate limited", + "attempt": 1, + } + # id_map pruned and num_requests decremented in lockstep (blocker #2). + assert cids[1] not in mb["id_map"] + assert mb["num_requests"] == len(mb["id_map"]) == 2 + + # outputs/ holds the succeeded records only (no errored line). + out_lines = Path(mb["output_path"]).read_text().splitlines() + assert len(out_lines) == 2 + out_cids = {json.loads(line)["custom_id"] for line in out_lines} + assert out_cids == {cids[0], cids[2]} + + # The derived accumulator carries the failed record's INPUT line. + acc = Path(run_dir) / "failed-records-batch-input.jsonl" + acc_lines = acc.read_text().splitlines() + assert len(acc_lines) == 1 + assert json.loads(acc_lines[0])["custom_id"] == cids[1] + + # The minibatch now reads as all-success (siphoned id no longer yielded). + pairs = dict(iter_batch_results(run_dir)) + assert set(pairs) == {"p0", "p2"} + assert all(r.succeeded for r in pairs.values()) + + def test_synthetic_missing_record_is_siphoned_too(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + # cids[1] omitted -> synthetic-missing -> siphoned. + results={"b0": [BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A")]}, + ) + sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + assert [e["problem_id"] for e in mb["failed_records"]] == ["p1"] + assert mb["num_requests"] == len(mb["id_map"]) == 1 + assert (Path(run_dir) / "failed-records-batch-input.jsonl").is_file() + + +class TestSiphonIdempotency: + def test_refetch_does_not_resiphon(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={"b0": [BatchResult(cids[0], BatchItemStatus.SUCCEEDED, text="A")]}, + ) + fetch_batch(client, run_dir, progress=False) + client.poll_calls.clear() + client.retrieve_calls.clear() + # A FETCHED minibatch is skipped; the siphon never runs twice. + sub2 = fetch_batch(client, run_dir, progress=False) + assert client.poll_calls == [] and client.retrieve_calls == [] + assert len(sub2.minibatches[0]["failed_records"]) == 1 # not doubled + + +class TestNoFailuresNoArtifact: + def test_all_succeeded_writes_no_accumulator_file(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + cids = _cids(BatchSubmission.load(run_dir)) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(c, BatchItemStatus.SUCCEEDED, text="x") for c in cids + ] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + assert "failed_records" not in sub.minibatches[0] + # The Stage-4 artifact appears only when record recovery is in play. + assert not (Path(run_dir) / "failed-records-batch-input.jsonl").exists() + + +class TestMaxAttemptsExhaustion: + def test_exhausted_record_rewritten_max_attempted_and_consolidated(self, tmp_path): + run_dir = _submit_run(tmp_path, n=2, minibatch_size=2, batch_ids=["b0"]) + # Simulate a minibatch already resubmitted to its bound (attempt == MAX_ATTEMPTS), + # so a fresh record's submissions = 0 + MAX_ATTEMPTS -> exhausted on this fetch. + sub0 = BatchSubmission.load(run_dir) + sub0.minibatches[0]["attempt"] = MAX_ATTEMPTS + sub0.save() + cids = _cids(sub0) + client = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.ERRORED, error="boom"), + BatchResult(cids[1], BatchItemStatus.SUCCEEDED, text="A1"), + ] + }, + ) + sub = fetch_batch(client, run_dir, progress=False) + mb = sub.minibatches[0] + # NOT siphoned: kept in id_map, num_requests unchanged, no failed_records. + assert mb.get("failed_records", []) == [] + assert cids[0] in mb["id_map"] + assert mb["num_requests"] == len(mb["id_map"]) == 2 + assert not (Path(run_dir) / "failed-records-batch-input.jsonl").exists() + + # Rewritten in outputs/ as max_attempted (distinct from a transient errored). + pairs = dict(iter_batch_results(run_dir)) + assert pairs["p0"].status == BatchItemStatus.MAX_ATTEMPTED + assert pairs["p1"].succeeded + + # Consolidates terminally; the run is legitimately fully consolidated. + consolidate_batch_results(run_dir) + rec = json.loads((Path(run_dir) / "results" / "p0.json").read_text()) + assert rec["status"] == "max_attempted" + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + assert manifest["fully_consolidated"] is True + assert manifest["pending_failed_records"] == 0 diff --git a/tests/prkit/batch/test_resubmit_records.py b/tests/prkit/batch/test_resubmit_records.py new file mode 100644 index 0000000..ac199bb --- /dev/null +++ b/tests/prkit/batch/test_resubmit_records.py @@ -0,0 +1,336 @@ +"""Tests for the Stage-4 record drain inside ``resubmit_failures``. + +Fully offline. A terminal run carrying pending siphoned ``failed_records`` is +built by submitting (SDK-patched real client) then fetching with a duck-typed +``_FetchClient`` that injects a per-record failure; the drain is then driven by a +duck-typed ``_Client`` whose ``submit_batch`` is a ``MagicMock``. Covers the +drain mechanics (a fresh retry minibatch with the right shape; ``minibatch_count`` +bumped; consumed accumulator), the headline end-to-end recovery loop, the +field-parse correlation regression (``custom_id`` and Gemini ``key``), the +MAX_ATTEMPTS off-by-one across whole-minibatch + record retries, a drain submit +error, and the ledger round-trip of the new minibatch keys. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from prkit.batch import ( + CONSOLIDATED, + SUBMIT_ERROR, + SUBMITTED, + BatchSubmission, + consolidate_batch_results, + fetch_batch, + iter_batch_results, + resubmit_failures, + submit_batch_physics_reasoning, +) +from prkit.core.domain import PhysicsDataset, PhysicsProblem +from prkit.core.model_clients.batch_types import ( + BatchItemStatus, + BatchResult, + BatchState, + BatchStatus, +) + + +# --------------------------------------------------------------------------- # +# Offline builders # +# --------------------------------------------------------------------------- # +def _openai_submit_client(model: str = "gpt-5.1"): + with patch("prkit.core.model_clients.openai.OpenAI") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.openai import OpenAIModel + + return OpenAIModel(model) + + +def _gemini_submit_client(model: str = "gemini-3.5-flash"): + with patch("prkit.core.model_clients.gemini.genai.Client") as mock_cls: + mock_cls.return_value = MagicMock() + from prkit.core.model_clients.gemini import GeminiModel + + return GeminiModel(model) + + +_SUBMIT_BUILDERS = {"openai": _openai_submit_client, "google": _gemini_submit_client} +_ID_FIELD = {"openai": "custom_id", "google": "key"} + + +def _dataset(n: int): + problems = [PhysicsProblem(problem_id=f"p{i}", question=f"Q{i}") for i in range(n)] + return PhysicsDataset(problems, info={"name": "physreason", "version": "1.0"}) + + +class _FetchClient: + def __init__(self, provider="openai", *, states=None, results=None): + self.provider = provider + self.model = "model-x" + self._states = states or {} + self._results = results or {} + + def poll_batch(self, batch_id: str) -> BatchStatus: + seq = self._states[batch_id] + state = seq.pop(0) if len(seq) > 1 else seq[0] + return BatchStatus( + batch_id=batch_id, + state=state, + provider=self.provider, + raw_status=str(state), + counts={"total": 1}, + ) + + def retrieve_batch_results(self, batch_id: str): + return iter(self._results.get(batch_id, [])) + + +class _Client: + """A duck-typed resubmit client: only ``provider`` + ``submit_batch`` are used.""" + + def __init__(self, provider="openai"): + self.provider = provider + self.model = "model-x" + + +def _run_with_siphoned_record(tmp_path, *, provider="openai", run_name="run"): + """Submit 2 problems in one minibatch, fetch with p0 ERRORED -> p0 siphoned. + + Returns the run_dir; afterwards minibatch 0 is FETCHED with one pending + ``failed_records`` entry (custom_id ``p0``, attempt 1) and p1 succeeded. + """ + client = _SUBMIT_BUILDERS[provider]() + client.submit_batch = MagicMock(side_effect=["b0"]) + run_dir = submit_batch_physics_reasoning( + client, + _dataset(2), + output_dir=tmp_path, + run_name=run_name, + minibatch_size=2, + ) + cids = list(BatchSubmission.load(run_dir).minibatches[0]["id_map"]) + fc = _FetchClient( + provider, + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.ERRORED, error="boom"), + BatchResult(cids[1], BatchItemStatus.SUCCEEDED, text="A1"), + ] + }, + ) + fetch_batch(fc, run_dir, progress=False) + return run_dir + + +# --------------------------------------------------------------------------- # +class TestDrainMechanics: + def test_pending_record_drained_into_fresh_minibatch(self, tmp_path): + run_dir = _run_with_siphoned_record(tmp_path) + sub0 = BatchSubmission.load(run_dir) + failed_cid = sub0.minibatches[0]["failed_records"][0]["custom_id"] + + client = _Client() + client.submit_batch = MagicMock(side_effect=["r0"]) + sub = resubmit_failures(client, run_dir) + + # A single fresh retry minibatch appended with the right shape. + assert len(sub.minibatches) == 2 + retry = sub.minibatches[1] + assert retry["index"] == 1 # max(existing index) + 1 + assert retry["status"] == SUBMITTED and retry["batch_id"] == "r0" + assert retry["is_retry"] is True + assert retry["attempt"] == 1 + assert retry["retry_sources"] == [0] + assert retry["record_attempts"] == {failed_cid: 1} + assert retry["id_map"] == {failed_cid: "p0"} + assert retry["num_requests"] == len(retry["id_map"]) == 1 + assert retry["endpoint"] is None and retry["completion_window"] is None + + # Ledger bookkeeping: minibatch_count bumped, source accumulator consumed. + assert sub.minibatch_count == 2 + assert sub.minibatches[0]["failed_records"] == [] + # The derived accumulator file is gone once nothing is pending. + assert not (Path(run_dir) / "failed-records-batch-input.jsonl").exists() + + # submit_batch got exactly the failed record's input line (by wire id). + client.submit_batch.assert_called_once() + sent = client.submit_batch.call_args.args[0] + assert [r["custom_id"] for r in sent] == [failed_cid] + # The retry input file was persisted for resubmittability. + assert (Path(run_dir) / "inputs" / "minibatch_0001.jsonl").is_file() + + def test_drain_submit_error_keeps_record_recoverable(self, tmp_path): + run_dir = _run_with_siphoned_record(tmp_path) + client = _Client() + client.submit_batch = MagicMock(side_effect=RuntimeError("nope")) + sub = resubmit_failures(client, run_dir) + retry = sub.minibatches[1] + # Recorded as SUBMIT_ERROR with an input file -> the whole-minibatch path can + # resubmit it later (records are not lost). + assert retry["status"] == SUBMIT_ERROR and retry["batch_id"] == "" + assert "nope" in retry["error"] + assert sub.minibatches[0]["failed_records"] == [] # still consumed + + +class TestEndToEndRecovery: + def test_full_loop_reaches_fully_consolidated(self, tmp_path): + run_dir = _run_with_siphoned_record(tmp_path) + # After the siphon pruned p0, minibatch 0's id_map holds only "p1"; the + # siphoned (failed) record is "p0". + assert list(BatchSubmission.load(run_dir).minibatches[0]["id_map"]) == ["p1"] + failed_cid = "p0" + + consolidate_batch_results(run_dir) + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + assert manifest["fully_consolidated"] is False # p0 still pending + assert manifest["pending_failed_records"] == 1 + assert not (Path(run_dir) / "results" / "p0.json").exists() + + client = _Client() + client.submit_batch = MagicMock(side_effect=["r0"]) + resubmit_failures(client, run_dir) + + fc = _FetchClient( + "openai", + states={"r0": [BatchState.COMPLETED]}, + results={ + "r0": [BatchResult(failed_cid, BatchItemStatus.SUCCEEDED, text="fixed")] + }, + ) + fetch_batch(fc, run_dir, progress=False) + sub = consolidate_batch_results(run_dir) + + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + assert manifest["fully_consolidated"] is True + assert manifest["pending_failed_records"] == 0 + assert all(mb["status"] == CONSOLIDATED for mb in sub.minibatches) + # The recovered result is the succeeded one. + rec = json.loads((Path(run_dir) / "results" / "p0.json").read_text()) + assert rec["status"] == "succeeded" and rec["text"] == "fixed" + + +class TestCorrelationFieldParse: + @pytest.mark.parametrize("provider", ["openai", "google"]) + def test_drained_record_maps_by_field_not_position(self, tmp_path, provider): + # Reorder the on-disk input lines so a positional zip would mis-correlate; + # the field-parse must still pick p0's own line for p0's failed record. + run_dir = _run_with_siphoned_record(tmp_path, provider=provider) + id_field = _ID_FIELD[provider] + input_path = Path(run_dir) / "inputs" / "minibatch_0000.jsonl" + lines = input_path.read_text().splitlines() + input_path.write_text("\n".join(reversed(lines)) + "\n", encoding="utf-8") + + client = _Client(provider) + client.submit_batch = MagicMock(side_effect=["r0"]) + sub = resubmit_failures(client, run_dir) + + sent = client.submit_batch.call_args.args[0] + assert len(sent) == 1 + # The submitted line is p0's own request (its wire id is p0), not p1's. + assert sent[0][id_field] == "p0" + assert sub.minibatches[1]["id_map"] == {"p0": "p0"} + + +class TestMaxAttemptsOffByOne: + def test_attempt_history_spans_whole_minibatch_and_record_retries(self, tmp_path): + # Owner's scenario guarding the off-by-one. Simulate a minibatch already + # whole-resubmitted once (attempt=2); a dropped record -> submissions 0+2=2 + # (siphoned). After the record drain (attempt=1, record_attempts[cid]=2), a + # further failure -> submissions 2+1=3 == MAX_ATTEMPTS -> MAX_ATTEMPTED. + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=["b0"]) + run_dir = submit_batch_physics_reasoning( + client, _dataset(2), output_dir=tmp_path, run_name="run", minibatch_size=2 + ) + sub0 = BatchSubmission.load(run_dir) + sub0.minibatches[0]["attempt"] = 2 # one prior whole-minibatch resubmit + sub0.save() + cids = list(sub0.minibatches[0]["id_map"]) + + fc = _FetchClient( + "openai", + states={"b0": [BatchState.COMPLETED]}, + results={ + "b0": [ + BatchResult(cids[0], BatchItemStatus.ERRORED, error="e1"), + BatchResult(cids[1], BatchItemStatus.SUCCEEDED, text="A1"), + ] + }, + ) + sub = fetch_batch(fc, run_dir, progress=False) + assert sub.minibatches[0]["failed_records"][0]["attempt"] == 2 # 0 + 2 + + rclient = _Client() + rclient.submit_batch = MagicMock(side_effect=["r0"]) + sub = resubmit_failures(rclient, run_dir) + retry = sub.minibatches[1] + assert retry["record_attempts"] == {cids[0]: 2} and retry["attempt"] == 1 + + fc2 = _FetchClient( + "openai", + states={"r0": [BatchState.COMPLETED]}, + results={"r0": [BatchResult(cids[0], BatchItemStatus.ERRORED, error="e2")]}, + ) + sub = fetch_batch(fc2, run_dir, progress=False) + retry = sub.minibatches[1] + # submissions = 2 + 1 = 3 == MAX_ATTEMPTS -> exhausted, NOT re-siphoned. + assert retry.get("failed_records", []) == [] + assert cids[0] in retry["id_map"] + pairs = dict(iter_batch_results(run_dir)) + assert pairs[cids[0]].status == BatchItemStatus.MAX_ATTEMPTED + + +class TestLedgerRoundTrip: + def test_stage4_minibatch_keys_survive_to_dict_from_dict(self, tmp_path): + run_dir = _run_with_siphoned_record(tmp_path) + client = _Client() + client.submit_batch = MagicMock(side_effect=["r0"]) + sub = resubmit_failures(client, run_dir) + + restored = BatchSubmission.from_dict(sub.to_dict()) + retry = restored.minibatches[1] + assert retry["is_retry"] is True + assert retry["attempt"] == 1 + assert retry["retry_sources"] == [0] + assert retry["record_attempts"] == {"p0": 1} + assert restored.minibatches[0]["failed_records"] == [] + + +class TestPreconditions: + def test_drain_requires_terminal_ledger(self, tmp_path): + # A run left RUNNING: resubmit refuses before any submit (terminal gate). + from prkit.batch import BatchNotTerminalError + + client = _openai_submit_client() + client.submit_batch = MagicMock(side_effect=["b0"]) + run_dir = submit_batch_physics_reasoning( + client, _dataset(1), output_dir=tmp_path, run_name="run", minibatch_size=1 + ) + fetch_batch( + _FetchClient(states={"b0": [BatchState.IN_PROGRESS]}), + run_dir, + progress=False, + ) + rclient = _Client() + rclient.submit_batch = MagicMock() + with pytest.raises(BatchNotTerminalError): + resubmit_failures(rclient, run_dir) + rclient.submit_batch.assert_not_called() + + +class TestNextCommand: + def test_consolidate_done_hint_points_at_resubmit_failures(self, tmp_path, caplog): + run_dir = _run_with_siphoned_record(tmp_path) + with caplog.at_level(logging.INFO, logger="prkit.batch"): + caplog.clear() + consolidate_batch_results(run_dir) + msgs = [r.getMessage() for r in caplog.records if r.name == "prkit.batch"] + assert any( + "records still failed" in m and "resubmit_failures" in m for m in msgs + ) From ca92711781dc629500161e7095fc8f065ca51624 Mon Sep 17 00:00:00 2001 From: Yinghuan Zhang Date: Mon, 22 Jun 2026 23:01:58 -0400 Subject: [PATCH 14/14] Refresh BATCH_MODE.md for Stage 4 (record recovery) Updates the rollup doc with Stage-4 design: siphon-at-fetch, record-drain in resubmit_failures, MAX_ATTEMPTED terminal status, CANCELLED reversal, the new end-to-end loop, and all vocabulary/state-machine sections. Co-Authored-By: Claude Sonnet 4.6 --- BATCH_MODE.md | 237 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 172 insertions(+), 65 deletions(-) diff --git a/BATCH_MODE.md b/BATCH_MODE.md index ed3c6dc..1daaae4 100644 --- a/BATCH_MODE.md +++ b/BATCH_MODE.md @@ -1,7 +1,7 @@ # Batch Mode (N4) — Design & Usage Reference **Single source of truth** for prkit's discounted provider **batch lane**. This doc is the rolled-up, -up-to-date design: it supersedes the per-stage design notes in `internal/N4_BATCH_DESIGN_STAGE_{1,2,3}.md` +up-to-date design: it supersedes the per-stage design notes in `internal/N4_BATCH_DESIGN_STAGE_{1,2,3,4}.md` (kept only as historical rationale). When a future stage is designed, **roll its committed decisions into this file** rather than leaving them stranded in a stage doc — see [Maintaining this doc](#maintaining-this-doc). @@ -11,23 +11,30 @@ file** rather than leaving them stranded in a stage doc — see [Maintaining thi |---|---|---|---| | **1 — Submit** | preprocess → split → write JSONL → submit → ledger | ✅ **Implemented & usable** | `submit_batch_physics_reasoning`, `build_problem_batch_request` | | **2 — Fetch** | poll → download → correlate → resume | ✅ **Implemented & usable** | `fetch_batch`, `iter_batch_results`, `batch_fetch_supported` | -| **3 — Finalize** | consolidate per-problem results · resubmit failed minibatches | ✅ **Implemented & usable** | `consolidate_batch_results`, `resubmit_failed_minibatches` | +| **3 — Finalize** | consolidate per-problem results · resubmit failed minibatches | ✅ **Implemented & usable** | `consolidate_batch_results`, `resubmit_failures` | +| **4 — Record recovery** | siphon per-record failures at fetch · drain them into fresh minibatches · `MAX_ATTEMPTS` bound | ✅ **Implemented & usable** | siphon inside `fetch_batch`, drain inside `resubmit_failures`, `BatchItemStatus.MAX_ATTEMPTED` | -> **All three stages are implemented and callable** as of 2026-06-22 (104 batch tests pass on -> `feat/n4-batch-mode`). Everything below describes the live API. +> **All four stages are implemented and callable** as of 2026-06-22 (120 batch tests pass on +> `feat/n4-batch-mode`). Everything below describes the live API. Stage 4 reuses the same two finalize +> verbs — there is no new public function; `fetch_batch` gained an internal siphon and `resubmit_failures` +> (renamed from Stage 3's `resubmit_failed_minibatches`) gained a record-drain. --- ## Concepts & vocabulary -The vocabulary was fixed in Stage 2 (it renamed Stage 1's terms — see [Design evolution](#design-evolution)). +The vocabulary was fixed in Stage 2 (it renamed Stage 1's terms — see [Design evolution](#design-evolution)); +Stage 4 added **record** as a recoverable unit. - **batch** — the *whole thing* a user triggers over a dataset. One `submit_batch_physics_reasoning(...)` call → one **run folder** → one `BatchSubmission` **ledger** (`metadata.json`). - **minibatch** — one `minibatch_size`-problem group = one provider batch job = one `minibatch_XXXX.jsonl` = - one element of `BatchSubmission.minibatches`. **The minibatch is the unit of success**: it either succeeded - (`FETCHED`) or failed as a whole. (Record-level failures within a fetched minibatch surface as `ERRORED` - results, not as minibatch failures.) + one element of `BatchSubmission.minibatches`. The minibatch is the unit for **whole-job** failures + (`FAILED` / `EXPIRED` / `CANCELLED` / `SUBMIT_ERROR`). +- **record** — one problem's single request inside a minibatch, correlated by its wire id + (`custom_id`, or `key` for Gemini) back to its `problem_id` via the minibatch `id_map`. **Since Stage 4 the + record is recoverable**: an individual failed record is *siphoned* out of an otherwise-good minibatch and + retried on its own (bounded by `MAX_ATTEMPTS`), rather than riding as a permanent `ERRORED` outcome. - **run folder** — the consumer-owned directory holding the ledger, the input JSONL, and (after fetch) the output JSONL. **Disk is the source of truth** across the provider's ~24h window. - The names `prkit.batch` / `submit_batch_*` / `fetch_batch` keep the word "batch" because it names the @@ -42,7 +49,8 @@ hub should own once — preprocessing parity, splitting, JSONL writing, submissi It deliberately does **not** own: - **An end-to-end runner / outer loop.** No dataset loading, no inference orchestration, no auto-chaining. The - consumer drives `submit → fetch → [resubmit → fetch]* → consolidate`. + consumer drives `submit → fetch → consolidate → [resubmit → fetch → consolidate]*` (the siphon happens + *inside* the fetch pass — it adds no new sweep). - **Scoring.** No scorer adapter. Batch mode stops at correlated `BatchResult`s; the consumer calls `prkit.api.Scorer` / `Verdict` directly (see [Scoring seam](#scoring-seam)). - **Pricing.** That is N6's job (the cost meter). Batch mode forks no pricing and imports no `prkit.cost`. The @@ -101,13 +109,16 @@ print(sub.is_complete(), sub.status_counts()) for problem_id, result in iter_batch_results(run_dir): # accepts run_dir or the ledger if result.succeeded: verdict = scorer.score(result.text, gold[problem_id]) # consumer's scoring seam (no adapter) - # non-success: result.status is ERRORED/EXPIRED/CANCELED, result.text is None, result.error set + # non-success: result.status is ERRORED/EXPIRED/CANCELED/MAX_ATTEMPTED, result.text is None, result.error set ``` `fetch_batch` without `wait=True` does **one** poll-and-download pass and returns — ideal for a cron / `/loop` driver that re-invokes until `sub.is_complete()`. Each pass logs a one-line INFO summary on the `prkit.batch` logger (suppress with `progress=False`). +To recover individual failed records (not just whole failed minibatches) and reach genuine 100%, drive the +Stage-4 loop in [Stage 4 — Record recovery](#stage-4--record-recovery) instead of stopping at step 3. + --- ## Public API (implemented) @@ -172,9 +183,15 @@ prkit.batch.fetch_batch( client.fetch_batch_physics_reasoning(run_dir_or_submission, **kwargs) -> BatchSubmission ``` -- Polls each non-final minibatch once; on `COMPLETED`/`EXPIRED` it downloads results to - `outputs/minibatch_XXXX.jsonl` and marks the minibatch `FETCHED`. The ledger is saved after **every** change - (crash-safe). +- Polls each non-final minibatch once; on `COMPLETED`/`EXPIRED` it downloads results, **siphons** any + per-record failures (Stage 4 — see below), writes succeeded-only to `outputs/minibatch_XXXX.jsonl`, and marks + the minibatch `FETCHED`. The ledger is saved after **every** change (crash-safe). +- **Stage-4 siphon (internal, no signature change):** when a minibatch is retrieved, each failed record (real + `ERRORED`/`EXPIRED`/`CANCELED`, or a synthetic-missing id) under `MAX_ATTEMPTS` is recorded on the minibatch's + `failed_records`, pruned from its `id_map`, and `num_requests` is decremented — so the minibatch reads as + *fully successful*. A record that has hit `MAX_ATTEMPTS` is instead kept and written with status + `MAX_ATTEMPTED` (a permanent give-up). The run-level `failed-records-batch-input.jsonl` accumulator is + refreshed; `resubmit_failures` later drains it. - **Idempotent/resumable:** `FETCHED` / `SUBMIT_ERROR` / `FAILED` / `CANCELLED` (and `CONSOLIDATED`) minibatches are skipped, so re-runs never re-hit the network for finished work. `EXPIRED` is still *re-polled* (cheap) and still *retrieved* (it can carry a completed subset). @@ -188,9 +205,13 @@ prkit.batch.iter_batch_results(submission) -> Iterator[tuple[str, BatchResult]] - Pure offline reader (no network, re-runnable for re-scoring). Accepts a `BatchSubmission` or a `run_dir`. - For every minibatch with a downloaded output file, correlates each line's `custom_id` → `problem_id` via the - minibatch's `id_map`, yielding **one `(problem_id, BatchResult)` per submitted problem in input order**. -- **Completeness:** any submitted id the provider never returned yields a synthetic `ERRORED` `BatchResult`; - extra/uncorrelated ids are dropped and counted on the minibatch (`uncorrelated_count`). + minibatch's `id_map`, yielding **one `(problem_id, BatchResult)` per remaining `id_map` entry in input order**. +- **Completeness:** any id still in `id_map` that the provider never returned yields a synthetic `ERRORED` + `BatchResult`; extra/uncorrelated ids are dropped and counted on the minibatch (`uncorrelated_count`). +- **After the Stage-4 siphon**, a fetched minibatch's failed records have been *pruned from `id_map`* (they now + live in the accumulator), so this reader yields them as results only once they are drained, retried, and + re-fetched in a new minibatch. Records that exhausted `MAX_ATTEMPTS` stay in `id_map` and yield a terminal + `MAX_ATTEMPTED` result. Net effect: every submitted problem yields **exactly one** final result. ### Capability check @@ -204,20 +225,23 @@ prkit.batch.batch_fetch_supported(client) -> bool # True for openai / anthropi @dataclass(frozen=True) class BatchResult: custom_id: str - status: BatchItemStatus # SUCCEEDED | ERRORED | EXPIRED | CANCELED + status: BatchItemStatus # SUCCEEDED | ERRORED | EXPIRED | CANCELED | MAX_ATTEMPTED text: str | None = None # the model's free-text output on success; None otherwise error: str | None = None # failure description on any non-success outcome @property def succeeded(self) -> bool: ... ``` +`MAX_ATTEMPTED` (value `"max_attempted"`) is the one prkit-synthesized status — a record that exhausted +`MAX_ATTEMPTS`, distinct from a transient `ERRORED` so a downstream scorer can tell "we gave up" from "errored." + ### Errors | Error | Base | Raised when | |---|---|---| | `BatchInputError` | `ValueError` | empty input, duplicate/illegal id, pre-existing non-empty run folder without `overwrite` | | `BatchFetchUnsupportedError` | `BatchInputError` | `fetch_batch` (or resubmit) called on a non-batch provider | -| `BatchNotTerminalError` | `BatchInputError` | *(Stage 3)* `resubmit_failed_minibatches` called on a non-terminal ledger | +| `BatchNotTerminalError` | `BatchInputError` | `resubmit_failures` called on a non-terminal ledger | --- @@ -279,24 +303,37 @@ Each `minibatches[i]` dict: "endpoint": str | None, # audit-only (provider parity) "completion_window": str | None, # audit-only # "uncorrelated_count": int # added by iter_batch_results when extra ids were dropped + # ---- Stage-4 additive keys (round-trip for free via dict(mb)) ---- + "attempt": int, # whole-minibatch submission count (submit => 1; resubmit => +1) + "failed_records": list[dict], # pending siphoned failures: {custom_id, problem_id, error, attempt} + # ---- on RETRY minibatches only (built by resubmit_failures' record-drain) ---- + "is_retry": bool, # True for a record-drain minibatch + "retry_sources": list[int], # source minibatch indices pooled into it + "record_attempts": dict[str, int], # per-record prior submission count, carried forward } ``` +**Invariant:** `num_requests == len(id_map)` on every minibatch (the siphon decrements both in lockstep). A +record's total submission count is `record_attempts.get(custom_id, 0) + attempt`; it is siphoned only while that +is `< MAX_ATTEMPTS`. `minibatch_count == len(minibatches)` counts submit **+** appended retry minibatches; +`total_problems` is unchanged (a retried problem is the same problem). + ### Run-folder layout ``` // - metadata.json # the BatchSubmission ledger (mutable; advanced by fetch) + metadata.json # the BatchSubmission ledger (mutable; advanced by fetch) inputs/ - minibatch_0000.jsonl # one provider-correct request line per problem + minibatch_0000.jsonl # one provider-correct request line per problem (never pruned) minibatch_0001.jsonl ... - outputs/ # written by fetch: one normalized-BatchResult JSONL per fetched minibatch - minibatch_0000.jsonl # lines: {"custom_id", "status", "text", "error"} + outputs/ # written by fetch: succeeded-only (+ MAX_ATTEMPTED) per fetched minibatch + minibatch_0000.jsonl # lines: {"custom_id", "status", "text", "error"} ... - results/ # written by consolidate (Stage 3): one file per problem - .json - results_manifest.json # written by consolidate (Stage 3): consolidation summary + results/ # written by consolidate (Stage 3): one file per problem + .json # a recovered (Stage-4) result supersedes a stale one for the same problem + results_manifest.json # written by consolidate: summary + pending_failed_records + fully_consolidated + failed-records-batch-input.jsonl # Stage 4: derived accumulator — request lines of every pending failed record ``` --- @@ -308,16 +345,20 @@ Minibatch status constants (`prkit.batch`, plain strings): `SUBMIT_ERROR`, `FETCH_ERROR`. ``` +# minibatch level SUBMITTED ─poll→ RUNNING ─poll→ COMPLETED ─retrieve→ FETCHED ─consolidate→ CONSOLIDATED (terminal-good) └poll→ EXPIRED(partial+results) → FETCHED (counts as success) - └poll→ EXPIRED(empty) ─────────┐ - └poll→ FAILED ─────────────────┤ resubmit (Stage 3) ─→ SUBMITTED (re-enters loop) -SUBMIT_ERROR ────────────────────────────────────────────┘ - └poll→ CANCELLED (terminal; NOT auto-resubmitted — see Stage 3 limitation) -FETCH_ERROR (non-terminal) ─re-run fetch_batch→ ... + └poll→ EXPIRED(empty)/FAILED/CANCELLED ┐ +SUBMIT_ERROR ──────────────────────────────────────────────────┤ resubmit_failures ─→ SUBMITTED (re-enters loop) +FETCH_ERROR (non-terminal) ─re-run fetch_batch→ ... ┘ (CANCELLED is now resubmitted too — Stage 4 D1) + +# record level (Stage 4) +record fail (errored/expired/canceled/synthetic) ─ submissions < MAX_ATTEMPTS ─ siphon → accumulator → new minibatch + └ submissions ≥ MAX_ATTEMPTS ─ written as MAX_ATTEMPTED (terminal) ``` -**Poll state → minibatch status** (the fetch-pass mapping): +**Poll state → minibatch status** (the fetch-pass mapping; on `COMPLETED`/`EXPIRED` the retrieved records are +additionally partitioned and per-record failures siphoned): | `BatchState` from poll | fetch action | resulting status | |---|---|---| @@ -337,6 +378,8 @@ Key status sets (internal, but they explain behavior): - `_COMPLETE_STATUSES = {FETCHED, CONSOLIDATED, FAILED, CANCELLED, SUBMIT_ERROR, EXPIRED}` — `is_complete()` / `wait=True` stop once every minibatch is one of these. - `_HAS_OUTPUT_STATUSES = {FETCHED, CONSOLIDATED}` — minibatches with a readable `outputs/` file. +- `_RESUBMIT_STATUSES = {FAILED, SUBMIT_ERROR, EXPIRED, CANCELLED}` — re-submitted by `resubmit_failures` + (Stage 4 added `CANCELLED`). `MAX_ATTEMPTS = 3` bounds total submissions per record. --- @@ -365,51 +408,106 @@ prkit.batch.consolidate_batch_results( - **Filename safety:** `problem_id` is sanitized to a filesystem-safe name; a collision (two problems mapping to the same file) raises `BatchInputError` rather than silently overwriting. -### `resubmit_failed_minibatches` (needs client, terminal-gated) +### `resubmit_failures` (needs client, terminal-gated) ```python -prkit.batch.resubmit_failed_minibatches(client, submission) -> BatchSubmission -client.resubmit_failed_minibatches(run_dir_or_submission, **kwargs) -> BatchSubmission # facade +prkit.batch.resubmit_failures(client, submission) -> BatchSubmission +client.resubmit_failures(run_dir_or_submission, **kwargs) -> BatchSubmission # facade ``` -- Re-submits each `FAILED` / `SUBMIT_ERROR` / `EXPIRED` minibatch (`_RESUBMIT_STATUSES`) by re-reading its - persisted `inputs/minibatch_XXXX.jsonl` and calling `client.submit_batch`, resetting the ledger entry to - `SUBMITTED` with a new `batch_id` **only after submit succeeds** (a failed submit → `SUBMIT_ERROR`, loop - continues). The consumer then re-runs `fetch_batch` on the new jobs. -- **Excludes `CANCELLED`** (a cancel can be deliberate; auto-resubmit would fight intent). +> Renamed from Stage 3's `resubmit_failed_minibatches` (the old name is gone — N4 is unmerged, no external +> consumers). One call now does **both** of the jobs below. + +- **(a) Whole-minibatch resubmit.** Re-submits each `FAILED` / `SUBMIT_ERROR` / `EXPIRED` / **`CANCELLED`** + minibatch (`_RESUBMIT_STATUSES`) by re-reading its persisted `inputs/minibatch_XXXX.jsonl` and calling + `client.submit_batch`, resetting the ledger entry to `SUBMITTED` with a new `batch_id` **only after submit + succeeds** (a failed submit → `SUBMIT_ERROR`, loop continues) and bumping its `attempt`. +- **(b) Record-drain (Stage 4).** Drains the run-level failed-records accumulator into **fresh** minibatches — + one chunk of ≤ `minibatch_size` per submit, with a new monotonic `index`, rebuilt `id_map` + `record_attempts`, + `is_retry=True` and `retry_sources`, appended to the ledger as `SUBMITTED` (`minibatch_count++`). The consumed + `failed_records` entries are cleared in the same atomic save. +- **`CANCELLED` is now resubmitted** (Stage 4 reversed Stage 3's exclusion — see [Design evolution](#design-evolution)), + so a cancelled minibatch is no longer a dead-end. - Requires a terminal ledger (`is_complete()`), else raises `BatchNotTerminalError` (run `fetch_batch` first); - raises `BatchFetchUnsupportedError` up front for a non-batch provider. + raises `BatchFetchUnsupportedError` up front for a non-batch provider. The consumer then re-runs `fetch_batch` + on the new jobs. -### Intended end-to-end loop (consumer-driven) +### Status → finalize action (single source of truth) -```python -run_dir = client.submit_batch_physics_reasoning(dataset) -sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) -while not all(mb["status"] in ("fetched", "consolidated") for mb in sub.minibatches): - sub = client.resubmit_failed_minibatches(run_dir) # Stage 3 - sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) -sub = consolidate_batch_results(run_dir) # Stage 3 -> results/.json -for problem_id, result in iter_batch_results(run_dir): - if result.succeeded: - verdict = scorer.score(result.text, gold[problem_id]) -``` +| minibatch `status` | `consolidate_batch_results` | `resubmit_failures` | +|---|---|---| +| `FETCHED` | write `results/.json` → `CONSOLIDATED` | skip (but its pending `failed_records` are drained) | +| `CONSOLIDATED` | skip (done) | skip (but its pending `failed_records` are drained) | +| `FAILED` / `SUBMIT_ERROR` / `EXPIRED` / **`CANCELLED`** | skip + warn | re-read input → `submit_batch` → `SUBMITTED` (`attempt++`) | +| `RUNNING` / `SUBMITTED` / `COMPLETED` / `FETCH_ERROR` | skip + warn | precondition `is_complete()` fails → raise | -### Known limitation (Stage 3) +--- -A `CANCELLED` minibatch is neither resubmitted nor consolidatable, so a batch containing one can never reach -"fully consolidated" and a naïve `while not all-done` loop would spin. Surfaced via the manifest's -`fully_consolidated` flag and the next-command guidance; a future stage may add an opt-in -`force_resubmit_cancelled` path. +## Stage 4 — Record recovery -### Status → finalize action (single source of truth) +Stages 1–3 close the lane at **minibatch** granularity: only a whole failed minibatch could be retried, so an +individual `ERRORED` record inside an otherwise-good minibatch stranded that one problem below 100% (and a +`CANCELLED` minibatch was a dead-end). Stage 4 adds **record-level** recovery without a new public function: a +siphon inside `fetch_batch` and a record-drain inside `resubmit_failures` (above). + +### The siphon (inside `fetch_batch`) + +When a minibatch is retrieved, its records are partitioned into succeeded vs failed: -| minibatch `status` | consolidate | resubmit | +- **Succeeded** records are written to `outputs/minibatch_XXXX.jsonl` as before. +- **Failed** records (real `ERRORED`/`EXPIRED`/`CANCELED`, or a synthetic-missing id) with total submissions + `< MAX_ATTEMPTS` are **siphoned**: appended to the minibatch's `failed_records`, pruned from its `id_map`, and + `num_requests` decremented. The minibatch is then `FETCHED` and reads as fully successful. +- Failed records that have hit `MAX_ATTEMPTS` (= 3) are **not** re-siphoned — they are written to `outputs/` + with the terminal status `MAX_ATTEMPTED`, kept in `id_map`, and consolidate normally. +- The derived `failed-records-batch-input.jsonl` accumulator (the failed records' original *input* lines) is + refreshed. Failures are correlated back to their input line by **parsing the provider id field** (`custom_id`, + or `key` for Gemini), never positional zip. + +`MAX_ATTEMPTS` is a fixed module constant (`prkit.batch.MAX_ATTEMPTS = 3`), not a per-call argument. A record's +count spans **both** whole-minibatch retries and record-level retries: +`submissions = record_attempts.get(custom_id, 0) + attempt`. + +### Record outcome → action (single source of truth) + +| record outcome on a `FETCHED` minibatch | `submissions` | action | |---|---|---| -| `FETCHED` | write `results/.json` → `CONSOLIDATED` | skip | -| `CONSOLIDATED` | skip (done) | skip | -| `FAILED` / `SUBMIT_ERROR` / `EXPIRED` | skip + warn | re-read input → `submit_batch` → `SUBMITTED` (or `SUBMIT_ERROR`) | -| `CANCELLED` | skip + warn (dead-end) | skip (known limitation) | -| `RUNNING` / `SUBMITTED` / `COMPLETED` / `FETCH_ERROR` | skip + warn | precondition `is_complete()` fails → raise | +| `SUCCEEDED` | — | write to `outputs/`; stays in `id_map`; consolidates normally | +| `ERRORED`/`EXPIRED`/`CANCELED`/synthetic-missing | `< MAX_ATTEMPTS` | **siphon** → `failed_records`; prune `id_map` + `num_requests--`; refresh accumulator | +| `ERRORED`/`EXPIRED`/`CANCELED`/synthetic-missing | `≥ MAX_ATTEMPTS` | write to `outputs/` as `MAX_ATTEMPTED`; **keep** in `id_map`; consolidates terminally | + +### Manifest gating & loop termination + +`consolidate_batch_results` writes `pending_failed_records = Σ len(mb["failed_records"])` into +`results_manifest.json` and sets `fully_consolidated = (every minibatch CONSOLIDATED) and pending_failed_records +== 0`. An exhausted `MAX_ATTEMPTED` record is **not** in `failed_records`, so a run with only permanent failures +is legitimately "done." The post-consolidate log points at `resubmit_failures` while records are still pending. + +The recovery loop **terminates**: every record either succeeds or reaches `MAX_ATTEMPTS` and becomes a terminal +`MAX_ATTEMPTED` result (no longer pending), so `pending_failed_records` strictly decreases to 0 — no spin. + +### End-to-end loop (consumer-driven, reaches genuine 100%) + +```python +import json +from pathlib import Path + +run_dir = client.submit_batch_physics_reasoning(dataset) +sub = client.fetch_batch_physics_reasoning(run_dir, wait=True) # siphons per-record failures +consolidate_batch_results(run_dir) # writes results_manifest.json +manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) +while not manifest["fully_consolidated"]: # False while minibatches OR records are pending + client.resubmit_failures(run_dir) # whole-minibatch retries + drain failed records + client.fetch_batch_physics_reasoning(run_dir, wait=True) + consolidate_batch_results(run_dir) + manifest = json.loads((Path(run_dir) / "results_manifest.json").read_text()) + +for problem_id, result in iter_batch_results(run_dir): # offline, exactly one result per problem + if result.succeeded: + verdict = scorer.score(result.text, gold[problem_id]) # consumer's scoring seam (no adapter) + elif str(result.status) == "max_attempted": + ... # a record we permanently gave up on +``` --- @@ -419,7 +517,9 @@ Batch mode stops at correlated `BatchResult`s. The consumer scores by reading `o `iter_batch_results` (or the per-problem `results/.json` files after consolidation) and calling `prkit.api.Scorer.score(prediction, reference) -> Verdict` directly. The reference/gold answers live on the consumer's `PhysicsProblem`s, not in the ledger. `prkit.batch` imports no `prkit.api`, accepts no `Scorer`, and -owns no references — scoring is genuinely the consumer's. +owns no references — scoring is genuinely the consumer's. A record may consolidate with status `max_attempted` +(a permanent give-up after `MAX_ATTEMPTS`) — the consumer can treat it distinctly from `errored`/`expired` +(e.g. report the give-up rate, or exclude it from scoring) without any prkit policy. ## Cost-meter (N6) seam @@ -453,6 +553,11 @@ Decisions that changed across stages — recorded so the renames and contract sh | xAI = "build-only" batch | xAI = **no batch surface** | Stage-1 cleanup deleted its only (unused, structured) batch method | | `iter_batch_results` gates on `== FETCHED` | *(Stage 3)* widens to `_HAS_OUTPUT_STATUSES` so re-scoring still works after consolidation | consolidate keeps `outputs/` files | | `submit(overwrite=True)` clears `inputs/`+`outputs/` | *(Stage 3)* also clears `results/`+`results_manifest.json` | a reused folder must not leave a stale results set | +| *(Stage 3)* `resubmit_failed_minibatches` **excludes** `CANCELLED` (a "dead-end" known limitation) | *(Stage 4)* `resubmit_failures` **includes** `CANCELLED` (`_RESUBMIT_STATUSES += CANCELLED`) | always-resubmit removes the dead-end; loops can reach 100% | +| *(Stage 3)* verb `resubmit_failed_minibatches` | *(Stage 4)* renamed `resubmit_failures` (+ facade); old name removed | it now drains record failures too, not just whole minibatches | +| *(Stage 3)* "the minibatch is the unit of success"; a failed record is permanently `ERRORED` | *(Stage 4)* **record-level recovery** — `fetch` siphons failed records, `resubmit_failures` drains them | reach genuine 100% without re-running 499 good records to recover 1 | +| *(Stage 3)* `fully_consolidated` = every minibatch `CONSOLIDATED` | *(Stage 4)* also requires `pending_failed_records == 0` | a `while not fully_consolidated` loop must keep going while records are pending | +| *(Stage 3)* `BatchItemStatus` = provider set only | *(Stage 4)* adds `MAX_ATTEMPTED` (record exhausted `MAX_ATTEMPTS=3`) | distinguish "we gave up" from a transient `errored` | Also removed during Stage 1 (no longer part of the contract): the reserved `api.Runner` Protocol, and the unused structured-batch request/response wrappers. The structured-output *engine* (`StructuredOutputPlan` / @@ -477,6 +582,8 @@ This file is the rollup. When a new stage is designed and its decisions are owne - `internal/N4_BATCH_DESIGN_STAGE_1.md` — submit half; Q1–Q6; `Runner` + structured-batch cleanup. - `internal/N4_BATCH_DESIGN_STAGE_2.md` — fetch half; ledger reshape; Stage-1 amendments. - `internal/N4_BATCH_DESIGN_STAGE_3.md` — finalize (consolidate + resubmit); Stage-1/2 amendments. +- `internal/N4_BATCH_DESIGN_STAGE_4.md` — record recovery (siphon-at-fetch + record-drain); CANCELLED reversal; + `resubmit_failures` rename; `MAX_ATTEMPTED`; Stage-1/2/3 amendments. - Code: `src/prkit/batch/__init__.py`, `src/prkit/core/model_clients/base.py` (facades), `src/prkit/core/model_clients/batch_types.py` (`BatchResult` / `BatchItemStatus` / `BatchState`). - Roadmap: `internal/DEVELOPMENT_ROADMAP.md` §N4, §N6.