From d6a796ccbd544789253bfb8d9f960e23237c7d19 Mon Sep 17 00:00:00 2001 From: YW <58594437+AnsonDev42@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:51:58 +0100 Subject: [PATCH] fix(runtime): poll batches during ingestion Reconcile active batches at durable chunk boundaries so completed remote work releases local enqueue capacity before source exhaustion. --- README.md | 4 +- docs/design_docs/ARCHITECTURE.md | 25 +- docs/design_docs/INGESTION_POLLING_FIX.md | 76 ++++++ docs/design_docs/OPENAI_BATCHING.md | 2 +- docs/design_docs/STORAGE_AND_RUNS.md | 2 + docs/smoke-test.md | 3 + pyproject.toml | 2 +- src/batchor/runtime/ingestion.py | 25 +- tests/integration/test_batchor_runner.py | 74 ++++++ tests/unit/test_batchor_runtime_ingestion.py | 253 +++++++++++++++++++ uv.lock | 2 +- 11 files changed, 455 insertions(+), 13 deletions(-) create mode 100644 docs/design_docs/INGESTION_POLLING_FIX.md diff --git a/README.md b/README.md index cf25017..ac91411 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ Durability is split on purpose: - the artifact store keeps replayable request JSONL and downloaded raw batch payloads That split is what allows retries and fresh-process resume without keeping every request inline in the control-plane store. -On resume, existing active provider batches are polled before `batchor` materializes or submits new local work, so completed batches and retry backoff are reconciled first. Deterministic sources also use stored checkpoint completion metadata to avoid opening a fully materialized Parquet or composite source just to discover there are no rows left. +On resume, existing active provider batches are polled before `batchor` materializes or submits new local work, so completed batches and retry backoff are reconciled first. During long ingestion, active batches are also reconciled at durable item-chunk boundaries when the configured provider polling interval is due. Polling happens before that chunk's submission attempt, so terminal remote work releases locally accounted enqueue tokens without waiting for the entire source to be materialized. Deterministic sources also use stored checkpoint completion metadata to avoid opening a fully materialized Parquet or composite source just to discover there are no rows left. Incomplete checkpointed ingestion keeps the run non-terminal. In the process that started the job, `Run.resume()` can continue from the attached source. After a fresh-process rehydration, call `start(job, run_id=...)`; `refresh()` otherwise @@ -338,7 +338,7 @@ run.wait() print(run.results()[0].output) ``` -While waiting, `batchor` polls active batches before submitting more work. If a refresh consumes a completed batch, records terminal item state, or submits another provider batch, `Run.wait()` immediately continues draining instead of sleeping for the next poll interval. +While ingesting or waiting, `batchor` polls active batches before submitting more work when the provider's monotonic polling cadence is due. This avoids one provider request per fast ingestion chunk while allowing long ingestion to reuse capacity released by completed batches. If a refresh consumes a completed batch, records terminal item state, or submits another provider batch, `Run.wait()` immediately continues draining instead of sleeping for the next poll interval. Structured-output models are validated up front against the OpenAI strict-schema subset used by `batchor`. diff --git a/docs/design_docs/ARCHITECTURE.md b/docs/design_docs/ARCHITECTURE.md index d932912..2eb391a 100644 --- a/docs/design_docs/ARCHITECTURE.md +++ b/docs/design_docs/ARCHITECTURE.md @@ -145,13 +145,16 @@ Internally that expands to: 1. Resolve provider and storage implementations. 2. Persist run config and ingest items into durable state. 3. On resume, poll any already-active provider batches before ingesting or submitting new work. -4. Claim a bounded submission window from pending items. -5. Build or replay request JSONL rows. -6. Persist request artifacts before upload. -7. Submit one or more provider batch files. -8. Poll active batches. -9. Download output/error files. -10. Parse terminal item results back into the state store. +4. Persist source items in durable chunks; at each chunk boundary, poll active batches before submission when the monotonic provider cadence is due. +5. Claim a bounded submission window from pending items. +6. Build or replay request JSONL rows. +7. Persist request artifacts before upload. +8. Submit one or more provider batch files. +9. Poll active batches. +10. Download output/error files. +11. Parse terminal item results back into the state store. + +Ingestion polling stays inside the ingestion module behind its existing poll callback. Its local monotonic cadence state does not add storage schema or background-thread coordination. After each ingestion poll, control state and retry backoff are re-read before any submission or further checkpointed-source materialization. When a caller uses `Run.wait()`, the runtime repeats that poll-and-submit pass until the run is terminal. A pass that changes durable work state, such as consuming completed batches or submitting more items, immediately triggers the next pass rather than sleeping for the configured poll interval. @@ -173,6 +176,14 @@ sequenceDiagram BatchRunner->>StateStore: append_items(materialized_items) BatchRunner->>StateStore: set_ingest_checkpoint() + opt active batch exists and monotonic poll cadence is due + BatchRunner->>Provider: get_batch(batch_id) + Provider-->>BatchRunner: BatchRemoteRecord + BatchRunner->>StateStore: consume terminal batch state + end + + BatchRunner->>StateStore: claim pending items within freed capacity + loop refresh() cycle (polling + submission) BatchRunner->>StateStore: get_active_batches() StateStore-->>BatchRunner: [ActiveBatchRecord...] diff --git a/docs/design_docs/INGESTION_POLLING_FIX.md b/docs/design_docs/INGESTION_POLLING_FIX.md new file mode 100644 index 0000000..93d3bf8 --- /dev/null +++ b/docs/design_docs/INGESTION_POLLING_FIX.md @@ -0,0 +1,76 @@ +# Periodic reconciliation during ingestion + +Status: implementation specification for the `0.0.6` bug-fix release. + +## Outcome + +Long-running item ingestion must continue reconciling active provider batches so +completed remote work releases locally accounted enqueue tokens and permits new +batches to be submitted before the entire source is materialized. + +## Non-goals + +- making ingestion asynchronous +- changing storage schemas or public APIs +- changing provider token-accounting rules +- introducing a background polling thread + +## Constraints and affected flow + +The affected path is `BatchRunner.start()` / resume -> checkpointed ingestion -> +chunk persistence -> submission. Reconciliation must happen only at durable chunk +boundaries, respect the provider polling interval, and preserve pause, cancellation, +retry-backoff, checkpoint, and fresh-process resume behavior. + +The implementation should keep orchestration in `runtime/ingestion.py` and use the +existing injected `poll_active_batches` callback. A monotonic clock may be injected +for deterministic cadence tests. Polling must occur before the chunk's submission +attempt when active batches exist and the cadence is due, so newly freed capacity +can be reused immediately. + +## Change list + +- Add cadence-aware active-batch polling at durable ingestion chunk boundaries. +- Add regression tests that reproduce a capacity-blocked multi-chunk ingest and + prove terminal reconciliation permits another submission before ingestion ends. +- Cover no-active-batch, cadence, pause/cancel, and retry-backoff safety as needed. +- Update README and runtime/provider/storage design documentation. +- Update the package and lock metadata from `0.0.5` to `0.0.6`. + +## Acceptance criteria and evidence + +1. During a multi-chunk ingest with active batches, a due reconciliation occurs + before the next submission attempt. Evidence: focused ingestion unit test. +2. A remotely completed batch releases locally counted enqueue tokens and another + pending batch can be submitted before all source rows are ingested. Evidence: + fake-provider integration regression test on head, plus a failing BASE + reproduction where practical. +3. Poll cadence avoids a provider request per fast chunk and uses monotonic time + while preserving the existing positive-interval provider configuration + contract. Evidence: deterministic cadence unit tests. +4. Poll-induced pause, cancellation, or retry backoff prevents unsafe subsequent + submission/materialization behavior and checkpoint completion remains correct. + Evidence: focused control/backoff regression tests and existing suite. +5. Public behavior and the `0.0.6` release are documented without a storage schema + change. Evidence: README/design-doc diff, strict MkDocs build, package build. +6. Coverage remains at least 85%, types pass, and the fake-provider/runtime suite + remains green. Evidence: required local smoke and GitHub CI on the immutable PR + head. + +## Environment contract + +All deterministic proof uses the in-memory or SQLite stores and fake providers; no +real provider credentials or Postgres instance are required. GitHub CI supplies the +ephemeral Postgres contract job. A live OpenAI smoke is post-merge only when the +documented credentials are available. + +## Local smoke decision + +`local_smoke: skipped` for paid live-provider execution because all fix criteria are +directly reproducible with deterministic fake-provider integration tests. The +repository-mandated pre-push smoke remains required: `uv run ty check src`, focused +runtime tests, `uv run pytest -q`, `uv run mkdocs build --strict`, and `uv build`. + +## Review and repair + +High-risk concurrency/run-lifecycle review is required. Repair cap: 3 cycles. diff --git a/docs/design_docs/OPENAI_BATCHING.md b/docs/design_docs/OPENAI_BATCHING.md index 2d338d3..33c89d4 100644 --- a/docs/design_docs/OPENAI_BATCHING.md +++ b/docs/design_docs/OPENAI_BATCHING.md @@ -57,7 +57,7 @@ Before request artifacts exist, built-in deterministic sources can resume ingest That currently includes CSV, JSONL, Parquet, and `CompositeItemSource` for ordered composition of checkpointed sources. When `CompositeItemSource` is used, child-source row IDs are auto-namespaced into run-unique `item_id` values, while the original row ID is preserved in lineage metadata. Custom non-file sources must implement a durable checkpoint contract explicitly; arbitrary iterables and live DB cursors are still `TBD`. -Before resumed ingestion continues, active OpenAI batches are polled once. If that poll consumes terminal batches or records enqueue/backoff failures, those state changes are applied before any new prompt rendering, request replay, or submission happens. +Before resumed ingestion continues, active OpenAI batches are polled once. If that poll consumes terminal batches or records enqueue/backoff failures, those state changes are applied before any new prompt rendering, request replay, or submission happens. Long-running ingestion then continues reconciliation at durable item-chunk boundaries whenever `poll_interval_sec` is due on a monotonic clock. The poll precedes the chunk's submission attempt, so completed OpenAI batches release their locally counted enqueue tokens immediately enough for the same boundary to submit more work. Fast chunks inside the positive interval do not trigger provider requests. The execution cycle reports durable progress explicitly to `Run.wait()`, so the wait loop immediately drains productive poll/submission cycles and sleeps only when a cycle makes no durable progress, respecting any persisted retry backoff. diff --git a/docs/design_docs/STORAGE_AND_RUNS.md b/docs/design_docs/STORAGE_AND_RUNS.md index d034fc8..c03ae89 100644 --- a/docs/design_docs/STORAGE_AND_RUNS.md +++ b/docs/design_docs/STORAGE_AND_RUNS.md @@ -181,6 +181,8 @@ This recovery happens when an execution owner advances the run, not when a caller merely reads a handle with `get_run()` or changes control state. When `start(job, run_id=...)` resumes a run with active provider batches, it first performs a poll-only reconciliation pass. Terminal provider batches are consumed, failed batches update retry backoff, and any active backoff prevents new source materialization or submission until the backoff expires. +While checkpointed ingestion is active, every materialized chunk and its next source checkpoint are persisted atomically before reconciliation is considered. At that durable boundary, active provider batches are polled before submission only when the provider's monotonic `poll_interval_sec` cadence is due. Terminal consumption therefore releases local inflight-token accounting before the same boundary submits pending work. A poll-induced pause, cancellation, or batch retry backoff blocks that submission and leaves the checkpoint incomplete at the last durable position. This scheduling state is process-local and requires no storage schema change; fresh-process resume still begins with its unconditional reconciliation pass. + Resume compatibility intentionally ignores non-persisted secret fields such as provider API keys. Rehydrated OpenAI runs need `OPENAI_API_KEY` when no explicit in-memory config is supplied. Gemini Developer API runs need `GEMINI_API_KEY`; Vertex runs need Application Default Credentials plus the persisted or ambient project, location, and GCS prefix configuration. For deterministic-source resume, the caller must also reuse the same `run_id` and provide the same source identity/fingerprint. diff --git a/docs/smoke-test.md b/docs/smoke-test.md index cdd5d50..649baaa 100644 --- a/docs/smoke-test.md +++ b/docs/smoke-test.md @@ -78,6 +78,7 @@ Use this when touching provider wiring, storage wiring, token budgeting, or run uv run ty check src uv run pytest tests/unit/test_batchor_tokens.py tests/unit/test_batchor_sqlite_storage_flow.py tests/unit/test_batchor_validation.py --no-cov -q uv run pytest tests/unit/test_batchor_artifacts.py tests/unit/test_batchor_storage_contracts.py --no-cov -q +uv run pytest tests/unit/test_batchor_runtime_ingestion.py --no-cov -q uv run pytest tests/integration/test_batchor_runner.py --no-cov -q uv run pytest tests/unit/test_batchor_gemini_provider.py tests/integration/test_batchor_gemini_runner.py --no-cov -q uv run pytest tests/unit/test_batchor_anthropic_provider.py --no-cov -q @@ -94,6 +95,8 @@ Expected: - Parquet source adapters can resume from opaque checkpoints and project only required columns - retry/resume from persisted request artifacts still works for SQLite-backed runs - resumed runs reconcile existing active batches and persisted backoff before materializing more source rows +- long-running ingestion polls active batches on a monotonic cadence at durable chunk boundaries, before submission, so completed remote work frees local enqueue-token capacity before full source materialization +- poll-induced pause, cancellation, or retry backoff leaves checkpointed ingestion incomplete at its last durable chunk and prevents unsafe submission - transient batch-poll failures do not block unrelated pending submissions from being sent when capacity remains - paused runs stop polling/submission until resumed - OpenAI control-plane and batch-level insufficient-quota provider failures auto-pause with a durable `control_reason` and do not consume item attempts diff --git a/pyproject.toml b/pyproject.toml index 1e4c272..54cf01a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "batchor" -version = "0.0.5" +version = "0.0.6" description = "Structured-first provider Batch runner with typed Pydantic results." readme = "README.md" license = { file = "LICENSE" } diff --git a/src/batchor/runtime/ingestion.py b/src/batchor/runtime/ingestion.py index c8cff7c..e99d077 100644 --- a/src/batchor/runtime/ingestion.py +++ b/src/batchor/runtime/ingestion.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import time from dataclasses import dataclass from typing import Any, Callable, Iterator, cast @@ -34,7 +35,9 @@ class IngestionDeps: emit_event: Observer callback used for coarse lifecycle events. submit_pending_items: Callback used to submit items during ingestion. poll_active_batches: Callback used to reconcile already-submitted - batches before resume materializes or submits new local work. + batches before resume or durable ingestion chunks submit new local + work. + monotonic: Monotonic clock used to enforce ingestion polling cadence. configs_match_for_resume: Predicate used to validate resume compatibility. """ @@ -44,6 +47,7 @@ class IngestionDeps: submit_pending_items: Callable[[str, RunContext], int] configs_match_for_resume: Callable[[PersistedRunConfig, PersistedRunConfig], bool] poll_active_batches: Callable[[str, RunContext], None] = _noop_poll_active_batches + monotonic: Callable[[], float] = time.monotonic def finalize_cancelled_ingestion(state: StateStore, *, run_id: str) -> None: @@ -149,6 +153,7 @@ def ingest_job_items( next_checkpoint_payload = checkpoint_payload checkpointed = checkpointed_source(job) is not None ingestion_complete = True + last_poll_at = deps.monotonic() for item_chunk, next_item_index, next_checkpoint_payload in materialize_item_chunks( job, start_index=start_index, @@ -177,6 +182,24 @@ def ingest_job_items( if control_state is RunControlState.CANCEL_REQUESTED: ingestion_complete = False break + if ( + control_state is RunControlState.RUNNING + and deps.state.get_batch_retry_backoff_remaining_sec(run_id=run_id) <= 0 + and deps.state.get_active_batches(run_id=run_id) + ): + now = deps.monotonic() + poll_interval_sec = context.config.provider_config.poll_interval_sec + if now - last_poll_at >= poll_interval_sec: + deps.poll_active_batches(run_id, context) + last_poll_at = now + control_state = deps.state.get_run_control_state(run_id=run_id) + backoff_remaining_sec = deps.state.get_batch_retry_backoff_remaining_sec(run_id=run_id) + if control_state is RunControlState.CANCEL_REQUESTED: + ingestion_complete = False + break + if checkpointed and (control_state is not RunControlState.RUNNING or backoff_remaining_sec > 0): + ingestion_complete = False + break deps.submit_pending_items(run_id, context) control_state = deps.state.get_run_control_state(run_id=run_id) if control_state is RunControlState.CANCEL_REQUESTED: diff --git a/tests/integration/test_batchor_runner.py b/tests/integration/test_batchor_runner.py index 2515743..e29897b 100644 --- a/tests/integration/test_batchor_runner.py +++ b/tests/integration/test_batchor_runner.py @@ -5,6 +5,7 @@ import subprocess import sys import textwrap +from dataclasses import replace from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Callable @@ -59,6 +60,14 @@ def sleep(self, seconds: float) -> None: self.current += timedelta(seconds=float(seconds)) +class _IngestionMonotonicClock: + def __init__(self, *readings: float) -> None: + self._readings = iter(readings) + + def __call__(self) -> float: + return next(self._readings) + + def _success_record(text: str) -> dict[str, object]: return { "custom_id": "", @@ -201,6 +210,23 @@ def estimate_request_tokens( return max(len(prompt), 1) +class _IngestionAwareBatchProvider(_FakeBatchProvider): + def __init__( + self, + *, + record_factory: Callable[[str], dict[str, object] | None], + materialized_count: Callable[[], int], + ) -> None: + super().__init__(record_factory=record_factory) + self.materialized_count = materialized_count + self.materialized_counts_at_create: list[int] = [] + + def create_batch(self, *, input_file_id: str, metadata: dict[str, str] | None = None) -> dict[str, object]: + batch = super().create_batch(input_file_id=input_file_id, metadata=metadata) + self.materialized_counts_at_create.append(self.materialized_count()) + return batch + + class _ArtifactOnlyBatchProvider(_FakeBatchProvider): def build_request_line( self, @@ -466,6 +492,54 @@ def test_transient_poll_failures_do_not_block_new_submissions(tmp_path: Path) -> assert summary.status_counts[ItemStatus.SUBMITTED] == 2 +def test_ingestion_poll_releases_capacity_before_source_is_fully_materialized( + tmp_path: Path, +) -> None: + materialized_count = 0 + + def build_prompt(item: BatchItem[str]) -> PromptParts: + nonlocal materialized_count + materialized_count += 1 + return PromptParts(prompt=item.payload) + + provider = _IngestionAwareBatchProvider( + record_factory=lambda custom_id: _success_record(f"text:{custom_id}"), + materialized_count=lambda: materialized_count, + ) + runner = BatchRunner( + storage="memory", + provider_factory=lambda _cfg: provider, + temp_root=tmp_path, + ) + runner._ingestion_deps = replace( + runner._ingestion_deps, + monotonic=_IngestionMonotonicClock(0.0, 1.0, 2.0), + ) + + runner.start( + BatchJob( + items=[BatchItem(item_id=f"row{index}", payload="x") for index in range(2001)], + build_prompt=build_prompt, + provider_config=OpenAIProviderConfig( + api_key="k", + model="gpt-4.1", + poll_interval_sec=1, + enqueue_limits=OpenAIEnqueueLimitConfig( + enqueued_token_limit=1000, + target_ratio=1.0, + headroom=0, + max_batch_enqueued_tokens=1000, + ), + ), + chunk_policy=ChunkPolicy(max_requests=1000, max_file_bytes=1024 * 1024), + ) + ) + + assert provider.created_batches == ["batch_0", "batch_1", "batch_2"] + assert provider.materialized_counts_at_create == [1000, 2000, 2001] + assert provider.materialized_counts_at_create[1] < 2001 + + def test_wait_keeps_draining_after_progress_without_idle_sleep(tmp_path: Path) -> None: clock = _FakeClock() provider = _FakeBatchProvider(record_factory=lambda custom_id: _success_record(f"text:{custom_id}")) diff --git a/tests/unit/test_batchor_runtime_ingestion.py b/tests/unit/test_batchor_runtime_ingestion.py index ad756bf..83f3927 100644 --- a/tests/unit/test_batchor_runtime_ingestion.py +++ b/tests/unit/test_batchor_runtime_ingestion.py @@ -13,11 +13,13 @@ MemoryStateStore, OpenAIProviderConfig, PromptParts, + RunControlState, ) from batchor.core.types import JSONValue from batchor.runtime.context import build_persisted_config, build_run_context from batchor.runtime.ingestion import ( IngestionDeps, + ingest_job_items, materialize_item_chunks, resume_existing_run, validate_checkpoint_source, @@ -30,6 +32,14 @@ class _NoopProvider: pass +class _MonotonicClock: + def __init__(self, *readings: float) -> None: + self._readings = iter(readings) + + def __call__(self) -> float: + return next(self._readings) + + class _CompleteCheckpointSource(CheckpointedItemSource[dict[str, str]]): def source_identity(self) -> SourceIdentity: return SourceIdentity( @@ -329,6 +339,249 @@ def poll_and_backoff(run_id: str, context: object) -> None: assert checkpoint.ingestion_complete is False +def test_ingest_job_items_polls_on_cadence_before_chunk_submission() -> None: + job = BatchJob( + items=[BatchItem(item_id=f"row{index}", payload="x") for index in range(2001)], + build_prompt=lambda item: PromptParts(prompt=item.payload), + provider_config=OpenAIProviderConfig( + api_key="k", + model="gpt-4.1", + poll_interval_sec=5.0, + ), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + run_id = "cadenced_ingestion_run" + storage.create_run(run_id=run_id, config=config, items=[]) + storage.register_batch( + run_id=run_id, + local_batch_id="local_1", + provider_batch_id="provider_1", + status="submitted", + custom_ids=[], + ) + calls: list[str] = [] + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda _run_id, _context: calls.append("submit") or 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=lambda _run_id, _context: calls.append("poll"), + monotonic=_MonotonicClock(0.0, 1.0, 6.0, 7.0), + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + + ingest_job_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=0, + checkpoint_payload=None, + ) + + assert calls == ["submit", "poll", "submit", "submit"] + + +def test_ingest_job_items_does_not_poll_without_active_batches() -> None: + job = BatchJob( + items=[BatchItem(item_id=f"row{index}", payload="x") for index in range(1001)], + build_prompt=lambda item: PromptParts(prompt=item.payload), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + run_id = "no_active_ingestion_run" + storage.create_run(run_id=run_id, config=config, items=[]) + poll_calls: list[str] = [] + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda _run_id, _context: 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=lambda polled_run_id, _context: poll_calls.append(polled_run_id), + monotonic=_MonotonicClock(0.0), + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + + ingest_job_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=0, + checkpoint_payload=None, + ) + + assert poll_calls == [] + + +@pytest.mark.parametrize( + "poll_control_state", + [RunControlState.PAUSED, RunControlState.CANCEL_REQUESTED], +) +def test_ingest_job_items_stops_checkpointed_ingestion_when_poll_changes_control( + tmp_path: Path, + poll_control_state: RunControlState, +) -> None: + path = tmp_path / "items.jsonl" + path.write_text( + "\n".join(json.dumps({"id": f"row{index}", "text": "x"}) for index in range(1500)) + "\n", + encoding="utf-8", + ) + source = JsonlItemSource( + path, + item_id_from_row=lambda row: str(row["id"]) if isinstance(row, dict) else "", + payload_from_row=lambda row: {"text": row["text"]} if isinstance(row, dict) else {}, + ) + job = BatchJob( + items=source, + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1", poll_interval_sec=1), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + run_id = f"poll_{poll_control_state.value}_run" + storage.create_run(run_id=run_id, config=config, items=[]) + identity = source.source_identity() + storage.set_ingest_checkpoint( + run_id=run_id, + checkpoint=IngestCheckpoint( + source_kind=identity.source_kind, + source_ref=identity.source_ref, + source_fingerprint=identity.source_fingerprint, + checkpoint_payload=source.initial_checkpoint(), + ), + ) + storage.register_batch( + run_id=run_id, + local_batch_id="local_1", + provider_batch_id="provider_1", + status="submitted", + custom_ids=[], + ) + submitted_runs: list[str] = [] + + def poll_and_change_control(polled_run_id: str, _context: object) -> None: + storage.set_run_control_state(run_id=polled_run_id, control_state=poll_control_state) + + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda submitted_run_id, _context: submitted_runs.append(submitted_run_id) or 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=poll_and_change_control, + monotonic=_MonotonicClock(0.0, 1.0), + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + + ingest_job_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=0, + checkpoint_payload=source.initial_checkpoint(), + ) + + checkpoint = storage.get_ingest_checkpoint(run_id=run_id) + assert checkpoint is not None + assert checkpoint.next_item_index == 1000 + assert checkpoint.ingestion_complete is False + assert submitted_runs == [] + + +def test_ingest_job_items_stops_checkpointed_ingestion_when_poll_records_backoff( + tmp_path: Path, +) -> None: + path = tmp_path / "items.jsonl" + path.write_text( + "\n".join(json.dumps({"id": f"row{index}", "text": "x"}) for index in range(1500)) + "\n", + encoding="utf-8", + ) + source = JsonlItemSource( + path, + item_id_from_row=lambda row: str(row["id"]) if isinstance(row, dict) else "", + payload_from_row=lambda row: {"text": row["text"]} if isinstance(row, dict) else {}, + ) + job = BatchJob( + items=source, + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1", poll_interval_sec=1), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + run_id = "poll_backoff_run" + storage.create_run(run_id=run_id, config=config, items=[]) + identity = source.source_identity() + storage.set_ingest_checkpoint( + run_id=run_id, + checkpoint=IngestCheckpoint( + source_kind=identity.source_kind, + source_ref=identity.source_ref, + source_fingerprint=identity.source_fingerprint, + checkpoint_payload=source.initial_checkpoint(), + ), + ) + storage.register_batch( + run_id=run_id, + local_batch_id="local_1", + provider_batch_id="provider_1", + status="submitted", + custom_ids=[], + ) + submitted_runs: list[str] = [] + + def poll_and_backoff(polled_run_id: str, _context: object) -> None: + storage.record_batch_retry_failure( + run_id=polled_run_id, + error_class="provider_unavailable", + base_delay_sec=10.0, + max_delay_sec=10.0, + ) + + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda submitted_run_id, _context: submitted_runs.append(submitted_run_id) or 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=poll_and_backoff, + monotonic=_MonotonicClock(0.0, 1.0), + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + + ingest_job_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=0, + checkpoint_payload=source.initial_checkpoint(), + ) + + checkpoint = storage.get_ingest_checkpoint(run_id=run_id) + assert checkpoint is not None + assert checkpoint.next_item_index == 1000 + assert checkpoint.ingestion_complete is False + assert submitted_runs == [] + + # --------------------------------------------------------------------------- # resume_existing_run: cancel and pause paths # --------------------------------------------------------------------------- diff --git a/uv.lock b/uv.lock index aa02fb7..241b67e 100644 --- a/uv.lock +++ b/uv.lock @@ -81,7 +81,7 @@ wheels = [ [[package]] name = "batchor" -version = "0.0.5" +version = "0.0.6" source = { editable = "." } dependencies = [ { name = "openai" },