From e642d0d01f5424880ad9c4234c8a63c6d9ae8e4c Mon Sep 17 00:00:00 2001 From: YW <58594437+AnsonDev42@users.noreply.github.com> Date: Sat, 18 Jul 2026 22:44:48 +0100 Subject: [PATCH 1/3] fix(runtime): harden durable run lifecycle - Keep terminal work replayable and submissions explicitly resolvable\n across crash windows.\n- Migrate durable stores from schema 4 to 6 with conservative generic\n ingest markers and additive capacity reservation tables. --- README.md | 9 +- docs/design_docs/ARCHITECTURE.md | 10 +- docs/design_docs/OPENAI_BATCHING.md | 11 +- docs/design_docs/RUNTIME_DURABILITY_0_0_7.md | 77 ++++ docs/design_docs/STORAGE_AND_RUNS.md | 16 +- docs/design_docs/STORAGE_MIGRATIONS.md | 50 ++- docs/doc-map.md | 1 + docs/smoke-test.md | 10 + mkdocs.yml | 1 + pyproject.toml | 2 +- src/batchor/__init__.py | 2 + src/batchor/core/exceptions.py | 11 + src/batchor/core/models.py | 19 + src/batchor/runtime/execution.py | 23 +- src/batchor/runtime/ingestion.py | 405 ++++++++++++++--- src/batchor/runtime/polling.py | 107 ++++- src/batchor/runtime/run_handle.py | 34 +- src/batchor/runtime/runner.py | 26 +- src/batchor/runtime/submission.py | 94 ++-- src/batchor/storage/memory.py | 200 ++++++++- src/batchor/storage/postgres_store.py | 1 + src/batchor/storage/sqlite_lifecycle.py | 417 +++++++++++++++++- src/batchor/storage/sqlite_queries.py | 60 ++- src/batchor/storage/sqlite_results.py | 8 +- src/batchor/storage/sqlite_schema.py | 35 +- src/batchor/storage/state_models.py | 59 +++ tests/integration/test_batchor_runner.py | 45 +- tests/unit/test_batchor_runtime_ingestion.py | 145 ++++++ tests/unit/test_batchor_runtime_polling.py | 176 ++++++++ tests/unit/test_batchor_runtime_submission.py | 18 +- .../unit/test_batchor_sqlite_storage_flow.py | 71 +++ tests/unit/test_batchor_storage.py | 2 +- tests/unit/test_batchor_storage_contracts.py | 94 ++++ uv.lock | 2 +- 34 files changed, 2057 insertions(+), 184 deletions(-) create mode 100644 docs/design_docs/RUNTIME_DURABILITY_0_0_7.md diff --git a/README.md b/README.md index ac91411..00a5048 100644 --- a/README.md +++ b/README.md @@ -79,8 +79,8 @@ 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. 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 +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 work-slice boundaries when the configured provider polling interval is due. Submission backoff never suppresses reconciliation. Fast sources retain the normal 1,000-item persistence path, while slow sources flush partial atomic slices on a time budget so pause, cancellation, polling, and deadlines can be observed. Provider-terminal batches remain locally consumable until every linked item transition and capacity release succeeds, making result consumption replayable after a crash. +Every newly created run has an incomplete ingestion marker, so even an empty run cannot appear terminal while its first source slice is still being rendered. 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 raises `RunIngestionSourceRequiredError` instead of silently completing only the @@ -329,6 +329,7 @@ run = runner.start( target_ratio=0.7, headroom=50_000, max_batch_enqueued_tokens=500_000, + quota_scope="research-account-a", ), ), ) @@ -338,7 +339,7 @@ run.wait() print(run.results()[0].output) ``` -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. +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. OpenAI capacity reservations are atomic and shared across runs by model and `quota_scope`; omit the scope for a conservative store-wide default, or assign stable labels when one store serves multiple accounts/projects. 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`. @@ -583,6 +584,8 @@ Current events include run creation/resume, automatic quota pause, item ingestio - Item `attempt_count` means consumed provider attempts: successful submitted completions increment it, counted item-level failures increment it, and local or batch-level retry/reset paths that did not consume an item attempt leave it unchanged. - Durable artifacts now go through an `ArtifactStore` seam. The built-in implementation is `LocalArtifactStore`, intended for local disk or a shared mounted volume. - Fresh-process resume requeues any locally claimed but not yet submitted items before continuing, so a process crash after request-artifact persistence does not strand work in `queued_local`. +- Batch creation uses a durable submission intent. If the process or connection fails after the provider may have created remote work, advancement raises `RunSubmissionIndeterminateError` instead of silently duplicating execution. After inspecting the provider, resolve the intent with `resolve_indeterminate_submission_as_created(provider_batch_id=..., status=...)` or `resolve_indeterminate_submission_as_not_created()`. +- `Run.wait(timeout=...)` propagates its deadline through library-controlled ingestion and polling boundaries and never starts another provider operation after expiry. An already-running arbitrary Python callback or provider SDK call remains subject to that operation's own timeout and cannot be forcibly interrupted safely. ## Durable artifacts diff --git a/docs/design_docs/ARCHITECTURE.md b/docs/design_docs/ARCHITECTURE.md index 2eb391a..bd93216 100644 --- a/docs/design_docs/ARCHITECTURE.md +++ b/docs/design_docs/ARCHITECTURE.md @@ -145,16 +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. Persist source items in durable chunks; at each chunk boundary, poll active batches before submission when the monotonic provider cadence is due. +4. Persist source items in durable work slices; keep the 1,000-item fast path, but flush a smaller atomic slice when the time/control budget requires it. 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. +8. Atomically reserve quota-scoped capacity and persist a submission intent before remote creation, then atomically register the returned batch with its item linkage. 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. +Ingestion polling stays inside the ingestion module behind its existing poll callback. Provider cadence and work-slice clocks are separate: fast sources retain normal batch granularity, while slow sources reach durable cooperative boundaries on a time budget. After each ingestion poll, control state and retry backoff are re-read before submission or further checkpointed-source materialization. Submission backoff blocks materialization/submission but never active-batch reconciliation. 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. @@ -473,6 +473,10 @@ That split gives `batchor`: can advance, unless cancellation intentionally abandons the unmaterialized source tail and finalizes the checkpoint. 21. `BatchItem.metadata["batchor_lineage"]` is reserved for lightweight source/join metadata when provided by built-in adapters or callers. +22. Provider-terminal status and local result consumption are distinct: a terminal batch remains locally active while submitted items or an unreleased capacity reservation still require replay. +23. Remote batch creation is preceded by a durable submission intent. An ambiguous create outcome requires explicit operator resolution as either the known remote batch or verified-not-created; the runtime never guesses and resubmits. +24. OpenAI enqueue capacity is reserved atomically across runs by a stable provider/model/quota scope without persisting credential-derived identities. +25. `wait(timeout=...)` is cooperative: the deadline bounds sleeps and prevents new library-controlled work, while an already-running user callback or provider SDK call is governed by its own timeout. ## Extension seams diff --git a/docs/design_docs/OPENAI_BATCHING.md b/docs/design_docs/OPENAI_BATCHING.md index 33c89d4..27c8d12 100644 --- a/docs/design_docs/OPENAI_BATCHING.md +++ b/docs/design_docs/OPENAI_BATCHING.md @@ -125,13 +125,14 @@ OpenAI-specific enqueue settings live on `OpenAIProviderConfig.enqueue_limits`: - `target_ratio` - `headroom` - `max_batch_enqueued_tokens` +- `quota_scope` From those settings, `batchor` derives: - an effective inflight token budget - an effective per-batch token limit -Submission is constrained by both token budget and generic chunking rules such as max request count and max request file size. +Submission is constrained by both token budget and generic chunking rules such as max request count and max request file size. Capacity is reserved atomically in durable storage before upload/create, so concurrent runs sharing the same model and quota scope cannot each claim the full account budget. The default scope is conservative and store-wide per model; set a stable `quota_scope` when one store serves multiple OpenAI accounts or projects. During `Run.wait()`, each drain cycle polls active batches before submitting new work, so capacity freed by completed batches is visible before the next submit attempt. If that cycle makes durable progress, the wait loop immediately runs another cycle instead of sleeping. ## Batch splitting @@ -155,7 +156,8 @@ Not all failures are treated the same way. Control-plane failures: -- transient upload/create/poll failures are treated as retryable batch-control-plane failures +- transient upload/poll failures are treated as retryable batch-control-plane failures +- a create-call outcome that may have reached the provider is recorded as an indeterminate durable intent and requires explicit operator resolution instead of automatic resubmission - retryable control-plane failures do not consume item attempts - batch submit failures can trigger batch-level backoff - transient poll failures do not stall unrelated submissions when other capacity remains @@ -174,9 +176,10 @@ Item-level failures: Cleanup behavior: -- if upload succeeds but batch creation fails, `batchor` makes a best-effort attempt to delete the uploaded OpenAI input file +- if upload succeeds but a known provider rejection proves batch creation did not occur, `batchor` makes a best-effort attempt to delete the uploaded OpenAI input file - if upload or batch creation fails because OpenAI reports insufficient quota, local queued items are released back to pending before the run is paused -- if a process dies after local artifact persistence but before durable batch registration, fresh-process resume requeues those items and resubmits from persisted request artifacts +- if creation may have succeeded but local registration did not, `RunSubmissionIndeterminateError` blocks duplicate execution until the operator attaches the confirmed remote batch with `resolve_indeterminate_submission_as_created(...)` or verifies absence with `resolve_indeterminate_submission_as_not_created()` +- terminal result download, parsing, item persistence, and capacity release are replayable; the batch remains locally reconcilable until all linked work is durably consumed ## Run control diff --git a/docs/design_docs/RUNTIME_DURABILITY_0_0_7.md b/docs/design_docs/RUNTIME_DURABILITY_0_0_7.md new file mode 100644 index 0000000..1e045bc --- /dev/null +++ b/docs/design_docs/RUNTIME_DURABILITY_0_0_7.md @@ -0,0 +1,77 @@ +# Runtime Durability 0.0.7 + +Status: released design and verification record for `0.0.7`. + +## Outcome + +`0.0.7` closes durability and liveness gaps in the run lifecycle without +changing the provider request formats or introducing background execution. The +release makes remote terminal consumption replayable, keeps reconciliation live +during ingestion and submit backoff, records submission ambiguity durably, and +enforces OpenAI enqueue capacity across runs sharing a quota scope. + +## Design boundaries + +- Batchor owns durable local reconciliation; providers remain the authority for + remote batch status and result payloads. +- Result consumption is at-least-once/replayable locally. A terminal remote + batch remains active until submitted items are terminalized and its capacity + reservation is released. +- Generic provider creation cannot promise exactly-once remote execution. A + timeout or crash after the create call becomes an explicit + `RunSubmissionIndeterminateError`, requiring operator resolution as created + or not created rather than an unsafe automatic retry. +- Cooperative deadlines prevent starting another Batchor-controlled operation + after expiry and cap sleeps. They cannot interrupt an arbitrary prompt + callback or provider SDK call already running in the current thread; those + calls retain their own timeout contract. +- Non-checkpointable iterators are retained only by their in-process execution + owner. Manual pause can stop and resume that iterator in the same process, + but fresh-process mid-ingestion recovery remains unsupported. +- Capacity is shared only among runs using the same durable store and explicit + quota scope. The default scope is conservative for one account/model; users + sharing a store across accounts or projects must configure distinct scopes. + +## Change list and acceptance criteria + +| Finding | Design/fix boundary | Acceptance evidence | +| --- | --- | --- | +| Terminal batch lost after a local failure | Terminal batches with submitted items or unreleased intent capacity remain pollable. Consumption filters already-terminal rows and persists item changes conditionally so it can replay after download, parse, or partial-write failure. | Fake-provider terminal-replay and partial-mutation tests in polling/storage suites. | +| Submit backoff suppresses polling | Submission backoff blocks materialization/submission only; active remote batches are reconciled at each due work boundary. | Ingestion regression with pre-existing submit backoff and a completed active batch. | +| A slow 1,000-row slice blocks progress | Materialization uses a normal 1,000-item fast path plus a cooperative time budget that flushes partial atomic slices. | Slow-source clock tests proving polling/control observation before another full chunk. | +| Empty non-checkpointed run appears terminal | Every new run receives an incomplete generic ingest marker atomically; v5 backfill treats unknown legacy state as incomplete. | Fresh SQLite rehydration and legacy migration tests. | +| Remote create and local linkage can split | Persist a submission intent before create; atomically finalize intent, batch, and item linkage after confirmed success. | Crash-window/indeterminate-resolution storage and submission tests. | +| Manual pause drains arbitrary ingestion | Manual pause stops at a cooperative boundary and retains the iterator for same-process resume; quota auto-pause still drains finite arbitrary input. | Non-checkpointed pause/resume and quota-pause tests. | +| OpenAI capacity is per-run | Atomically reserve tokens in a durable provider/model/quota-scope counter and include active legacy submitted rows. | Multi-run scope and legacy-reservation tests. | +| A success clears another failure's backoff | Polling aggregates its cycle before retry-state mutation; a completed batch does not clear unrelated retry delay. | Mixed failed/completed batch backoff tests. | +| `wait(timeout)` overruns | Thread an absolute deadline through refresh, ingestion, polling, and sleep decisions. | Blocked-ingestion/provider and sleep deadline tests. | +| Cancellation waits only on stale backoff | Cancellation clears retry backoff after its final active batch drains. | Cancelled terminal quota/error regression test. | + +## Storage migration contract + +The schema advances from 4 to 6. Version 5 establishes the generic +`run_ingest_state` invariant and conservative marker backfill. Version 6 adds +`submission_intents` and `capacity_scopes`. Both upgrades are additive, +performed during storage startup, and retain all existing rows. See +[`STORAGE_MIGRATIONS.md`](STORAGE_MIGRATIONS.md) for operator guidance. + +## Non-goals + +- Background polling, asynchronous ingestion, or forcibly cancelling arbitrary + user/provider calls. +- Provider-side remote cancellation guarantees. +- Fresh-process resume for arbitrary non-checkpointable iterables. +- Universal provider idempotency/discovery for `create_batch`. +- Cross-store or credential-derived capacity coordination. + +## Evidence contract + +The deterministic evidence for this release is the fake-provider, memory, +SQLite, and storage-contract suite. It must retain at least 85% coverage and +pass type checking, linting, strict documentation build, package build, and +`git diff --check` on the release commit. + +No paid live-provider smoke is required before push: all release criteria are +directly represented by deterministic fake-provider/storage regressions. A live +provider smoke, when credentials are available, remains an operational +post-merge check rather than proof of these local state-machine invariants. diff --git a/docs/design_docs/STORAGE_AND_RUNS.md b/docs/design_docs/STORAGE_AND_RUNS.md index c03ae89..3dec0b5 100644 --- a/docs/design_docs/STORAGE_AND_RUNS.md +++ b/docs/design_docs/STORAGE_AND_RUNS.md @@ -155,7 +155,9 @@ There is still one ingest-checkpoint row per run, not one row per child source. This enables `start(job, run_id=...)` to resume ingestion rather than rematerializing already-ingested rows, as long as the source still matches the stored identity. -Non-checkpointable iterables do not yet support the same mid-ingest crash recovery behavior. +Every run receives a generic incomplete-ingestion row atomically with run creation, including arbitrary iterables. This prevents the empty-run terminal calculation from winning while the first slice is still being rendered. Non-checkpointable iterables still do not support fresh-process mid-ingest recovery; manual pause retains their iterator only in the execution owner process and resumes it exactly once. + +Schema upgrade backfills every run missing a generic marker as ingestion-incomplete and records its current materialized item count. Storage cannot distinguish a completed arbitrary iterable from one paused between chunks, even when some items already exist, so assuming completion would preserve the false-terminal race. An empty source can be safely reattached from index zero; a partially materialized non-checkpointable legacy source must be treated as an abandoned tail and explicitly cancelled after any active batches drain. ## Rehydration rules @@ -181,7 +183,13 @@ 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. +While ingestion is active, the normal fast path persists 1,000 items at once. A separate work clock flushes a smaller atomic slice when slow source or prompt work crosses the cooperative time budget. At each durable boundary, active provider batches are polled when the provider cadence is due even if submission retry backoff is active. Poll-induced pause, cancellation, or retry backoff blocks further checkpointed materialization/submission and leaves the source incomplete at its last durable position. + +Provider batch creation uses a durable intent written before the remote call. Successful creation is finalized in one local transaction that registers the batch and links every submitted item. A crash or ambiguous connection failure between the remote call and finalization leaves an indeterminate intent; advancement raises `RunSubmissionIndeterminateError` until an operator either attaches the confirmed remote batch or verifies that no batch exists. This is the explicit boundary where a generic provider cannot supply exactly-once creation or discovery. + +Finalized intents also reserve OpenAI enqueue capacity in an atomic counter keyed by provider, model, and configured quota scope. The reservation spans runs and is released idempotently only after terminal batch consumption/reset succeeds. Provider-terminal batches with submitted items or unreleased capacity remain locally reconcilable, so download, parsing, partial item writes, and post-consumption crashes replay safely. + +Active submitted rows created before submission intents existed are included dynamically as legacy reservations when a scope admits new work. This prevents an upgrade from granting the full configured budget on top of already-enqueued provider batches. 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. @@ -263,7 +271,9 @@ Auto-pause must not overwrite `cancel_requested`. Cancellation is non-reversible For non-checkpointed finite iterables, ingestion continues materializing remaining items after an auto-pause so `resume()` cannot silently lose input rows that were not yet persisted. Checkpointed sources may stop on pause because their stored checkpoint can resume materialization later. -`wait()` fails fast on paused runs instead of sleeping indefinitely. The raised `RunPausedError` carries the same `control_reason` as the summary. +Manual pause stops both checkpointed and arbitrary ingestion at a cooperative boundary. Arbitrary iterators are retained only in-process; automatic quota pause still materializes their finite tail to avoid losing rows that cannot be reconstructed. + +`wait()` fails fast on paused runs instead of sleeping indefinitely. The raised `RunPausedError` carries the same `control_reason` as the summary. Its optional timeout is propagated through ingestion and polling, caps sleeps, and prevents starting another library-controlled operation after expiry. It cannot safely preempt an arbitrary callback or provider SDK call already executing in the current thread; those operations retain their own timeout contract. Provider-side remote cancellation is still `TBD`. ## Python API versus CLI diff --git a/docs/design_docs/STORAGE_MIGRATIONS.md b/docs/design_docs/STORAGE_MIGRATIONS.md index bb54bcc..b6df9b6 100644 --- a/docs/design_docs/STORAGE_MIGRATIONS.md +++ b/docs/design_docs/STORAGE_MIGRATIONS.md @@ -7,19 +7,65 @@ This document describes the current SQLite schema-versioning story for `batchor` - SQLite remains the default durable backend - the SQLite schema is additive and self-healing for supported columns/tables - storage metadata now persists a `schema_version` -- the current published schema version is `4` -- schema version `4` adds the nullable `runs.control_reason` column for durable pause reasons +- the current published schema version is `6` +- schema version `4` added the nullable `runs.control_reason` column for durable pause reasons +- schema version `5` makes the existing `run_ingest_state` row a generic, mandatory + incomplete-ingestion marker for every run; startup backfills one for every legacy + run that lacks it +- schema version `6` adds durable `submission_intents` and shared + `capacity_scopes` reservations for crash-safe submission reconciliation and + cross-run OpenAI enqueue accounting + +## Version 5: generic ingestion markers + +Version 5 changes the meaning of `run_ingest_state` from a checkpoint-only +record into a generic ingestion lifecycle marker. New runs create the marker in +the same transaction as the run row. This prevents an empty run from being +calculated as terminal while its first item slice is still being rendered. + +On startup, SQLite and Postgres backfill every run without a marker as +`ingestion_complete = 0`, with `next_item_index` set to its durable item count +and an unattached source identity. This is deliberately conservative: pre-v5 +state has no durable fact that distinguishes a completed arbitrary iterable +from an interrupted one. An empty source can be reattached at index zero; a +partially materialized non-checkpointable source must be cancelled/abandoned +explicitly after its active batches drain. Deterministic sources retain their +normal checkpoint-resume behavior. + +## Version 6: submission intents and capacity scopes + +Version 6 adds two additive tables: + +- `submission_intents` is written before a provider `create_batch` call. A + `creating` intent represents an ambiguous remote outcome and blocks automatic + retry until an operator resolves it as created or not created. A successful + create finalizes the batch registration and item linkage together. +- `capacity_scopes` holds an atomic OpenAI enqueue-token reservation keyed by + provider/model/quota scope. The reservation spans runs and is released only + after terminal consumption or reset succeeds. + +The tables are created by the normal additive metadata startup path. Existing +submitted rows that predate intents are counted dynamically when a scope admits +new work, so an upgrade cannot over-admit capacity on top of an active legacy +batch. No data is deleted or rewritten during either migration. ## Compatibility Rules - opening an existing SQLite database with the current package will create any missing additive tables or columns - `SQLiteStorage.schema_version` reports the effective schema version after startup checks +- opening an existing Postgres schema applies the same additive metadata tables + and generic-ingest-marker backfill - durable run data should be considered compatible only within the documented supported package line - secret provider fields are not part of durable schema compatibility ## Current Guidance - for local upgrades, open the database with the newer package before relying on it in automation +- before upgrading a partially ingested legacy arbitrary iterable, decide whether + to reattach an empty source or cancel/abandon the legacy tail; it is not safe + to infer completion from its pre-v5 rows +- inspect and resolve any `RunSubmissionIndeterminateError` rather than retrying + the run blindly after an ambiguous provider create - treat schema updates as one-way unless explicit downgrade guidance is published - keep artifact directories beside the same SQLite database during upgrades - if you need hard rollback guarantees, export terminal artifacts and results before changing package versions diff --git a/docs/doc-map.md b/docs/doc-map.md index f76141b..d8dbf87 100644 --- a/docs/doc-map.md +++ b/docs/doc-map.md @@ -23,6 +23,7 @@ This page explains what each document is for so readers do not have to guess whi | `design_docs/GEMINI_BATCHING.md` | Gemini text-only request construction, batch polling, response parsing, and current limits. | | `design_docs/STORAGE_AND_RUNS.md` | Durable `Run` lifecycle, rehydration, checkpoints, control state, artifact retention, and operator semantics. | | `design_docs/STORAGE_MIGRATIONS.md` | SQLite schema-versioning and migration guidance. | +| `design_docs/RUNTIME_DURABILITY_0_0_7.md` | Release design, boundaries, and verification contract for the 0.0.7 runtime-durability changes. | | `design_docs/ROADMAP.md` | Intentionally unimplemented areas and planned work. | ## Validation and project policy diff --git a/docs/smoke-test.md b/docs/smoke-test.md index 649baaa..2ee1558 100644 --- a/docs/smoke-test.md +++ b/docs/smoke-test.md @@ -91,14 +91,19 @@ Expected: - replaying multiple items from the same persisted request artifact does not require rereading that artifact file for each item - deterministic source ingestion can resume from a persisted checkpoint when rerun with the same `run_id` - incomplete checkpointed ingestion cannot become terminal, and fresh-process advancement requires reattaching the source with `start(job, run_id=...)` +- generic ingestion markers keep empty and partially materializing arbitrary-iterable runs non-terminal +- schema upgrade conservatively backfills missing empty and partially materialized ingestion markers - composite deterministic sources can namespace duplicate row IDs across explicit inputs and resume across source boundaries - 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 +- slow ingestion flushes bounded partial work slices while fast ingestion retains the normal 1,000-item batch granularity +- submission backoff blocks new materialization/submission without suppressing active-batch reconciliation - 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 +- manual pause stops a non-checkpointed iterator and same-process resume continues it exactly once - OpenAI control-plane and batch-level insufficient-quota provider failures auto-pause with a durable `control_reason` and do not consume item attempts - OpenAI row-level insufficient-quota records in completed batch output remain retryable, do not consume attempts, and back off without pausing the run - quota auto-pause preserves `cancel_requested` and does not strand non-checkpointed input rows during initial ingestion @@ -111,9 +116,14 @@ Expected: - shared storage-contract behavior remains aligned across SQLite and opt-in Postgres - completed submitted items report consumed attempts consistently across storage backends - OpenAI request splitting and enqueue-limit logic still behave as expected +- concurrent runs atomically share OpenAI capacity by quota scope, and terminal replay releases reservations idempotently +- pre-intent active batches continue consuming shared capacity after schema upgrade +- terminal batches replay after download/parse/partial-state crashes without duplicating consumed attempts +- ambiguous provider creation requires explicit created/not-created resolution and never silently submits a duplicate - Gemini text-only request construction, batch polling normalization, response parsing, and structured-output validation still behave as expected - Anthropic request construction, safe correlation IDs, polling normalization, result parsing, and structured-output validation still behave as expected - wait-mode refresh cycles keep draining immediately after poll/submission progress, without introducing idle poll sleeps while more local work can be sent +- wait deadlines cap sleeps and prevent new ingestion/poll/download work after expiry, subject to already-running callback/SDK operation timeouts - structured-output parsing remains stable Notes: diff --git a/mkdocs.yml b/mkdocs.yml index aff304e..8479d04 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -92,6 +92,7 @@ nav: - Anthropic Batching: design_docs/ANTHROPIC_BATCHING.md - Storage & Runs: design_docs/STORAGE_AND_RUNS.md - Storage Migrations: design_docs/STORAGE_MIGRATIONS.md + - Runtime Durability 0.0.7: design_docs/RUNTIME_DURABILITY_0_0_7.md - Roadmap: design_docs/ROADMAP.md - Project: - Support Policy: policies/support.md diff --git a/pyproject.toml b/pyproject.toml index 54cf01a..3ff83b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "batchor" -version = "0.0.6" +version = "0.0.7" description = "Structured-first provider Batch runner with typed Pydantic results." readme = "README.md" license = { file = "LICENSE" } diff --git a/src/batchor/__init__.py b/src/batchor/__init__.py index f759f94..73b8425 100644 --- a/src/batchor/__init__.py +++ b/src/batchor/__init__.py @@ -36,6 +36,7 @@ RunIngestionSourceRequiredError, RunNotFinishedError, RunPausedError, + RunSubmissionIndeterminateError, StructuredOutputSchemaError, ) from batchor.core.models import ( @@ -129,6 +130,7 @@ "RunEvent", "RunLifecycleStatus", "RunIngestionSourceRequiredError", + "RunSubmissionIndeterminateError", "RunNotFinishedError", "RunPausedError", "RunSnapshot", diff --git a/src/batchor/core/exceptions.py b/src/batchor/core/exceptions.py index 3c780fe..48ce13b 100644 --- a/src/batchor/core/exceptions.py +++ b/src/batchor/core/exceptions.py @@ -81,6 +81,17 @@ def __init__(self, run_id: str) -> None: self.run_id = run_id +class RunSubmissionIndeterminateError(RuntimeError): + """Raised when a provider batch-create call may have succeeded before a crash.""" + + def __init__(self, run_id: str) -> None: + super().__init__( + f"run {run_id} has an indeterminate provider batch creation; inspect the provider before " + "calling resolve_indeterminate_submission_as_not_created()" + ) + self.run_id = run_id + + class StructuredOutputSchemaError(ValueError): """Raised when a Pydantic model produces a schema incompatible with OpenAI. diff --git a/src/batchor/core/models.py b/src/batchor/core/models.py index e670914..bde02bc 100644 --- a/src/batchor/core/models.py +++ b/src/batchor/core/models.py @@ -157,12 +157,16 @@ class OpenAIEnqueueLimitConfig: computing the effective budget. Defaults to ``0``. max_batch_enqueued_tokens: Optional per-batch token ceiling. ``0`` means no per-batch limit beyond the inflight budget. + quota_scope: Optional stable account/project label. Set this when + multiple OpenAI accounts share one durable store; it is persisted + instead of deriving a scope from the credential digest. """ enqueued_token_limit: int = 0 target_ratio: float = 0.7 headroom: int = 0 max_batch_enqueued_tokens: int = 0 + quota_scope: str = "" def __post_init__(self) -> None: if self.enqueued_token_limit < 0: @@ -175,6 +179,8 @@ def __post_init__(self) -> None: raise ValueError("headroom must be < enqueued_token_limit") if self.max_batch_enqueued_tokens < 0: raise ValueError("max_batch_enqueued_tokens must be >= 0") + if not isinstance(self.quota_scope, str): + raise TypeError("quota_scope must be a string") if self.enqueued_token_limit > 0: effective_budget = min( int(self.enqueued_token_limit * self.target_ratio), @@ -202,6 +208,7 @@ def to_payload(self) -> JSONObject: "target_ratio": self.target_ratio, "headroom": self.headroom, "max_batch_enqueued_tokens": self.max_batch_enqueued_tokens, + "quota_scope": self.quota_scope, } @classmethod @@ -221,6 +228,7 @@ def from_payload(cls, payload: JSONObject) -> OpenAIEnqueueLimitConfig: target_ratio = payload.get("target_ratio", 0.7) headroom = payload.get("headroom", 0) max_batch_enqueued_tokens = payload.get("max_batch_enqueued_tokens", 0) + quota_scope = payload.get("quota_scope", "") if not isinstance(enqueued_token_limit, int): raise TypeError("enqueued_token_limit must be an int") if not isinstance(target_ratio, int | float): @@ -229,11 +237,14 @@ def from_payload(cls, payload: JSONObject) -> OpenAIEnqueueLimitConfig: raise TypeError("headroom must be an int") if not isinstance(max_batch_enqueued_tokens, int): raise TypeError("max_batch_enqueued_tokens must be an int") + if not isinstance(quota_scope, str): + raise TypeError("quota_scope must be a string") return cls( enqueued_token_limit=enqueued_token_limit, target_ratio=float(target_ratio), headroom=headroom, max_batch_enqueued_tokens=max_batch_enqueued_tokens, + quota_scope=quota_scope, ) @@ -551,6 +562,14 @@ def from_payload(cls, payload: JSONObject) -> OpenAIProviderConfig: ) +def openai_enqueue_quota_scope(provider_config: OpenAIProviderConfig) -> str: + """Return a stable persisted capacity scope without credential material.""" + configured_scope = provider_config.enqueue_limits.quota_scope.strip() + if configured_scope: + return f"openai:{provider_config.model}:label:{configured_scope}" + return f"openai:{provider_config.model}:default" + + @dataclass(frozen=True) class ChunkPolicy: """Submission chunking limits applied before provider batches are created. diff --git a/src/batchor/runtime/execution.py b/src/batchor/runtime/execution.py index 28e22e3..1c849ce 100644 --- a/src/batchor/runtime/execution.py +++ b/src/batchor/runtime/execution.py @@ -2,13 +2,14 @@ from __future__ import annotations +import time from dataclasses import dataclass from typing import Any, Callable from pydantic import BaseModel from batchor.core.enums import RunControlState -from batchor.core.exceptions import RunIngestionSourceRequiredError +from batchor.core.exceptions import RunIngestionSourceRequiredError, RunSubmissionIndeterminateError from batchor.core.models import BatchJob, RunSummary from batchor.runtime.context import RunContext, build_persisted_config from batchor.runtime.ingestion import IngestionDeps, resume_existing_run @@ -48,7 +49,7 @@ def start_or_attach(self, *, run_id: str, job: BatchJob[Any, BaseModel]) -> None """Attach the source-bearing job to an execution owner in this process.""" self._jobs[run_id] = job - def resume_attached_ingestion(self, run_id: str) -> None: + def resume_attached_ingestion(self, run_id: str, *, deadline: float | None = None) -> None: """Continue incomplete ingestion when this process still owns its source.""" job = self._jobs.get(run_id) checkpoint = self._state.get_ingest_checkpoint(run_id=run_id) @@ -62,6 +63,7 @@ def resume_attached_ingestion(self, run_id: str) -> None: job=job, config=build_persisted_config(job), context=self._context_for_run(run_id), + deadline=deadline, ) self._owners.add(run_id) @@ -71,9 +73,13 @@ def require_attached_source(self, run_id: str) -> None: if checkpoint is not None and not checkpoint.ingestion_complete and run_id not in self._jobs: raise RunIngestionSourceRequiredError(run_id) - def advance(self, run_id: str) -> CycleOutcome: + def advance(self, run_id: str, *, deadline: float | None = None) -> CycleOutcome: """Advance one run in recovery, poll, control, ingest, submit order.""" before = self._state.get_run_summary(run_id=run_id) + if _deadline_expired(deadline): + return CycleOutcome(summary=before, progressed=False) + if self._state.has_indeterminate_submission_intents(run_id=run_id): + raise RunSubmissionIndeterminateError(run_id) if run_id not in self._owners: self._state.requeue_local_items(run_id=run_id) self._owners.add(run_id) @@ -84,13 +90,18 @@ def advance(self, run_id: str) -> CycleOutcome: return CycleOutcome(summary=summary, progressed=False) if control_state is not RunControlState.CANCEL_REQUESTED: self.require_attached_source(run_id) - self.resume_attached_ingestion(run_id) + self.resume_attached_ingestion(run_id, deadline=deadline) + + if _deadline_expired(deadline): + summary = self._state.get_run_summary(run_id=run_id) + return CycleOutcome(summary=summary, progressed=_summary_made_progress(before, summary)) summary = refresh_run( self._polling_deps, run_id=run_id, context=self._context_for_run(run_id), submit_pending_items=self._submit_pending_items, + deadline=deadline, ) progressed = _summary_made_progress(before, summary) return CycleOutcome(summary=summary, progressed=progressed) @@ -103,3 +114,7 @@ def _summary_made_progress(before: RunSummary, after: RunSummary) -> bool: or after.active_batches != before.active_batches or after.status_counts != before.status_counts ) + + +def _deadline_expired(deadline: float | None) -> bool: + return deadline is not None and time.monotonic() >= deadline diff --git a/src/batchor/runtime/ingestion.py b/src/batchor/runtime/ingestion.py index e99d077..46f5c11 100644 --- a/src/batchor/runtime/ingestion.py +++ b/src/batchor/runtime/ingestion.py @@ -4,12 +4,13 @@ import json import time -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Callable, Iterator, cast from pydantic import BaseModel from batchor.core.enums import RunControlState +from batchor.core.exceptions import RunIngestionSourceRequiredError, RunSubmissionIndeterminateError from batchor.core.models import BatchJob, PromptParts from batchor.core.types import JSONObject, JSONValue from batchor.runtime.context import RunContext @@ -48,6 +49,24 @@ class IngestionDeps: configs_match_for_resume: Callable[[PersistedRunConfig, PersistedRunConfig], bool] poll_active_batches: Callable[[str, RunContext], None] = _noop_poll_active_batches monotonic: Callable[[], float] = time.monotonic + work_slice_max_items: int = 1000 + work_slice_max_seconds: float | None = None + work_monotonic: Callable[[], float] = time.monotonic + non_checkpointed_sessions: dict[str, NonCheckpointedIngestionSession] = field(default_factory=dict) + + +@dataclass +class NonCheckpointedIngestionSession: + """In-process cursor for an arbitrary iterable paused at a safe boundary. + + This is deliberately not durable: arbitrary iterables cannot be replayed + in a fresh process. It only lets a manual pause stop before consuming the + next source item and lets the same execution owner continue later. + """ + + iterator: Iterator[Any] + next_item_index: int + seen_ids: set[str] = field(default_factory=set) def finalize_cancelled_ingestion(state: StateStore, *, run_id: str) -> None: @@ -70,6 +89,7 @@ def resume_existing_run( job: BatchJob[Any, BaseModel], config: PersistedRunConfig, context: RunContext, + deadline: float | None = None, ) -> None: """Resume an existing durable run, including checkpointed ingestion. @@ -87,10 +107,14 @@ def resume_existing_run( stored_config = deps.state.get_run_config(run_id=run_id) if not deps.configs_match_for_resume(stored_config, config): raise ValueError(f"existing run config does not match supplied job: {run_id}") + if deps.state.has_indeterminate_submission_intents(run_id=run_id): + raise RunSubmissionIndeterminateError(run_id) deps.state.requeue_local_items(run_id=run_id) control_state = deps.state.get_run_control_state(run_id=run_id) if control_state is RunControlState.CANCEL_REQUESTED: return + if _deadline_reached(deps, deadline): + return if control_state is RunControlState.RUNNING and deps.state.get_active_batches(run_id=run_id): deps.poll_active_batches(run_id, context) summary = deps.state.get_run_summary(run_id=run_id) @@ -103,10 +127,23 @@ def resume_existing_run( return checkpoint = deps.state.get_ingest_checkpoint(run_id=run_id) if checkpoint is not None and not checkpoint.ingestion_complete: - source = require_checkpointed_source(job, run_id=run_id) - validate_checkpoint_source(run_id=run_id, source=source, checkpoint=checkpoint) if control_state is RunControlState.PAUSED: return + source = checkpointed_source(job) + if source is None: + if run_id not in deps.non_checkpointed_sessions and checkpoint.next_item_index != 0: + raise RunIngestionSourceRequiredError(run_id) + ingest_job_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=checkpoint.next_item_index, + checkpoint_payload=None, + deadline=deadline, + ) + return + validate_checkpoint_source(run_id=run_id, source=source, checkpoint=checkpoint) if checkpoint.checkpoint_payload is not None and source.checkpoint_is_complete(checkpoint.checkpoint_payload): deps.state.update_ingest_checkpoint( run_id=run_id, @@ -123,6 +160,7 @@ def resume_existing_run( context=context, start_index=checkpoint.next_item_index, checkpoint_payload=checkpoint.checkpoint_payload, + deadline=deadline, ) return if control_state is RunControlState.PAUSED: @@ -138,6 +176,7 @@ def ingest_job_items( context: RunContext, start_index: int, checkpoint_payload: JSONValue | None, + deadline: float | None = None, ) -> None: """Materialize source items, persist them, and submit as ingestion proceeds. @@ -149,72 +188,320 @@ def ingest_job_items( start_index: Item index to resume from. checkpoint_payload: Source-owned checkpoint payload, if any. """ + if deps.work_slice_max_items <= 0: + raise ValueError("work_slice_max_items must be > 0") + if deps.work_slice_max_seconds is not None and deps.work_slice_max_seconds <= 0: + raise ValueError("work_slice_max_seconds must be > 0 when provided") + + source = checkpointed_source(job) + if source is not None: + _ingest_checkpointed_items( + deps, + run_id=run_id, + job=job, + context=context, + source=source, + start_index=start_index, + checkpoint_payload=checkpoint_payload, + deadline=deadline, + ) + return + _ingest_non_checkpointed_items( + deps, + run_id=run_id, + job=job, + context=context, + start_index=start_index, + deadline=deadline, + ) + + +def _ingest_checkpointed_items( + deps: IngestionDeps, + *, + run_id: str, + job: BatchJob[Any, BaseModel], + context: RunContext, + source: CheckpointedItemSource[Any], + start_index: int, + checkpoint_payload: JSONValue | None, + deadline: float | None, +) -> None: + """Ingest a replayable source in small, atomically checkpointed slices.""" next_item_index = start_index next_checkpoint_payload = checkpoint_payload - checkpointed = checkpointed_source(job) is not None - ingestion_complete = True + source_checkpoint = checkpoint_payload if checkpoint_payload is not None else source.initial_checkpoint() + item_chunk: list[MaterializedItem] = [] + seen_ids: set[str] = set() last_poll_at = deps.monotonic() - for item_chunk, next_item_index, next_checkpoint_payload in materialize_item_chunks( - job, - start_index=start_index, - checkpoint_payload=checkpoint_payload, - ): - if checkpointed: - deps.state.append_items_with_ingest_checkpoint( - run_id=run_id, - items=item_chunk, - next_item_index=next_item_index, - checkpoint_payload=next_checkpoint_payload, - ingestion_complete=False, - ) - else: - deps.state.append_items(run_id=run_id, items=item_chunk) - deps.emit_event( - "items_ingested", - run_id=run_id, - provider_kind=context.config.provider_config.provider_kind, - data={ - "chunk_item_count": len(item_chunk), - "last_item_index": item_chunk[-1].item_index, - }, - ) - control_state = deps.state.get_run_control_state(run_id=run_id) - if control_state is RunControlState.CANCEL_REQUESTED: - ingestion_complete = False + slice_started_at = _next_slice_started_at(deps) if deps.work_slice_max_seconds is not None else last_poll_at + complete = True + for source_item in source.iter_from_checkpoint(source_checkpoint): + if _deadline_reached(deps, deadline): + 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) + item_chunk.append(_materialize_item(job, source_item.item, next_item_index, seen_ids)) + next_item_index += 1 + next_checkpoint_payload = source_item.next_checkpoint + if not _slice_due( + deps, item_count=len(item_chunk), slice_started_at=slice_started_at + ) and not _cooperative_boundary_due( + deps, + run_id=run_id, + item_count=len(item_chunk), ): - 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: - ingestion_complete = False - break - if checkpointed and control_state is not RunControlState.RUNNING: - ingestion_complete = False + continue + should_continue, last_poll_at = _persist_and_advance_boundary( + deps, + run_id=run_id, + context=context, + item_chunk=item_chunk, + next_item_index=next_item_index, + checkpoint_payload=next_checkpoint_payload, + checkpointed=True, + last_poll_at=last_poll_at, + deadline=deadline, + ) + item_chunk = [] + slice_started_at = _next_slice_started_at(deps) + if not should_continue: + complete = False break - if checkpointed: - deps.state.update_ingest_checkpoint( + if complete and item_chunk: + should_continue, _ = _persist_and_advance_boundary( + deps, run_id=run_id, + context=context, + item_chunk=item_chunk, next_item_index=next_item_index, checkpoint_payload=next_checkpoint_payload, - ingestion_complete=ingestion_complete, + checkpointed=True, + last_poll_at=last_poll_at, + deadline=deadline, + ) + if not should_continue: + complete = False + deps.state.update_ingest_checkpoint( + run_id=run_id, + next_item_index=next_item_index, + checkpoint_payload=next_checkpoint_payload, + ingestion_complete=complete, + ) + + +def _ingest_non_checkpointed_items( + deps: IngestionDeps, + *, + run_id: str, + job: BatchJob[Any, BaseModel], + context: RunContext, + start_index: int, + deadline: float | None, +) -> None: + """Ingest an arbitrary iterable, retaining its cursor across manual pause.""" + session = deps.non_checkpointed_sessions.get(run_id) + if session is None: + if start_index != 0: + raise ValueError("non-resumable item sources cannot start from a checkpoint") + session = NonCheckpointedIngestionSession(iterator=iter(job.items), next_item_index=0) + deps.non_checkpointed_sessions[run_id] = session + item_chunk: list[MaterializedItem] = [] + last_poll_at = deps.monotonic() + slice_started_at = _next_slice_started_at(deps) if deps.work_slice_max_seconds is not None else last_poll_at + exhausted = False + while True: + if _deadline_reached(deps, deadline): + break + try: + item = next(session.iterator) + except StopIteration: + exhausted = True + break + item_chunk.append(_materialize_item(job, item, session.next_item_index, session.seen_ids)) + session.next_item_index += 1 + if not _slice_due( + deps, item_count=len(item_chunk), slice_started_at=slice_started_at + ) and not _cooperative_boundary_due( + deps, + run_id=run_id, + item_count=len(item_chunk), + ): + continue + should_continue, last_poll_at = _persist_and_advance_boundary( + deps, + run_id=run_id, + context=context, + item_chunk=item_chunk, + next_item_index=session.next_item_index, + checkpoint_payload=None, + checkpointed=False, + last_poll_at=last_poll_at, + deadline=deadline, ) + item_chunk = [] + slice_started_at = _next_slice_started_at(deps) + if not should_continue: + break + if item_chunk: + should_continue, _ = _persist_and_advance_boundary( + deps, + run_id=run_id, + context=context, + item_chunk=item_chunk, + next_item_index=session.next_item_index, + checkpoint_payload=None, + checkpointed=False, + last_poll_at=last_poll_at, + deadline=deadline, + ) + if not should_continue: + exhausted = False + deps.state.update_ingest_checkpoint( + run_id=run_id, + next_item_index=session.next_item_index, + checkpoint_payload=None, + ingestion_complete=exhausted, + ) + if exhausted or deps.state.get_run_control_state(run_id=run_id) is RunControlState.CANCEL_REQUESTED: + deps.non_checkpointed_sessions.pop(run_id, None) + + +def _materialize_item( + job: BatchJob[Any, BaseModel], item: Any, item_index: int, seen_ids: set[str] +) -> MaterializedItem: + if item.item_id in seen_ids: + raise ValueError(f"duplicate item_id: {item.item_id}") + seen_ids.add(item.item_id) + prompt_parts = normalize_prompt_parts(job.build_prompt(item)) + return MaterializedItem( + item_id=item.item_id, + item_index=item_index, + payload=json_value(item.payload, label=f"payload for {item.item_id}"), + metadata=json_object(item.metadata, label=f"metadata for {item.item_id}"), + prompt=prompt_parts.prompt, + system_prompt=prompt_parts.system_prompt, + ) + + +def _slice_due(deps: IngestionDeps, *, item_count: int, slice_started_at: float) -> bool: + if item_count >= deps.work_slice_max_items: + return True + return ( + deps.work_slice_max_seconds is not None + and deps.work_monotonic() - slice_started_at >= deps.work_slice_max_seconds + ) + + +def _next_slice_started_at(deps: IngestionDeps) -> float: + """Avoid unnecessary clock reads when only an item budget is active.""" + return deps.work_monotonic() if deps.work_slice_max_seconds is not None else 0.0 + + +def _cooperative_boundary_due( + deps: IngestionDeps, + *, + run_id: str, + item_count: int, +) -> bool: + """Observe control/poll cadence inside a large fast-item slice. + + The normal path stays at the configured item chunk size. This only + creates a partial durable append when a manual control change or + cancellation needs it. Slow source work is independently bounded by the time slice, + which then runs the normal active-batch poll cadence. + """ + if deps.work_slice_max_seconds is None or item_count % 100 != 0: + return False + control_state = deps.state.get_run_control_state(run_id=run_id) + return control_state is RunControlState.CANCEL_REQUESTED or _is_manual_pause( + deps, run_id=run_id, control_state=control_state + ) + + +def _deadline_reached(deps: IngestionDeps, deadline: float | None) -> bool: + return deadline is not None and deps.monotonic() >= deadline + + +def _is_manual_pause( + deps: IngestionDeps, + *, + run_id: str, + control_state: RunControlState | None = None, +) -> bool: + if control_state is None: + control_state = deps.state.get_run_control_state(run_id=run_id) + if control_state is not RunControlState.PAUSED: + return False + summary = deps.state.get_run_summary(run_id=run_id) + return summary.control_reason == "manual" + + +def _persist_and_advance_boundary( + deps: IngestionDeps, + *, + run_id: str, + context: RunContext, + item_chunk: list[MaterializedItem], + next_item_index: int, + checkpoint_payload: JSONValue | None, + checkpointed: bool, + last_poll_at: float, + deadline: float | None, +) -> tuple[bool, float]: + """Durably append one slice, reconcile active work, then submit if allowed.""" + deps.state.append_items_with_ingest_checkpoint( + run_id=run_id, + items=item_chunk, + next_item_index=next_item_index, + checkpoint_payload=checkpoint_payload, + ingestion_complete=False, + ) + deps.emit_event( + "items_ingested", + run_id=run_id, + provider_kind=context.config.provider_config.provider_kind, + data={"chunk_item_count": len(item_chunk), "last_item_index": item_chunk[-1].item_index}, + ) + if _deadline_reached(deps, deadline): + return False, last_poll_at + control_state = deps.state.get_run_control_state(run_id=run_id) + if control_state is RunControlState.CANCEL_REQUESTED or _is_manual_pause( + deps, run_id=run_id, control_state=control_state + ): + return False, last_poll_at + # Submission retry backoff never suppresses reconciliation of active work. + active_batches = deps.state.get_active_batches(run_id=run_id) + now = deps.monotonic() if active_batches else last_poll_at + if ( + control_state is RunControlState.RUNNING + and active_batches + and now - last_poll_at >= context.config.provider_config.poll_interval_sec + ): + if _deadline_reached(deps, deadline): + return False, last_poll_at + deps.poll_active_batches(run_id, context) + last_poll_at = now + if _deadline_reached(deps, deadline): + return False, last_poll_at + control_state = deps.state.get_run_control_state(run_id=run_id) + if control_state is RunControlState.CANCEL_REQUESTED or _is_manual_pause( + deps, run_id=run_id, control_state=control_state + ): + return False, last_poll_at + if deps.state.get_batch_retry_backoff_remaining_sec(run_id=run_id) > 0: + return False, last_poll_at + if control_state is RunControlState.RUNNING: + if _deadline_reached(deps, deadline): + return False, last_poll_at + 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 or _is_manual_pause( + deps, run_id=run_id, control_state=control_state + ): + return False, last_poll_at + if checkpointed and control_state is not RunControlState.RUNNING: + return False, last_poll_at + return True, last_poll_at def materialize_item_chunks( diff --git a/src/batchor/runtime/polling.py b/src/batchor/runtime/polling.py index 1a10ea1..8f68ed4 100644 --- a/src/batchor/runtime/polling.py +++ b/src/batchor/runtime/polling.py @@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from time import monotonic from typing import Callable, cast from batchor.artifacts import ArtifactStore @@ -50,12 +51,21 @@ class CompletedBatchConsumeResult: quota_error_count: int = 0 +class _PollingDeadlineExpired(Exception): + """Internal sentinel used to stop before beginning another provider call.""" + + +def _deadline_expired(deadline: float | None) -> bool: + return deadline is not None and monotonic() >= deadline + + def refresh_run( deps: PollingDeps, *, run_id: str, context: RunContext, submit_pending_items: Callable[[str, RunContext], int], + deadline: float | None = None, ) -> RunSummary: """Perform one poll-and-submit cycle and return the updated summary. @@ -72,7 +82,7 @@ def refresh_run( control_state = deps.state.get_run_control_state(run_id=run_id) if control_state is RunControlState.PAUSED: return deps.state.get_run_summary(run_id=run_id) - poll_once(deps, run_id=run_id, context=context) + poll_once(deps, run_id=run_id, context=context, deadline=deadline) if control_state is RunControlState.CANCEL_REQUESTED: if not deps.state.get_active_batches(run_id=run_id): deps.state.mark_nonterminal_items_cancelled( @@ -85,8 +95,13 @@ def refresh_run( ), ) finalize_cancelled_ingestion(deps.state, run_id=run_id) + # Cancellation is irreversible once no remote work remains. A + # retry delay left by the final terminal batch must not keep an + # otherwise terminal cancelled run artificially running. + deps.state.clear_batch_retry_backoff(run_id=run_id) return deps.state.get_run_summary(run_id=run_id) - submit_pending_items(run_id, context) + if not _deadline_expired(deadline): + submit_pending_items(run_id, context) return deps.state.get_run_summary(run_id=run_id) @@ -95,6 +110,7 @@ def poll_once( *, run_id: str, context: RunContext, + deadline: float | None = None, ) -> None: """Poll all active provider batches and consume any terminal ones. @@ -103,13 +119,26 @@ def poll_once( run_id: Durable run identifier. context: Runtime context for the run. """ + if _deadline_expired(deadline): + return batches = deps.state.get_active_batches(run_id=run_id) if not batches: return remote_by_batch_id: dict[str, BatchRemoteRecord] = {} poll_errors: dict[str, Exception] = {} - if len(batches) == 1: + if deadline is not None: + # Deadline-aware polling is deliberately sequential so the check is + # made immediately before every provider call without spawning work + # that may outlive ``Run.wait(timeout=...)``. + for batch in batches: + if _deadline_expired(deadline): + break + try: + remote_by_batch_id[batch.provider_batch_id] = context.provider.get_batch(batch.provider_batch_id) + except Exception as exc: # noqa: BLE001 + poll_errors[batch.provider_batch_id] = exc + elif len(batches) == 1: batch = batches[0] try: remote_by_batch_id[batch.provider_batch_id] = context.provider.get_batch(batch.provider_batch_id) @@ -130,8 +159,18 @@ def poll_once( except Exception as exc: # noqa: BLE001 poll_errors[provider_batch_id] = exc + # A polling pass must never clear a retry delay merely because a different + # batch happened to succeed. Retry state is run-wide today, so a success + # has no unambiguous failure domain to acknowledge. Delays expire on + # their own; ``record_batch_retry_failure`` starts a fresh streak after an + # expired delay. row_quota_backoff_recorded = False for batch in batches: + if _deadline_expired(deadline): + return + if batch.provider_batch_id not in remote_by_batch_id and batch.provider_batch_id not in poll_errors: + # Deadline-aware fetch stopped before this batch. + return if batch.provider_batch_id in poll_errors: exc = poll_errors[batch.provider_batch_id] if is_insufficient_quota_error(exc): @@ -177,13 +216,25 @@ def poll_once( error_file_id=remote.get("error_file_id"), ) if status == "completed": - consume_result = consume_completed_batch( - deps, + try: + consume_result = consume_completed_batch( + deps, + run_id=run_id, + context=context, + provider_batch_id=batch.provider_batch_id, + output_file_id=remote.get("output_file_id"), + error_file_id=remote.get("error_file_id"), + deadline=deadline, + ) + except _PollingDeadlineExpired: + return + # The shared enqueue reservation is held until the terminal + # result has been durably consumed. Releasing it earlier would + # let a failed download/parse strand submitted work while another + # run spends the same capacity. + deps.state.release_submission_capacity_for_batch( run_id=run_id, - context=context, provider_batch_id=batch.provider_batch_id, - output_file_id=remote.get("output_file_id"), - error_file_id=remote.get("error_file_id"), ) if consume_result.quota_error_count > 0 and not row_quota_backoff_recorded: row_quota_backoff_recorded = True @@ -193,8 +244,6 @@ def poll_once( base_delay_sec=context.config.retry_policy.base_backoff_sec, max_delay_sec=context.config.retry_policy.max_backoff_sec, ) - if not row_quota_backoff_recorded: - deps.state.clear_batch_retry_backoff(run_id=run_id) deps.emit_event( "batch_completed", run_id=run_id, @@ -202,11 +251,15 @@ def poll_once( data={"provider_batch_id": batch.provider_batch_id}, ) elif status in {"failed", "cancelled", "expired"}: - output_content, error_content = download_batch_file_contents( - context=context, - output_file_id=remote.get("output_file_id"), - error_file_id=remote.get("error_file_id"), - ) + try: + output_content, error_content = download_batch_file_contents( + context=context, + output_file_id=remote.get("output_file_id"), + error_file_id=remote.get("error_file_id"), + deadline=deadline, + ) + except _PollingDeadlineExpired: + return output_artifact_path, error_artifact_path = write_batch_result_artifacts( artifact_store=deps.artifact_store, run_id=run_id, @@ -232,6 +285,10 @@ def poll_once( provider_batch_id=batch.provider_batch_id, error=error, ) + deps.state.release_submission_capacity_for_batch( + run_id=run_id, + provider_batch_id=batch.provider_batch_id, + ) deps.emit_event( "batch_terminal_failure", run_id=run_id, @@ -265,6 +322,7 @@ def download_batch_file_contents( context: RunContext, output_file_id: object, error_file_id: object, + deadline: float | None = None, ) -> tuple[str | None, str | None]: """Download the output and error files for a provider batch. @@ -279,15 +337,20 @@ def download_batch_file_contents( """ output_id = output_file_id if isinstance(output_file_id, str) else None error_id = error_file_id if isinstance(error_file_id, str) else None + if _deadline_expired(deadline): + raise _PollingDeadlineExpired if output_id is None and error_id is None: return None, None - if output_id is not None and error_id is not None: + if output_id is not None and error_id is not None and deadline is None: with ThreadPoolExecutor(max_workers=2) as executor: output_future = executor.submit(context.provider.download_file_content, output_id) error_future = executor.submit(context.provider.download_file_content, error_id) return output_future.result(), error_future.result() if output_id is not None: - return context.provider.download_file_content(output_id), None + output_content = context.provider.download_file_content(output_id) + if _deadline_expired(deadline): + raise _PollingDeadlineExpired + return output_content, (context.provider.download_file_content(error_id) if error_id is not None else None) if error_id is None: return None, None return None, context.provider.download_file_content(error_id) @@ -301,6 +364,7 @@ def consume_completed_batch( provider_batch_id: str, output_file_id: object, error_file_id: object, + deadline: float | None = None, ) -> CompletedBatchConsumeResult: """Download, persist, and parse a completed provider batch. @@ -320,6 +384,7 @@ def consume_completed_batch( context=context, output_file_id=output_file_id, error_file_id=error_file_id, + deadline=deadline, ) output_artifact_path, error_artifact_path = write_batch_result_artifacts( artifact_store=deps.artifact_store, @@ -356,7 +421,13 @@ def consume_completed_batch( processed_custom_ids: set[str] = set() quota_error_count = 0 + # A terminal batch may be replayed after any failure between recording the + # remote terminal state and mutating local items. Only consume records + # that still belong to submitted items: already terminalized rows must not + # be parsed again or have their attempt count incremented a second time. for custom_id, record in successes.items(): + if custom_id not in submitted_custom_ids: + continue processed_custom_ids.add(custom_id) if context.output_model is None: completions.append( @@ -401,6 +472,8 @@ def consume_completed_batch( ) for custom_id, error_record in errors.items(): + if custom_id not in submitted_custom_ids: + continue processed_custom_ids.add(custom_id) quota_error = is_insufficient_quota_error(error_record) if quota_error: diff --git a/src/batchor/runtime/run_handle.py b/src/batchor/runtime/run_handle.py index a268e56..6b74856 100644 --- a/src/batchor/runtime/run_handle.py +++ b/src/batchor/runtime/run_handle.py @@ -97,13 +97,13 @@ def is_finished(self) -> bool: """ return self.status in (RunLifecycleStatus.COMPLETED, RunLifecycleStatus.COMPLETED_WITH_FAILURES) - def refresh(self) -> RunSummary: + def refresh(self, *, deadline: float | None = None) -> RunSummary: """Perform one poll-and-submit cycle and return the updated summary. Returns: Updated run summary. """ - outcome = self._runner._advance_run(self.run_id) + outcome = self._runner._advance_run(self.run_id, deadline=deadline) self._summary = outcome.summary self._last_cycle_progressed = outcome.progressed return self._summary @@ -129,7 +129,9 @@ def wait( """ deadline = None if timeout is None else time.monotonic() + timeout while True: - self.refresh() + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError(f"timed out waiting for run {self.run_id}") + self.refresh(deadline=deadline) if self.is_finished: return self if self.control_state is RunControlState.PAUSED: @@ -146,7 +148,10 @@ def wait( else: sleep_for = self._summary.backoff_remaining_sec if sleep_for > 0: - self._runner.sleep(sleep_for) + if deadline is not None: + sleep_for = min(sleep_for, max(0.0, deadline - time.monotonic())) + if sleep_for > 0: + self._runner.sleep(sleep_for) def summary(self) -> RunSummary: """Read the latest persisted summary for the run from storage. @@ -234,6 +239,27 @@ def pause(self) -> RunSummary: self._summary = self._runner.pause_run(self.run_id).summary() return self._summary + def resolve_indeterminate_submission_as_not_created(self) -> Run: + """Acknowledge an operator-verified absent remote batch and allow retry.""" + resolved = self._runner.resolve_indeterminate_submission_as_not_created(self.run_id) + self._summary = resolved.summary() + return self + + def resolve_indeterminate_submission_as_created( + self, + *, + provider_batch_id: str, + status: str = "submitted", + ) -> Run: + """Link an operator-confirmed remote batch and resume normal polling.""" + resolved = self._runner.resolve_indeterminate_submission_as_created( + self.run_id, + provider_batch_id=provider_batch_id, + status=status, + ) + self._summary = resolved.summary() + return self + def resume(self) -> RunSummary: """Resume this run after it has been paused. diff --git a/src/batchor/runtime/runner.py b/src/batchor/runtime/runner.py index 2ad46d8..2a7a55e 100644 --- a/src/batchor/runtime/runner.py +++ b/src/batchor/runtime/runner.py @@ -111,6 +111,8 @@ def __init__( submit_pending_items=self._submit_pending_items, configs_match_for_resume=self._configs_match_for_resume, poll_active_batches=self._poll_active_batches, + work_slice_max_items=1000, + work_slice_max_seconds=0.25, ) self._executor = RunExecutor( state=self.state, @@ -272,6 +274,26 @@ def cancel_run(self, run_id: str) -> Run: ) return self.get_run(run_id) + def resolve_indeterminate_submission_as_not_created(self, run_id: str) -> Run: + """Release an intent only after the operator verified no remote batch exists.""" + self.state.abandon_indeterminate_submission_intents(run_id=run_id) + return self.get_run(run_id) + + def resolve_indeterminate_submission_as_created( + self, + run_id: str, + *, + provider_batch_id: str, + status: str = "submitted", + ) -> Run: + """Link an operator-confirmed remote batch to its persisted intent.""" + self.state.finalize_indeterminate_submission_as_created( + run_id=run_id, + provider_batch_id=provider_batch_id, + status=status, + ) + return self.get_run(run_id) + def read_terminal_results( self, run_id: str, @@ -558,8 +580,8 @@ def _submit_pending_items( context=context or self._context_for_run(run_id), ) - def _advance_run(self, run_id: str) -> CycleOutcome: - return self._executor.advance(run_id) + def _advance_run(self, run_id: str, *, deadline: float | None = None) -> CycleOutcome: + return self._executor.advance(run_id, deadline=deadline) def _poll_active_batches( self, diff --git a/src/batchor/runtime/submission.py b/src/batchor/runtime/submission.py index 84e1729..4030e33 100644 --- a/src/batchor/runtime/submission.py +++ b/src/batchor/runtime/submission.py @@ -10,7 +10,8 @@ from batchor.artifacts import ArtifactStore from batchor.core.enums import RunControlState -from batchor.core.models import ItemFailure, OpenAIProviderConfig, PromptParts +from batchor.core.exceptions import RunSubmissionIndeterminateError +from batchor.core.models import ItemFailure, OpenAIProviderConfig, PromptParts, openai_enqueue_quota_scope from batchor.core.types import BatchRequestLine, JSONObject from batchor.runtime.artifacts import ( RequestArtifactCache, @@ -22,8 +23,8 @@ from batchor.runtime.context import RunContext from batchor.runtime.retry import ( classify_batch_error, + is_enqueue_token_limit_error, is_insufficient_quota_error, - is_retryable_batch_control_plane_error, ) from batchor.runtime.tokens import ( chunk_request_rows, @@ -194,14 +195,10 @@ def submit_pending_items( ) submitted_item_ids: set[str] = set() submitted_count = 0 - active_tokens = deps.state.get_active_submitted_token_estimate(run_id=run_id) if inflight_budget is not None else 0 - for chunk_index, chunk in enumerate(chunks, start=1): if deps.state.get_run_control_state(run_id=run_id) is not RunControlState.RUNNING: break chunk_tokens = sum(int(row["submission_tokens"]) for row in chunk) - if inflight_budget is not None and chunk_tokens > max(inflight_budget - active_tokens, 0): - break request_lines = [cast(JSONObject, row["request_line"]) for row in chunk] request_relative_path = request_artifact_relative_path(run_id) deps.artifact_store.write_text( @@ -223,11 +220,29 @@ def submit_pending_items( ) if deps.state.get_run_control_state(run_id=run_id) is not RunControlState.RUNNING: break + intent_id = f"submission-{uuid4().hex}" + if not deps.state.begin_submission_intent( + run_id=run_id, + intent_id=intent_id, + submissions=[ + PreparedSubmission( + item_id=str(row["item_id"]), + custom_id=str(row["custom_id"]), + submission_tokens=int(row["submission_tokens"]), + ) + for row in chunk + ], + quota_scope=submission_quota_scope(config.provider_config) if inflight_budget is not None else None, + submission_tokens=chunk_tokens, + capacity_limit=inflight_budget, + ): + break with ExitStack() as stack: request_file = stack.enter_context(deps.artifact_store.stage_local_copy(request_relative_path.as_posix())) try: remote_input_file_id = context.provider.upload_input_file(request_file) except Exception as exc: # noqa: BLE001 + deps.state.abandon_submission_intent(intent_id=intent_id) if is_insufficient_quota_error(exc): _release_unsubmitted_claimed_items( deps, @@ -251,8 +266,9 @@ def submit_pending_items( metadata={"run_id": run_id, **context.config.batch_metadata}, ) except Exception as exc: # noqa: BLE001 - cleanup_uploaded_input_file(context, remote_input_file_id) if is_insufficient_quota_error(exc): + deps.state.abandon_submission_intent(intent_id=intent_id) + cleanup_uploaded_input_file(context, remote_input_file_id) _release_unsubmitted_claimed_items( deps, run_id=run_id, @@ -268,35 +284,49 @@ def submit_pending_items( phase="batch_create", ) return submitted_count - if not is_retryable_batch_control_plane_error(exc): - raise - deps.state.record_batch_retry_failure( - run_id=run_id, - error_class=classify_batch_error(exc), - base_delay_sec=config.retry_policy.base_backoff_sec, - max_delay_sec=config.retry_policy.max_backoff_sec, - ) - deps.emit_event( - "batch_submit_retry", - run_id=run_id, - provider_kind=context.config.provider_config.provider_kind, - data={"error_class": classify_batch_error(exc)}, - ) - break + if is_enqueue_token_limit_error(exc): + # The provider explicitly rejected this create for an + # enqueue-budget violation; no remote batch exists. + deps.state.abandon_submission_intent(intent_id=intent_id) + cleanup_uploaded_input_file(context, remote_input_file_id) + deps.state.record_batch_retry_failure( + run_id=run_id, + error_class=classify_batch_error(exc), + base_delay_sec=config.retry_policy.base_backoff_sec, + max_delay_sec=config.retry_policy.max_backoff_sec, + ) + deps.emit_event( + "batch_submit_retry", + run_id=run_id, + provider_kind=context.config.provider_config.provider_kind, + data={"error_class": classify_batch_error(exc)}, + ) + break + # Generic providers cannot prove a failed create call did not + # create remote work (timeouts are especially ambiguous). + # Leave the durable intent and input artifact for operator + # inspection rather than silently issuing a duplicate batch. + intent_item_ids = {str(row["item_id"]) for row in chunk} + safe_unsubmitted = [ + item.item_id + for item in claimed + if item.item_id not in submitted_item_ids + and item.item_id not in failed_item_ids + and item.item_id not in intent_item_ids + ] + if safe_unsubmitted: + deps.state.release_items_to_pending(run_id=run_id, item_ids=safe_unsubmitted) + raise RunSubmissionIndeterminateError(run_id) from exc - deps.state.clear_batch_retry_backoff(run_id=run_id) provider_batch_id = str(batch["id"]) local_batch_id = f"batch-{chunk_index:04d}-{uuid4().hex[:8]}" - deps.state.register_batch( + deps.state.finalize_submission_intent( run_id=run_id, + intent_id=intent_id, local_batch_id=local_batch_id, provider_batch_id=provider_batch_id, status=str(batch.get("status", "submitted")), custom_ids=[str(row["custom_id"]) for row in chunk], - ) - deps.state.mark_items_submitted( - run_id=run_id, - provider_batch_id=provider_batch_id, submissions=[ PreparedSubmission( item_id=str(row["item_id"]), @@ -318,7 +348,6 @@ def submit_pending_items( ) submitted_count += len(chunk) submitted_item_ids.update(str(row["item_id"]) for row in chunk) - active_tokens += chunk_tokens unsent = [ item.item_id @@ -330,6 +359,13 @@ def submit_pending_items( return submitted_count +def submission_quota_scope(provider_config: object) -> str: + """Return the durable capacity scope for an OpenAI submission.""" + if not isinstance(provider_config, OpenAIProviderConfig): + raise TypeError("OpenAI enqueue capacity requires OpenAIProviderConfig") + return openai_enqueue_quota_scope(provider_config) + + def _release_unsubmitted_claimed_items( deps: SubmissionDeps, *, diff --git a/src/batchor/storage/memory.py b/src/batchor/storage/memory.py index 2189c6b..be8bac2 100644 --- a/src/batchor/storage/memory.py +++ b/src/batchor/storage/memory.py @@ -13,10 +13,11 @@ from dataclasses import asdict, dataclass, field from datetime import UTC, datetime, timedelta +from threading import RLock from typing import Callable from batchor.core.enums import ItemStatus, RunControlState, RunLifecycleStatus -from batchor.core.models import ItemFailure, RunSummary +from batchor.core.models import ItemFailure, OpenAIProviderConfig, RunSummary, openai_enqueue_quota_scope from batchor.core.types import JSONObject, JSONValue from batchor.runtime.retry import compute_backoff_delay from batchor.storage.state_models import ( @@ -74,6 +75,20 @@ class _StoredBatch: error_artifact_path: str | None = None +@dataclass +class _StoredSubmissionIntent: + intent_id: str + run_id: str + custom_ids: list[str] + item_ids: list[str] + submissions: list[PreparedSubmission] + quota_scope: str | None + submission_tokens: int + status: str = "creating" + provider_batch_id: str | None = None + capacity_released: bool = False + + @dataclass class _StoredRun: run_id: str @@ -109,6 +124,8 @@ def __init__( now: Callable[[], datetime] | None = None, ) -> None: self._runs: dict[str, _StoredRun] = {} + self._submission_intents: dict[str, _StoredSubmissionIntent] = {} + self._submission_intent_lock = RLock() self._now = now or (lambda: datetime.now(UTC)) def has_run(self, *, run_id: str) -> bool: @@ -123,7 +140,16 @@ def create_run( ) -> None: if run_id in self._runs: raise ValueError(f"run already exists: {run_id}") - run = _StoredRun(run_id=run_id, config=config) + run = _StoredRun( + run_id=run_id, + config=config, + ingest_checkpoint=IngestCheckpoint( + source_kind="unattached", + source_ref="", + source_fingerprint="", + ingestion_complete=bool(items), + ), + ) self._runs[run_id] = run self.append_items(run_id=run_id, items=items) self._refresh_run_status(run) @@ -442,6 +468,151 @@ def mark_items_submitted( item.error = None self._refresh_run_status(run) + def begin_submission_intent( + self, + *, + run_id: str, + intent_id: str, + submissions: list[PreparedSubmission], + quota_scope: str | None, + submission_tokens: int, + capacity_limit: int | None, + ) -> bool: + with self._submission_intent_lock: + if quota_scope is not None and capacity_limit is not None: + reserved = self._legacy_active_submitted_tokens(quota_scope=quota_scope) + sum( + intent.submission_tokens + for intent in self._submission_intents.values() + if intent.quota_scope == quota_scope and not intent.capacity_released + ) + if reserved + submission_tokens > capacity_limit: + return False + self._submission_intents[intent_id] = _StoredSubmissionIntent( + intent_id=intent_id, + run_id=run_id, + custom_ids=[submission.custom_id for submission in submissions], + item_ids=[submission.item_id for submission in submissions], + submissions=list(submissions), + quota_scope=quota_scope, + submission_tokens=submission_tokens, + ) + return True + + def finalize_submission_intent( + self, + *, + run_id: str, + intent_id: str, + local_batch_id: str, + provider_batch_id: str, + status: str, + custom_ids: list[str], + submissions: list[PreparedSubmission], + ) -> None: + with self._submission_intent_lock: + intent = self._submission_intents.get(intent_id) + if intent is None or intent.run_id != run_id or intent.status != "creating": + raise ValueError(f"submission intent is not creating: {intent_id}") + run = self._get_run(run_id) + run.batches[provider_batch_id] = _StoredBatch( + local_batch_id=local_batch_id, + provider_batch_id=provider_batch_id, + status=status, + custom_ids=list(custom_ids), + ) + for submission in submissions: + item = run.items[submission.item_id] + item.status = ItemStatus.SUBMITTED + item.active_batch_id = provider_batch_id + item.active_custom_id = submission.custom_id + item.active_submission_tokens = submission.submission_tokens + item.error = None + intent.status = "active" + intent.provider_batch_id = provider_batch_id + self._refresh_run_status(run) + + def abandon_submission_intent(self, *, intent_id: str) -> None: + with self._submission_intent_lock: + intent = self._submission_intents.get(intent_id) + if intent is not None: + intent.status = "abandoned" + intent.capacity_released = True + + def has_indeterminate_submission_intents(self, *, run_id: str) -> bool: + with self._submission_intent_lock: + return any( + intent.run_id == run_id and intent.status == "creating" for intent in self._submission_intents.values() + ) + + def release_submission_capacity_for_batch(self, *, run_id: str, provider_batch_id: str) -> None: + with self._submission_intent_lock: + for intent in self._submission_intents.values(): + if ( + intent.run_id == run_id + and intent.provider_batch_id == provider_batch_id + and intent.status == "active" + ): + intent.capacity_released = True + + def abandon_indeterminate_submission_intents(self, *, run_id: str) -> None: + with self._submission_intent_lock: + abandoned_item_ids: set[str] = set() + for intent in self._submission_intents.values(): + if intent.run_id == run_id and intent.status == "creating": + abandoned_item_ids.update(intent.item_ids) + intent.status = "operator_abandoned" + intent.capacity_released = True + run = self._get_run(run_id) + for item_id in abandoned_item_ids: + item = run.items[item_id] + if item.status is ItemStatus.QUEUED_LOCAL: + item.status = ItemStatus.PENDING + item.active_batch_id = None + item.active_custom_id = None + item.active_submission_tokens = 0 + self._refresh_run_status(run) + + def _legacy_active_submitted_tokens(self, *, quota_scope: str) -> int: + """Count submitted rows that predate durable submission intents.""" + reserved = 0 + intent_batch_ids = { + (intent.run_id, intent.provider_batch_id) + for intent in self._submission_intents.values() + if intent.provider_batch_id is not None + } + for run in self._runs.values(): + provider_config = run.config.provider_config + if not isinstance(provider_config, OpenAIProviderConfig): + continue + if openai_enqueue_quota_scope(provider_config) != quota_scope: + continue + reserved += sum( + item.active_submission_tokens + for item in run.items.values() + if item.status is ItemStatus.SUBMITTED and (run.run_id, item.active_batch_id) not in intent_batch_ids + ) + return reserved + + def finalize_indeterminate_submission_as_created(self, *, run_id: str, provider_batch_id: str, status: str) -> None: + with self._submission_intent_lock: + intents = [ + intent + for intent in self._submission_intents.values() + if intent.run_id == run_id and intent.status == "creating" + ] + if len(intents) != 1: + raise ValueError("expected exactly one indeterminate submission intent") + intent = intents[0] + self.finalize_submission_intent( + run_id=run_id, + intent_id=intent.intent_id, + local_batch_id=f"recovered-{intent.intent_id[-8:]}", + provider_batch_id=provider_batch_id, + status=status, + custom_ids=intent.custom_ids, + submissions=intent.submissions, + ) + def update_batch_status( self, *, @@ -468,7 +639,20 @@ def get_active_batches(self, *, run_id: str) -> list[ActiveBatchRecord]: error_file_id=batch.error_file_id, ) for batch in run.batches.values() + # A provider-terminal batch remains active locally until every + # submitted item has been consumed. This makes result download + # and item-state mutations safely replayable after a crash. if batch.status not in self.TERMINAL_BATCH_STATUSES + or any( + item.status is ItemStatus.SUBMITTED and item.active_batch_id == batch.provider_batch_id + for item in run.items.values() + ) + or any( + intent.run_id == run_id + and intent.provider_batch_id == batch.provider_batch_id + and not intent.capacity_released + for intent in self._submission_intents.values() + ) ] def get_submitted_custom_ids_for_batch( @@ -497,6 +681,8 @@ def mark_items_completed( run = self._get_run(run_id) for completion in completions: item = self._item_for_custom_id(run, completion.custom_id) + if item.status is not ItemStatus.SUBMITTED: + continue item.attempt_count += 1 item.status = ItemStatus.COMPLETED item.terminal_result_sequence = self._next_terminal_sequence(run) @@ -520,6 +706,8 @@ def mark_items_failed( run = self._get_run(run_id) for failure in failures: item = self._item_for_custom_id(run, failure.custom_id) + if item.status is not ItemStatus.SUBMITTED: + continue if failure.count_attempt: item.attempt_count += 1 item.status = self._failed_status( @@ -597,7 +785,11 @@ def record_batch_retry_failure( max_delay_sec: float, ) -> RetryBackoffState: run = self._get_run(run_id) - consecutive = run.backoff.consecutive_failures + 1 + # Successful work in another batch must not clear this delay. A new + # failure after the previous delay has elapsed begins a fresh streak, + # keeping exponential backoff bounded without cross-batch coupling. + previous_is_active = run.backoff.next_retry_at is not None and run.backoff.next_retry_at > self._now() + consecutive = (run.backoff.consecutive_failures if previous_is_active else 0) + 1 total = run.backoff.total_failures + 1 backoff_sec = compute_backoff_delay( consecutive_failures=consecutive, @@ -735,7 +927,7 @@ def _item_for_custom_id(run: _StoredRun, custom_id: str) -> _StoredItem: def _refresh_run_status(self, run: _StoredRun) -> None: all_terminal = all(run.items[item_id].status in self.TERMINAL_ITEM_STATUSES for item_id in run.item_ids) - active_batches = any(batch.status not in self.TERMINAL_BATCH_STATUSES for batch in run.batches.values()) + active_batches = bool(self.get_active_batches(run_id=run.run_id)) backoff_remaining = self.get_batch_retry_backoff_remaining_sec(run_id=run.run_id) ingestion_complete = run.ingest_checkpoint is None or run.ingest_checkpoint.ingestion_complete if all_terminal and not active_batches and backoff_remaining <= 0 and ingestion_complete: diff --git a/src/batchor/storage/postgres_store.py b/src/batchor/storage/postgres_store.py index e392a28..e3b19c3 100644 --- a/src/batchor/storage/postgres_store.py +++ b/src/batchor/storage/postgres_store.py @@ -81,6 +81,7 @@ def _ensure_schema(self) -> None: } if "control_reason" not in run_columns: conn.execute(text(f'ALTER TABLE "{self.schema}"."runs" ADD COLUMN control_reason VARCHAR')) + self._backfill_missing_ingest_markers(conn) existing_schema_row = conn.execute( select(STORAGE_METADATA_TABLE.c.value).where(STORAGE_METADATA_TABLE.c.key == "schema_version") ).first() diff --git a/src/batchor/storage/sqlite_lifecycle.py b/src/batchor/storage/sqlite_lifecycle.py index 72d65bb..a90f36c 100644 --- a/src/batchor/storage/sqlite_lifecycle.py +++ b/src/batchor/storage/sqlite_lifecycle.py @@ -2,10 +2,12 @@ from dataclasses import asdict -from sqlalchemy import and_, bindparam, select, update +from sqlalchemy import and_, bindparam, exists, select, update from sqlalchemy.engine import RowMapping +from sqlalchemy.exc import IntegrityError from batchor.core.enums import ItemStatus, RunControlState, RunLifecycleStatus +from batchor.core.models import OpenAIProviderConfig, openai_enqueue_quota_scope from batchor.core.types import JSONValue from batchor.storage.sqlite_codec import ( _decode_json, @@ -18,10 +20,12 @@ from batchor.storage.sqlite_protocol import SQLiteStorageProtocol from batchor.storage.sqlite_schema import ( BATCHES_TABLE, + CAPACITY_SCOPES_TABLE, ITEMS_TABLE, RUN_INGEST_STATE_TABLE, RUN_RETRY_STATE_TABLE, RUNS_TABLE, + SUBMISSION_INTENTS_TABLE, ) from batchor.storage.state import ( ActiveBatchRecord, @@ -82,6 +86,27 @@ def create_run( } ], ) + # Every run starts life with an incomplete ingest marker. It is + # deliberately generic: arbitrary iterables cannot be resumed in + # a fresh process, but they must never make an empty run appear + # terminal while their first chunk is still being materialized. + conn.execute( + RUN_INGEST_STATE_TABLE.insert(), + [ + { + "run_id": run_id, + "source_kind": "unattached", + "source_ref": "", + "source_fingerprint": "", + "next_item_index": 0, + "checkpoint_payload_json": None, + # Initial rows are already fully materialized; runner + # creation deliberately passes an empty list and remains + # incomplete until its source finishes. + "ingestion_complete": 1 if items else 0, + } + ], + ) if items: conn.execute(ITEMS_TABLE.insert(), self._item_rows(run_id=run_id, items=items)) @@ -169,22 +194,22 @@ def set_ingest_checkpoint( checkpoint: IngestCheckpoint, ) -> None: with self.engine.begin() as conn: - conn.execute( - RUN_INGEST_STATE_TABLE.insert(), - [ - { - "run_id": run_id, - "source_kind": checkpoint.source_kind, - "source_ref": checkpoint.source_ref, - "source_fingerprint": checkpoint.source_fingerprint, - "next_item_index": checkpoint.next_item_index, - "checkpoint_payload_json": _encode_json(checkpoint.checkpoint_payload) - if checkpoint.checkpoint_payload is not None - else None, - "ingestion_complete": 1 if checkpoint.ingestion_complete else 0, - } - ], + values = { + "run_id": run_id, + "source_kind": checkpoint.source_kind, + "source_ref": checkpoint.source_ref, + "source_fingerprint": checkpoint.source_fingerprint, + "next_item_index": checkpoint.next_item_index, + "checkpoint_payload_json": _encode_json(checkpoint.checkpoint_payload) + if checkpoint.checkpoint_payload is not None + else None, + "ingestion_complete": 1 if checkpoint.ingestion_complete else 0, + } + updated = conn.execute( + update(RUN_INGEST_STATE_TABLE).where(RUN_INGEST_STATE_TABLE.c.run_id == run_id).values(**values) ) + if updated.rowcount == 0: + conn.execute(RUN_INGEST_STATE_TABLE.insert(), [values]) def update_ingest_checkpoint( self, @@ -392,6 +417,343 @@ def mark_items_submitted( ], ) + def begin_submission_intent( + self, + *, + run_id: str, + intent_id: str, + submissions: list[PreparedSubmission], + quota_scope: str | None, + submission_tokens: int, + capacity_limit: int | None, + ) -> bool: + with self.engine.begin() as conn: + if quota_scope is not None and capacity_limit is not None: + legacy_reserved = self._legacy_active_submitted_tokens(conn, quota_scope=quota_scope) + try: + # A savepoint keeps a Postgres transaction usable after a + # concurrent unique-key winner creates the scope first. + with conn.begin_nested(): + conn.execute( + CAPACITY_SCOPES_TABLE.insert(), + [{"quota_scope": quota_scope, "reserved_tokens": 0}], + ) + except IntegrityError: + # A competing transaction created the counter first. + pass + reserved = conn.execute( + update(CAPACITY_SCOPES_TABLE) + .where( + (CAPACITY_SCOPES_TABLE.c.quota_scope == quota_scope) + & ( + CAPACITY_SCOPES_TABLE.c.reserved_tokens + <= capacity_limit - submission_tokens - legacy_reserved + ) + ) + .values(reserved_tokens=CAPACITY_SCOPES_TABLE.c.reserved_tokens + submission_tokens) + ) + if reserved.rowcount != 1: + return False + conn.execute( + SUBMISSION_INTENTS_TABLE.insert(), + [ + { + "intent_id": intent_id, + "run_id": run_id, + "status": "creating", + "provider_batch_id": None, + "custom_ids_json": _encode_json([submission.custom_id for submission in submissions]), + "item_ids_json": _encode_json([submission.item_id for submission in submissions]), + "submissions_json": _encode_json( + [ + { + "item_id": submission.item_id, + "custom_id": submission.custom_id, + "submission_tokens": submission.submission_tokens, + } + for submission in submissions + ] + ), + "quota_scope": quota_scope, + "submission_tokens": submission_tokens, + "capacity_released": 0, + "created_at": _encode_datetime(self._now()), + } + ], + ) + return True + + def finalize_submission_intent( + self, + *, + run_id: str, + intent_id: str, + local_batch_id: str, + provider_batch_id: str, + status: str, + custom_ids: list[str], + submissions: list[PreparedSubmission], + ) -> None: + with self.engine.begin() as conn: + intent = conn.execute( + select(SUBMISSION_INTENTS_TABLE.c.status).where( + (SUBMISSION_INTENTS_TABLE.c.intent_id == intent_id) & (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + ) + ).first() + if intent is None or str(intent[0]) != "creating": + raise ValueError(f"submission intent is not creating: {intent_id}") + conn.execute( + BATCHES_TABLE.insert(), + [ + { + "run_id": run_id, + "provider_batch_id": provider_batch_id, + "local_batch_id": local_batch_id, + "status": status, + "custom_ids_json": _encode_json(custom_ids), + "output_file_id": None, + "error_file_id": None, + } + ], + ) + if submissions: + statement = ( + update(ITEMS_TABLE) + .where( + and_( + ITEMS_TABLE.c.run_id == bindparam("b_run_id"), + ITEMS_TABLE.c.item_id == bindparam("b_item_id"), + ) + ) + .values( + status=ItemStatus.SUBMITTED, + active_batch_id=bindparam("b_provider_batch_id"), + active_custom_id=bindparam("b_custom_id"), + active_submission_tokens=bindparam("b_submission_tokens"), + error_json=None, + ) + ) + conn.execute( + statement, + [ + { + "b_run_id": run_id, + "b_item_id": value.item_id, + "b_provider_batch_id": provider_batch_id, + "b_custom_id": value.custom_id, + "b_submission_tokens": value.submission_tokens, + } + for value in submissions + ], + ) + conn.execute( + update(SUBMISSION_INTENTS_TABLE) + .where(SUBMISSION_INTENTS_TABLE.c.intent_id == intent_id) + .values(status="active", provider_batch_id=provider_batch_id) + ) + + def abandon_submission_intent(self, *, intent_id: str) -> None: + with self.engine.begin() as conn: + self._release_submission_intents(conn, intent_ids=[intent_id], status="abandoned") + + def has_indeterminate_submission_intents(self, *, run_id: str) -> bool: + with self.engine.begin() as conn: + return ( + conn.execute( + select(SUBMISSION_INTENTS_TABLE.c.intent_id) + .where( + (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + & (SUBMISSION_INTENTS_TABLE.c.status == "creating") + ) + .limit(1) + ).first() + is not None + ) + + def release_submission_capacity_for_batch(self, *, run_id: str, provider_batch_id: str) -> None: + with self.engine.begin() as conn: + intent_ids = list( + conn.execute( + select(SUBMISSION_INTENTS_TABLE.c.intent_id).where( + (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + & (SUBMISSION_INTENTS_TABLE.c.provider_batch_id == provider_batch_id) + & (SUBMISSION_INTENTS_TABLE.c.status == "active") + ) + ).scalars() + ) + self._release_submission_intents(conn, intent_ids=[str(value) for value in intent_ids], status=None) + + def abandon_indeterminate_submission_intents(self, *, run_id: str) -> None: + with self.engine.begin() as conn: + intent_rows = list( + conn.execute( + select( + SUBMISSION_INTENTS_TABLE.c.intent_id, + SUBMISSION_INTENTS_TABLE.c.item_ids_json, + ).where( + (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + & (SUBMISSION_INTENTS_TABLE.c.status == "creating") + ) + ).mappings() + ) + self._release_submission_intents( + conn, + intent_ids=[str(row["intent_id"]) for row in intent_rows], + status="operator_abandoned", + ) + item_ids = {str(item_id) for row in intent_rows for item_id in _decode_json(row["item_ids_json"])} + if item_ids: + conn.execute( + update(ITEMS_TABLE) + .where( + (ITEMS_TABLE.c.run_id == run_id) + & (ITEMS_TABLE.c.item_id.in_(item_ids)) + & (ITEMS_TABLE.c.status == ItemStatus.QUEUED_LOCAL) + ) + .values( + status=ItemStatus.PENDING, + active_batch_id=None, + active_custom_id=None, + active_submission_tokens=0, + ) + ) + + def _legacy_active_submitted_tokens(self, conn, *, quota_scope: str) -> int: # noqa: ANN001 + """Count active submitted tokens not represented by a v6 intent.""" + rows = conn.execute( + select( + RUNS_TABLE.c.provider_config_json, + ITEMS_TABLE.c.active_submission_tokens, + ) + .select_from(ITEMS_TABLE.join(RUNS_TABLE, ITEMS_TABLE.c.run_id == RUNS_TABLE.c.run_id)) + .where( + (ITEMS_TABLE.c.status == ItemStatus.SUBMITTED) + & ~exists( + select(SUBMISSION_INTENTS_TABLE.c.intent_id).where( + (SUBMISSION_INTENTS_TABLE.c.run_id == ITEMS_TABLE.c.run_id) + & (SUBMISSION_INTENTS_TABLE.c.provider_batch_id == ITEMS_TABLE.c.active_batch_id) + ) + ) + ) + ).mappings() + reserved = 0 + for row in rows: + provider_config = self.provider_registry.load_config(_decode_object(row["provider_config_json"])) + if not isinstance(provider_config, OpenAIProviderConfig): + continue + if openai_enqueue_quota_scope(provider_config) == quota_scope: + reserved += int(row["active_submission_tokens"]) + return reserved + + def finalize_indeterminate_submission_as_created(self, *, run_id: str, provider_batch_id: str, status: str) -> None: + with self.engine.begin() as conn: + row = ( + conn.execute( + select(SUBMISSION_INTENTS_TABLE).where( + (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + & (SUBMISSION_INTENTS_TABLE.c.status == "creating") + ) + ) + .mappings() + .all() + ) + if len(row) != 1: + raise ValueError("expected exactly one indeterminate submission intent") + intent = row[0] + submissions_payload = _decode_json(intent["submissions_json"]) + if not isinstance(submissions_payload, list): + raise ValueError("invalid persisted submission intent") + submissions = [ + PreparedSubmission( + item_id=str(value["item_id"]), + custom_id=str(value["custom_id"]), + submission_tokens=int(value["submission_tokens"]), + ) + for value in submissions_payload + if isinstance(value, dict) + ] + if len(submissions) != len(submissions_payload): + raise ValueError("invalid persisted submission intent") + conn.execute( + BATCHES_TABLE.insert(), + [ + { + "run_id": run_id, + "provider_batch_id": provider_batch_id, + "local_batch_id": f"recovered-{str(intent['intent_id'])[-8:]}", + "status": status, + "custom_ids_json": intent["custom_ids_json"], + "output_file_id": None, + "error_file_id": None, + } + ], + ) + statement = ( + update(ITEMS_TABLE) + .where( + and_(ITEMS_TABLE.c.run_id == bindparam("b_run_id"), ITEMS_TABLE.c.item_id == bindparam("b_item_id")) + ) + .values( + status=ItemStatus.SUBMITTED, + active_batch_id=bindparam("b_provider_batch_id"), + active_custom_id=bindparam("b_custom_id"), + active_submission_tokens=bindparam("b_submission_tokens"), + error_json=None, + ) + ) + conn.execute( + statement, + [ + { + "b_run_id": run_id, + "b_item_id": value.item_id, + "b_provider_batch_id": provider_batch_id, + "b_custom_id": value.custom_id, + "b_submission_tokens": value.submission_tokens, + } + for value in submissions + ], + ) + conn.execute( + update(SUBMISSION_INTENTS_TABLE) + .where(SUBMISSION_INTENTS_TABLE.c.intent_id == intent["intent_id"]) + .values(status="active", provider_batch_id=provider_batch_id) + ) + + @staticmethod + def _release_submission_intents(conn, *, intent_ids: list[str], status: str | None) -> None: # noqa: ANN001 + for intent_id in intent_ids: + row = ( + conn.execute( + select( + SUBMISSION_INTENTS_TABLE.c.quota_scope, + SUBMISSION_INTENTS_TABLE.c.submission_tokens, + SUBMISSION_INTENTS_TABLE.c.capacity_released, + ).where(SUBMISSION_INTENTS_TABLE.c.intent_id == intent_id) + ) + .mappings() + .first() + ) + if row is None or int(row["capacity_released"]) != 0: + continue + values: dict[str, object] = {"capacity_released": 1} + if status is not None: + values["status"] = status + updated = conn.execute( + update(SUBMISSION_INTENTS_TABLE) + .where( + (SUBMISSION_INTENTS_TABLE.c.intent_id == intent_id) + & (SUBMISSION_INTENTS_TABLE.c.capacity_released == 0) + ) + .values(**values) + ) + if updated.rowcount == 1 and row["quota_scope"] is not None: + conn.execute( + update(CAPACITY_SCOPES_TABLE) + .where(CAPACITY_SCOPES_TABLE.c.quota_scope == str(row["quota_scope"])) + .values(reserved_tokens=CAPACITY_SCOPES_TABLE.c.reserved_tokens - int(row["submission_tokens"])) + ) + def update_batch_status( self, *, @@ -622,7 +984,28 @@ def get_active_batches(self, *, run_id: str) -> list[ActiveBatchRecord]: .where( and_( BATCHES_TABLE.c.run_id == run_id, - BATCHES_TABLE.c.status.not_in(tuple(self.TERMINAL_BATCH_STATUSES)), + ( + BATCHES_TABLE.c.status.not_in(tuple(self.TERMINAL_BATCH_STATUSES)) + | exists( + select(ITEMS_TABLE.c.item_id).where( + and_( + ITEMS_TABLE.c.run_id == run_id, + ITEMS_TABLE.c.active_batch_id == BATCHES_TABLE.c.provider_batch_id, + ITEMS_TABLE.c.status == self.ACTIVE_ITEM_STATUS_SUBMITTED, + ) + ) + ) + | exists( + select(SUBMISSION_INTENTS_TABLE.c.intent_id).where( + and_( + SUBMISSION_INTENTS_TABLE.c.run_id == run_id, + SUBMISSION_INTENTS_TABLE.c.provider_batch_id + == BATCHES_TABLE.c.provider_batch_id, + SUBMISSION_INTENTS_TABLE.c.capacity_released == 0, + ) + ) + ) + ), ) ) .order_by(BATCHES_TABLE.c.local_batch_id) diff --git a/src/batchor/storage/sqlite_queries.py b/src/batchor/storage/sqlite_queries.py index 51819ec..4de712c 100644 --- a/src/batchor/storage/sqlite_queries.py +++ b/src/batchor/storage/sqlite_queries.py @@ -1,6 +1,6 @@ from __future__ import annotations -from sqlalchemy import func, select, update +from sqlalchemy import exists, func, select, update from sqlalchemy.engine import Connection, RowMapping from batchor.core.enums import ItemStatus, RunControlState, RunLifecycleStatus @@ -28,6 +28,7 @@ RUNS_TABLE, SQLITE_SCHEMA_VERSION, STORAGE_METADATA_TABLE, + SUBMISSION_INTENTS_TABLE, ) from batchor.storage.state import ( IngestCheckpoint, @@ -149,6 +150,7 @@ def _ensure_schema(self) -> None: } if "checkpoint_payload_json" not in ingest_columns: conn.exec_driver_sql("ALTER TABLE run_ingest_state ADD COLUMN checkpoint_payload_json TEXT") + self._backfill_missing_ingest_markers(conn) existing_schema_row = conn.execute( select(STORAGE_METADATA_TABLE.c.value).where(STORAGE_METADATA_TABLE.c.key == "schema_version") ).first() @@ -164,6 +166,44 @@ def _ensure_schema(self) -> None: .values(value=str(SQLITE_SCHEMA_VERSION)) ) + @staticmethod + def _backfill_missing_ingest_markers(conn: Connection) -> None: + """Conservatively migrate runs created before generic ingest markers.""" + rows = conn.execute( + select( + RUNS_TABLE.c.run_id, + select(func.count()) + .select_from(ITEMS_TABLE) + .where(ITEMS_TABLE.c.run_id == RUNS_TABLE.c.run_id) + .scalar_subquery() + .label("item_count"), + ).where( + ~exists( + select(RUN_INGEST_STATE_TABLE.c.run_id).where( + RUN_INGEST_STATE_TABLE.c.run_id == RUNS_TABLE.c.run_id + ) + ) + ) + ).mappings() + for row in rows: + conn.execute( + RUN_INGEST_STATE_TABLE.insert(), + [ + { + "run_id": str(row["run_id"]), + "source_kind": "unattached", + "source_ref": "", + "source_fingerprint": "", + "next_item_index": int(row["item_count"]), + "checkpoint_payload_json": None, + # No pre-v6 durable fact distinguishes a completed + # arbitrary iterable from one paused between chunks. + # Conservatively require explicit recovery for both. + "ingestion_complete": 0, + } + ], + ) + def _fetch_retry_state(self, conn: Connection, run_id: str) -> RetryBackoffState: row = ( conn.execute(select(RUN_RETRY_STATE_TABLE).where(RUN_RETRY_STATE_TABLE.c.run_id == run_id)).mappings().one() @@ -291,7 +331,23 @@ def _summary_for_run( .select_from(BATCHES_TABLE) .where( (BATCHES_TABLE.c.run_id == run_id) - & BATCHES_TABLE.c.status.not_in(tuple(self.TERMINAL_BATCH_STATUSES)) + & ( + BATCHES_TABLE.c.status.not_in(tuple(self.TERMINAL_BATCH_STATUSES)) + | exists( + select(ITEMS_TABLE.c.item_id).where( + (ITEMS_TABLE.c.run_id == run_id) + & (ITEMS_TABLE.c.active_batch_id == BATCHES_TABLE.c.provider_batch_id) + & (ITEMS_TABLE.c.status == self.ACTIVE_ITEM_STATUS_SUBMITTED) + ) + ) + | exists( + select(SUBMISSION_INTENTS_TABLE.c.intent_id).where( + (SUBMISSION_INTENTS_TABLE.c.run_id == run_id) + & (SUBMISSION_INTENTS_TABLE.c.provider_batch_id == BATCHES_TABLE.c.provider_batch_id) + & (SUBMISSION_INTENTS_TABLE.c.capacity_released == 0) + ) + ) + ) ) ).scalar_one() ) diff --git a/src/batchor/storage/sqlite_results.py b/src/batchor/storage/sqlite_results.py index a61074b..dfa5763 100644 --- a/src/batchor/storage/sqlite_results.py +++ b/src/batchor/storage/sqlite_results.py @@ -39,6 +39,7 @@ def mark_items_completed( and_( ITEMS_TABLE.c.run_id == bindparam("b_run_id"), ITEMS_TABLE.c.active_custom_id == bindparam("b_custom_id"), + ITEMS_TABLE.c.status == self.ACTIVE_ITEM_STATUS_SUBMITTED, ) ) .values( @@ -98,6 +99,7 @@ def mark_items_failed( and_( ITEMS_TABLE.c.run_id == run_id, ITEMS_TABLE.c.active_custom_id.in_(custom_ids), + ITEMS_TABLE.c.status == self.ACTIVE_ITEM_STATUS_SUBMITTED, ) ) ).mappings() @@ -280,7 +282,11 @@ def record_batch_retry_failure( ) -> RetryBackoffState: with self.engine.begin() as conn: current = self._fetch_retry_state(conn, run_id) - consecutive = current.consecutive_failures + 1 + # A success elsewhere must not clear this run-level delay. Once + # its deadline has elapsed, the next failure starts a fresh + # streak instead of retaining an unrelated historical exponent. + previous_is_active = current.next_retry_at is not None and current.next_retry_at > self._now() + consecutive = (current.consecutive_failures if previous_is_active else 0) + 1 total = current.total_failures + 1 backoff_sec = compute_backoff_delay( consecutive_failures=consecutive, diff --git a/src/batchor/storage/sqlite_schema.py b/src/batchor/storage/sqlite_schema.py index a89c8e7..6d6b82a 100644 --- a/src/batchor/storage/sqlite_schema.py +++ b/src/batchor/storage/sqlite_schema.py @@ -1,7 +1,7 @@ from sqlalchemy import Column, Float, Index, Integer, MetaData, String, Table, Text METADATA = MetaData() -SQLITE_SCHEMA_VERSION = 4 +SQLITE_SCHEMA_VERSION = 6 STORAGE_METADATA_TABLE = Table( "storage_metadata", @@ -100,3 +100,36 @@ Column("checkpoint_payload_json", Text, nullable=True), Column("ingestion_complete", Integer, nullable=False), ) + +# A submission intent is written before the provider-side create call. A +# ``creating`` row means the remote outcome is indeterminate and must not be +# silently retried after a crash. It also acts as the durable reservation for +# shared OpenAI enqueue capacity until the linked batch drains. +SUBMISSION_INTENTS_TABLE = Table( + "submission_intents", + METADATA, + Column("intent_id", String, primary_key=True), + Column("run_id", String, nullable=False), + Column("status", String, nullable=False), + Column("provider_batch_id", String, nullable=True), + Column("custom_ids_json", Text, nullable=False), + Column("item_ids_json", Text, nullable=False), + Column("submissions_json", Text, nullable=False), + Column("quota_scope", String, nullable=True), + Column("submission_tokens", Integer, nullable=False), + Column("capacity_released", Integer, nullable=False, default=0), + Column("created_at", String, nullable=False), +) +Index( + "ix_submission_intents_scope_active", + SUBMISSION_INTENTS_TABLE.c.quota_scope, + SUBMISSION_INTENTS_TABLE.c.capacity_released, +) +Index("ix_submission_intents_run_status", SUBMISSION_INTENTS_TABLE.c.run_id, SUBMISSION_INTENTS_TABLE.c.status) + +CAPACITY_SCOPES_TABLE = Table( + "capacity_scopes", + METADATA, + Column("quota_scope", String, primary_key=True), + Column("reserved_tokens", Integer, nullable=False), +) diff --git a/src/batchor/storage/state_models.py b/src/batchor/storage/state_models.py index 2228499..39ed794 100644 --- a/src/batchor/storage/state_models.py +++ b/src/batchor/storage/state_models.py @@ -663,6 +663,65 @@ def mark_items_submitted( """ ... + @abstractmethod + def begin_submission_intent( + self, + *, + run_id: str, + intent_id: str, + submissions: list[PreparedSubmission], + quota_scope: str | None, + submission_tokens: int, + capacity_limit: int | None, + ) -> bool: + """Persist intent before remote creation and reserve shared capacity. + + Returns ``False`` when the durable quota scope cannot admit the + reservation. A successful call is intentionally not retryable until + the intent is finalized or explicitly abandoned. + """ + ... + + @abstractmethod + def finalize_submission_intent( + self, + *, + run_id: str, + intent_id: str, + local_batch_id: str, + provider_batch_id: str, + status: str, + custom_ids: list[str], + submissions: list[PreparedSubmission], + ) -> None: + """Atomically register a remote batch and link its submitted items.""" + ... + + @abstractmethod + def abandon_submission_intent(self, *, intent_id: str) -> None: + """Release a pre-create intent after a known-safe remote failure.""" + ... + + @abstractmethod + def has_indeterminate_submission_intents(self, *, run_id: str) -> bool: + """Return whether recovery would risk duplicate remote execution.""" + ... + + @abstractmethod + def release_submission_capacity_for_batch(self, *, run_id: str, provider_batch_id: str) -> None: + """Release a finalized reservation once every linked item is drained.""" + ... + + @abstractmethod + def abandon_indeterminate_submission_intents(self, *, run_id: str) -> None: + """Mark creating intents operator-verified as not remotely created.""" + ... + + @abstractmethod + def finalize_indeterminate_submission_as_created(self, *, run_id: str, provider_batch_id: str, status: str) -> None: + """Atomically link the sole operator-confirmed creating intent.""" + ... + @abstractmethod def update_batch_status( self, diff --git a/tests/integration/test_batchor_runner.py b/tests/integration/test_batchor_runner.py index e29897b..48a06b2 100644 --- a/tests/integration/test_batchor_runner.py +++ b/tests/integration/test_batchor_runner.py @@ -35,6 +35,7 @@ RunLifecycleStatus, RunNotFinishedError, RunPausedError, + RunSubmissionIndeterminateError, SQLiteStorage, ) from batchor.providers.openai import OpenAIBatchProvider @@ -443,10 +444,11 @@ def estimate_request_tokens(self, request_line, *, chars_per_token): ), temp_root=artifact_root, ).get_run(run_id) - resumed.wait(poll_interval=0) - result = resumed.results()[0] - assert result.status is ItemStatus.COMPLETED - assert result.output_text == '{"label": "row1", "score": 0.9}' or result.output_text is not None + with pytest.raises(RunSubmissionIndeterminateError): + resumed.wait(poll_interval=0) + resumed.resolve_indeterminate_submission_as_not_created() + with pytest.raises(RunIngestionSourceRequiredError): + resumed.wait(poll_interval=0) def test_transient_poll_failures_do_not_block_new_submissions(tmp_path: Path) -> None: @@ -1279,15 +1281,18 @@ def test_sqlite_resume_retries_from_persisted_request_artifact(tmp_path: Path) - storage=storage, provider_factory=lambda _cfg: first_provider, ) - started = first_runner.start( - BatchJob( - items=[BatchItem(item_id="row1", payload={"text": "hello"})], - build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), - structured_output=ClassificationResult, - provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), - retry_policy=RetryPolicy(max_attempts=2, base_backoff_sec=0, max_backoff_sec=0), + run_id = "artifact_resume_indeterminate" + with pytest.raises(RunSubmissionIndeterminateError): + first_runner.start( + BatchJob( + items=[BatchItem(item_id="row1", payload={"text": "hello"})], + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + structured_output=ClassificationResult, + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), + retry_policy=RetryPolicy(max_attempts=2, base_backoff_sec=0, max_backoff_sec=0), + ), + run_id=run_id, ) - ) with storage.engine.begin() as conn: row = ( @@ -1298,7 +1303,7 @@ def test_sqlite_resume_retries_from_persisted_request_artifact(tmp_path: Path) - storage_sqlite.ITEMS_TABLE.c.request_artifact_line, storage_sqlite.ITEMS_TABLE.c.request_sha256, storage_sqlite.ITEMS_TABLE.c.status, - ).where(storage_sqlite.ITEMS_TABLE.c.run_id == started.run_id) + ).where(storage_sqlite.ITEMS_TABLE.c.run_id == run_id) ) .mappings() .one() @@ -1307,7 +1312,7 @@ def test_sqlite_resume_retries_from_persisted_request_artifact(tmp_path: Path) - assert row["request_artifact_path"] is not None assert row["request_artifact_line"] == 1 assert row["request_sha256"] is not None - assert row["status"] == ItemStatus.PENDING + assert row["status"] == ItemStatus.QUEUED_LOCAL second_provider = _ArtifactOnlyBatchProvider( record_factory=lambda custom_id: _success_record(json.dumps({"label": custom_id.split(":")[0], "score": 0.9})) @@ -1315,12 +1320,12 @@ def test_sqlite_resume_retries_from_persisted_request_artifact(tmp_path: Path) - resumed = BatchRunner( storage=SQLiteStorage(path=storage.path), provider_factory=lambda _cfg: second_provider, - ).get_run(started.run_id) - resumed.wait(poll_interval=0) - result = resumed.results()[0] - assert result.status is ItemStatus.COMPLETED - assert result.output is not None - assert result.output.label == "row1" + ).get_run(run_id) + with pytest.raises(RunSubmissionIndeterminateError): + resumed.wait(poll_interval=0) + resumed.resolve_indeterminate_submission_as_not_created() + with pytest.raises(RunIngestionSourceRequiredError): + resumed.wait(poll_interval=0) def test_completed_run_can_prune_persisted_request_artifacts(tmp_path: Path) -> None: diff --git a/tests/unit/test_batchor_runtime_ingestion.py b/tests/unit/test_batchor_runtime_ingestion.py index 83f3927..3b9bc38 100644 --- a/tests/unit/test_batchor_runtime_ingestion.py +++ b/tests/unit/test_batchor_runtime_ingestion.py @@ -1186,3 +1186,148 @@ def test_materialize_item_chunks_rejects_non_resumable_source_with_nonzero_start with pytest.raises(ValueError, match="non-resumable item sources cannot start from a checkpoint"): list(materialize_item_chunks(job, start_index=1)) + + +def test_ingestion_time_budget_persists_a_partial_slice_before_old_chunk_limit() -> None: + job = BatchJob( + items=[BatchItem(item_id=f"row{index}", payload="x") for index in range(5)], + 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 = "time_sliced_ingestion" + storage.create_run(run_id=run_id, config=config, items=[]) + chunk_sizes: list[int] = [] + deps = IngestionDeps( + state=storage, + emit_event=lambda _event_type, **kwargs: chunk_sizes.append(kwargs["data"]["chunk_item_count"]), + submit_pending_items=lambda _run_id, _context: 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + work_monotonic=_MonotonicClock(0.0, 0.1, 0.2, 0.3, 0.3, 0.4, 0.5), + work_slice_max_items=1000, + work_slice_max_seconds=0.25, + ) + 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 chunk_sizes == [3, 2] + assert [record.item_id for record in storage.get_item_records(run_id=run_id)] == [ + "row0", + "row1", + "row2", + "row3", + "row4", + ] + + +def test_manual_pause_stops_noncheckpointed_source_and_same_process_resume_is_exactly_once() -> None: + source = (BatchItem(item_id=f"row{index}", payload="x") for index in range(5)) + job = BatchJob( + items=source, + 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 = "manual_pause_noncheckpointed" + storage.create_run(run_id=run_id, config=config, items=[]) + submitted: list[str] = [] + + def pause_after_first_slice(submitted_run_id: str, _context: object) -> int: + submitted.append(submitted_run_id) + if len(submitted) > 1: + return 0 + storage.set_run_control_state( + run_id=submitted_run_id, + control_state=RunControlState.PAUSED, + control_reason="manual", + ) + return 0 + + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=pause_after_first_slice, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + work_slice_max_items=2, + ) + 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, + ) + checkpoint = storage.get_ingest_checkpoint(run_id=run_id) + assert checkpoint is not None + assert checkpoint.next_item_index == 2 + assert checkpoint.ingestion_complete is False + assert [record.item_id for record in storage.get_item_records(run_id=run_id)] == ["row0", "row1"] + + storage.set_run_control_state(run_id=run_id, control_state=RunControlState.RUNNING) + resume_existing_run(deps, run_id=run_id, job=job, config=config, context=context) + + assert [record.item_id for record in storage.get_item_records(run_id=run_id)] == [ + "row0", + "row1", + "row2", + "row3", + "row4", + ] + assert storage.get_ingest_checkpoint(run_id=run_id).ingestion_complete is True # type: ignore[union-attr] + assert submitted == [run_id, run_id, run_id] + + +def test_submission_backoff_does_not_suppress_ingestion_reconciliation() -> None: + job = BatchJob( + items=[BatchItem(item_id=f"row{index}", payload="x") for index in range(1000)], + build_prompt=lambda item: PromptParts(prompt=item.payload), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1", poll_interval_sec=1), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + run_id = "backoff_still_polls" + 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=[], + ) + storage.record_batch_retry_failure( + run_id=run_id, error_class="provider_unavailable", base_delay_sec=10, max_delay_sec=10 + ) + 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), + ) + 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 == ["poll"] diff --git a/tests/unit/test_batchor_runtime_polling.py b/tests/unit/test_batchor_runtime_polling.py index 2471004..692e92d 100644 --- a/tests/unit/test_batchor_runtime_polling.py +++ b/tests/unit/test_batchor_runtime_polling.py @@ -16,6 +16,7 @@ PromptParts, RetryPolicy, RunControlState, + SQLiteStorage, ) from batchor.providers.openai import OpenAIBatchProvider from batchor.runtime.context import RunContext, build_persisted_config, build_run_context @@ -56,6 +57,7 @@ def __init__( output_content: str = "", error_content: str = "", poll_exception: Exception | None = None, + download_exception: Exception | None = None, ) -> None: self.batch_factory = batch_factory or ( lambda batch_id: { @@ -68,6 +70,7 @@ def __init__( self.output_content = output_content self.error_content = error_content self.poll_exception = poll_exception + self.download_exception = download_exception self._parser = OpenAIBatchProvider( OpenAIProviderConfig(api_key="k", model="gpt-4.1"), client=object(), @@ -81,6 +84,10 @@ def get_batch(self, batch_id: str) -> dict[str, object]: return self.batch_factory(batch_id) def download_file_content(self, file_id: str) -> str: + if self.download_exception is not None: + exc = self.download_exception + self.download_exception = None + raise exc if file_id.startswith("output_"): return self.output_content return self.error_content @@ -351,6 +358,125 @@ def test_consume_completed_batch_retries_missing_output_rows_without_consuming_a assert record.error.error_class == "batch_output_missing_row" +def test_terminal_batch_is_replayed_after_download_failure_across_sqlite_reopen(tmp_path: Path) -> None: + provider = _FakePollingProvider( + output_content=json.dumps(_success_record("done") | {"custom_id": "row1:a1"}) + "\n", + download_exception=ConnectionError("download interrupted"), + ) + storage, artifact_store, deps, context, run_id = _polling_setup(tmp_path, provider=provider) + sqlite_path = tmp_path / "durable.sqlite3" + durable = SQLiteStorage(path=sqlite_path) + config = storage.get_run_config(run_id=run_id) + durable.create_run( + run_id=run_id, + config=config, + items=[ + MaterializedItem( + item_id="row1", + item_index=0, + payload={"text": "hello"}, + metadata={}, + prompt="hello", + ) + ], + ) + durable.register_batch( + run_id=run_id, + local_batch_id="local_batch_0", + provider_batch_id="batch_0", + status="submitted", + custom_ids=["row1:a1"], + ) + durable.mark_items_submitted( + run_id=run_id, + provider_batch_id="batch_0", + submissions=[PreparedSubmission(item_id="row1", custom_id="row1:a1", submission_tokens=1)], + ) + durable_deps = PollingDeps(state=durable, artifact_store=artifact_store, emit_event=lambda *_args, **_kwargs: None) + + try: + poll_once(durable_deps, run_id=run_id, context=context) + except ConnectionError: + pass + else: # pragma: no cover - makes the crash window assertion explicit + raise AssertionError("expected the initial result download to fail") + + # The remote terminal status was committed before the failure, but the + # still-submitted item keeps the batch locally active and therefore + # replayable by a fresh storage instance. + assert durable.get_item_records(run_id=run_id)[0].status is ItemStatus.SUBMITTED + assert len(durable.get_active_batches(run_id=run_id)) == 1 + durable.close() + + reopened = SQLiteStorage(path=sqlite_path) + replay_deps = PollingDeps(state=reopened, artifact_store=artifact_store, emit_event=lambda *_args, **_kwargs: None) + poll_once(replay_deps, run_id=run_id, context=context) + + record = reopened.get_item_records(run_id=run_id)[0] + assert record.status is ItemStatus.COMPLETED + assert record.attempt_count == 1 + assert reopened.get_active_batches(run_id=run_id) == [] + + +def test_terminal_replay_skips_already_consumed_rows_after_partial_item_mutation( + tmp_path: Path, + monkeypatch, +) -> None: + output = "\n".join( + [ + json.dumps(_success_record("done") | {"custom_id": "row1:a1"}), + json.dumps( + { + "custom_id": "row2:a1", + "response": {"status_code": 500, "body": {"error": {"message": "retry"}}}, + } + ), + ] + ) + provider = _FakePollingProvider(output_content=output + "\n") + storage, _artifact_store, deps, context, run_id = _polling_setup(tmp_path, provider=provider) + storage.append_items( + run_id=run_id, + items=[ + MaterializedItem( + item_id="row2", + item_index=1, + payload={"text": "later"}, + metadata={}, + prompt="later", + ) + ], + ) + storage.register_batch( + run_id=run_id, + local_batch_id="local_batch_0", + provider_batch_id="batch_0", + status="submitted", + custom_ids=["row1:a1", "row2:a1"], + ) + storage.mark_items_submitted( + run_id=run_id, + provider_batch_id="batch_0", + submissions=[PreparedSubmission(item_id="row2", custom_id="row2:a1", submission_tokens=1)], + ) + original_mark_failed = storage.mark_items_failed + monkeypatch.setattr(storage, "mark_items_failed", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("crash"))) + try: + poll_once(deps, run_id=run_id, context=context) + except RuntimeError as exc: + assert str(exc) == "crash" + else: # pragma: no cover + raise AssertionError("expected partial consumption failure") + monkeypatch.setattr(storage, "mark_items_failed", original_mark_failed) + + poll_once(deps, run_id=run_id, context=context) + + records = storage.get_item_records(run_id=run_id) + assert [record.status for record in records] == [ItemStatus.COMPLETED, ItemStatus.FAILED_RETRYABLE] + assert records[0].attempt_count == 1 + assert records[1].attempt_count == 1 + + def test_poll_once_records_backoff_for_item_level_insufficient_quota(tmp_path: Path) -> None: quota_record = { "custom_id": "row1:a1", @@ -384,6 +510,56 @@ def test_poll_once_records_backoff_for_item_level_insufficient_quota(tmp_path: P assert summary.backoff_remaining_sec > 0 +def test_completed_batch_does_not_clear_existing_retry_backoff(tmp_path: Path) -> None: + provider = _FakePollingProvider( + output_content=json.dumps(_success_record("done") | {"custom_id": "row1:a1"}) + "\n", + ) + storage, _artifact_store, deps, context, run_id = _polling_setup( + tmp_path, + provider=provider, + retry_policy=RetryPolicy(max_attempts=2, base_backoff_sec=60.0, max_backoff_sec=60.0), + ) + storage.record_batch_retry_failure( + run_id=run_id, + error_class="batch_create_connection_error", + base_delay_sec=60.0, + max_delay_sec=60.0, + ) + + poll_once(deps, run_id=run_id, context=context) + + assert storage.get_item_records(run_id=run_id)[0].status is ItemStatus.COMPLETED + assert storage.get_run_summary(run_id=run_id).backoff_remaining_sec > 0 + + +def test_cancellation_clears_retry_backoff_after_final_batch_drains(tmp_path: Path) -> None: + quota_record = { + "custom_id": "row1:a1", + "response": { + "status_code": 429, + "body": {"error": {"code": "insufficient_quota", "message": "quota exhausted"}}, + }, + } + provider = _FakePollingProvider(output_content=json.dumps(quota_record) + "\n") + storage, _artifact_store, deps, context, run_id = _polling_setup( + tmp_path, + provider=provider, + retry_policy=RetryPolicy(max_attempts=2, base_backoff_sec=60.0, max_backoff_sec=60.0), + ) + storage.set_run_control_state(run_id=run_id, control_state=RunControlState.CANCEL_REQUESTED) + + summary = refresh_run( + deps, + run_id=run_id, + context=context, + submit_pending_items=lambda _run_id, _context: 0, + ) + + assert summary.status.value == "completed_with_failures" + assert summary.backoff_remaining_sec == 0.0 + assert storage.get_item_records(run_id=run_id)[0].status is ItemStatus.FAILED_PERMANENT + + def test_consume_completed_batch_reports_item_level_quota_without_pausing(tmp_path: Path) -> None: quota_record = { "custom_id": "row1:a1", diff --git a/tests/unit/test_batchor_runtime_submission.py b/tests/unit/test_batchor_runtime_submission.py index 6c23687..829264c 100644 --- a/tests/unit/test_batchor_runtime_submission.py +++ b/tests/unit/test_batchor_runtime_submission.py @@ -3,6 +3,8 @@ import json from pathlib import Path +import pytest + from batchor import ( BatchItem, BatchJob, @@ -16,6 +18,7 @@ PromptParts, RetryPolicy, RunControlState, + RunSubmissionIndeterminateError, ) from batchor.runtime.artifacts import request_sha256 from batchor.runtime.context import build_persisted_config, build_run_context @@ -337,7 +340,7 @@ def test_submit_pending_items_marks_oversized_rows_and_releases_unsent_items(tmp assert provider.created_batches == ["batch_0"] -def test_submit_pending_items_cleans_up_uploaded_input_on_retryable_create_failure(tmp_path: Path) -> None: +def test_submit_pending_items_preserves_indeterminate_create_intent(tmp_path: Path) -> None: provider = _FakeSubmissionProvider( create_failures=[RuntimeError("temporary service unavailable")], ) @@ -367,13 +370,16 @@ def test_submit_pending_items_cleans_up_uploaded_input_on_retryable_create_failu items=[_materialized_text_item("row1", 0, "hello")], ) - submitted = submit_pending_items(deps, run_id=run_id, context=context) + with pytest.raises(RunSubmissionIndeterminateError): + submit_pending_items(deps, run_id=run_id, context=context) record = storage.get_item_records(run_id=run_id)[0] - assert submitted == 0 - assert record.status is ItemStatus.PENDING - assert provider.deleted_files == ["file_0"] - assert storage.get_run_summary(run_id=run_id).backoff_remaining_sec > 0 + assert record.status is ItemStatus.QUEUED_LOCAL + assert provider.deleted_files == [] + assert storage.has_indeterminate_submission_intents(run_id=run_id) + storage.abandon_indeterminate_submission_intents(run_id=run_id) + storage.requeue_local_items(run_id=run_id) + assert storage.get_item_records(run_id=run_id)[0].status is ItemStatus.PENDING def test_submit_pending_items_auto_pauses_on_insufficient_quota_create_failure(tmp_path: Path) -> None: diff --git a/tests/unit/test_batchor_sqlite_storage_flow.py b/tests/unit/test_batchor_sqlite_storage_flow.py index 1327ce5..1566050 100644 --- a/tests/unit/test_batchor_sqlite_storage_flow.py +++ b/tests/unit/test_batchor_sqlite_storage_flow.py @@ -1,7 +1,9 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from pathlib import Path +from threading import Barrier import pytest from sqlalchemy import select @@ -169,6 +171,75 @@ def test_sqlite_storage_submission_and_terminal_summary(tmp_path: Path) -> None: assert provider_config.enqueue_limits.max_batch_enqueued_tokens == 500 +def test_sqlite_submission_capacity_reservation_is_atomic_across_store_instances(tmp_path: Path) -> None: + path = tmp_path / "shared.sqlite3" + first = SQLiteStorage(path=path) + second = SQLiteStorage(path=path) + first.create_run(run_id="capacity_a", config=_config(), items=_items()[:1]) + second.create_run(run_id="capacity_b", config=_config(), items=_items()[:1]) + barrier = Barrier(2) + + def reserve(storage: SQLiteStorage, run_id: str) -> bool: + barrier.wait() + return storage.begin_submission_intent( + run_id=run_id, + intent_id=f"intent-{run_id}", + submissions=[PreparedSubmission(item_id="row1", custom_id=f"{run_id}:a1", submission_tokens=10)], + quota_scope="openai:gpt-4.1:team-a", + submission_tokens=10, + capacity_limit=10, + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(lambda args: reserve(*args), [(first, "capacity_a"), (second, "capacity_b")])) + + assert sorted(results) == [False, True] + + +def test_sqlite_upgrade_backfills_empty_legacy_run_as_ingestion_incomplete(tmp_path: Path) -> None: + path = tmp_path / "legacy-empty.sqlite3" + storage = SQLiteStorage(path=path) + storage.create_run(run_id="legacy_empty", config=_config(), items=[]) + with storage.engine.begin() as conn: + conn.execute( + storage_sqlite.RUN_INGEST_STATE_TABLE.delete().where( + storage_sqlite.RUN_INGEST_STATE_TABLE.c.run_id == "legacy_empty" + ) + ) + conn.execute( + storage_sqlite.RUNS_TABLE.update() + .where(storage_sqlite.RUNS_TABLE.c.run_id == "legacy_empty") + .values(status=RunLifecycleStatus.COMPLETED) + ) + storage.close() + + upgraded = SQLiteStorage(path=path) + checkpoint = upgraded.get_ingest_checkpoint(run_id="legacy_empty") + assert checkpoint is not None + assert checkpoint.ingestion_complete is False + assert upgraded.get_run_summary(run_id="legacy_empty").status is RunLifecycleStatus.RUNNING + + +def test_sqlite_upgrade_backfills_partial_legacy_run_as_ingestion_incomplete(tmp_path: Path) -> None: + path = tmp_path / "legacy-partial.sqlite3" + storage = SQLiteStorage(path=path) + storage.create_run(run_id="legacy_partial", config=_config(), items=_items()[:1]) + with storage.engine.begin() as conn: + conn.execute( + storage_sqlite.RUN_INGEST_STATE_TABLE.delete().where( + storage_sqlite.RUN_INGEST_STATE_TABLE.c.run_id == "legacy_partial" + ) + ) + storage.close() + + upgraded = SQLiteStorage(path=path) + checkpoint = upgraded.get_ingest_checkpoint(run_id="legacy_partial") + assert checkpoint is not None + assert checkpoint.next_item_index == 1 + assert checkpoint.ingestion_complete is False + assert upgraded.get_run_summary(run_id="legacy_partial").status is RunLifecycleStatus.RUNNING + + def test_sqlite_storage_reset_and_backoff(tmp_path: Path) -> None: clock = _FrozenClock() storage = SQLiteStorage(path=tmp_path / "backoff.sqlite3", now=clock.now) diff --git a/tests/unit/test_batchor_storage.py b/tests/unit/test_batchor_storage.py index 8b41547..66f96a2 100644 --- a/tests/unit/test_batchor_storage.py +++ b/tests/unit/test_batchor_storage.py @@ -20,7 +20,7 @@ def test_sqlite_storage_explicit_path_overrides_name(tmp_path: Path) -> None: def test_sqlite_storage_exposes_schema_version(tmp_path: Path) -> None: storage = SQLiteStorage(path=tmp_path / "versioned.sqlite3") - assert storage.schema_version == 4 + assert storage.schema_version == 6 def test_sqlite_storage_migrates_missing_control_reason_column(tmp_path: Path) -> None: diff --git a/tests/unit/test_batchor_storage_contracts.py b/tests/unit/test_batchor_storage_contracts.py index 7f5e73a..0a8c307 100644 --- a/tests/unit/test_batchor_storage_contracts.py +++ b/tests/unit/test_batchor_storage_contracts.py @@ -8,6 +8,7 @@ from sqlalchemy import text from batchor import ( + ItemStatus, MemoryStateStore, OpenAIEnqueueLimitConfig, OpenAIProviderConfig, @@ -227,6 +228,99 @@ def test_storage_contract_incomplete_ingestion_prevents_terminal_status(storage) assert storage.get_run_summary(run_id="run_incomplete").status is RunLifecycleStatus.COMPLETED +def test_storage_contract_attaches_operator_confirmed_indeterminate_batch(storage) -> None: + storage.create_run(run_id="run_indeterminate", config=_config(), items=_items()[:1]) + claimed = storage.claim_items_for_submission(run_id="run_indeterminate", max_attempts=2) + submissions = [ + PreparedSubmission( + item_id=claimed[0].item_id, + custom_id="row1:a1", + submission_tokens=10, + ) + ] + assert storage.begin_submission_intent( + run_id="run_indeterminate", + intent_id="intent-indeterminate", + submissions=submissions, + quota_scope="openai:gpt-4.1:contract", + submission_tokens=10, + capacity_limit=10, + ) + assert storage.has_indeterminate_submission_intents(run_id="run_indeterminate") + + storage.finalize_indeterminate_submission_as_created( + run_id="run_indeterminate", + provider_batch_id="provider-recovered", + status="submitted", + ) + + assert not storage.has_indeterminate_submission_intents(run_id="run_indeterminate") + assert storage.get_submitted_custom_ids_for_batch( + run_id="run_indeterminate", + provider_batch_id="provider-recovered", + ) == ["row1:a1"] + assert storage.get_active_submitted_token_estimate(run_id="run_indeterminate") == 10 + + +def test_storage_contract_not_created_resolution_requeues_claimed_items(storage) -> None: + storage.create_run(run_id="run_not_created", config=_config(), items=_items()) + claimed = storage.claim_items_for_submission(run_id="run_not_created", max_attempts=2) + assert storage.begin_submission_intent( + run_id="run_not_created", + intent_id="intent-not-created", + submissions=[ + PreparedSubmission( + item_id=claimed[0].item_id, + custom_id="row1:a1", + submission_tokens=10, + ) + ], + quota_scope="openai:gpt-4.1:default", + submission_tokens=10, + capacity_limit=10, + ) + + storage.abandon_indeterminate_submission_intents(run_id="run_not_created") + + reclaimed = storage.claim_items_for_submission(run_id="run_not_created", max_attempts=2) + assert [item.item_id for item in reclaimed] == ["row1"] + records = {record.item_id: record for record in storage.get_item_records(run_id="run_not_created")} + assert records["row2"].status is ItemStatus.QUEUED_LOCAL + + +def test_storage_contract_capacity_includes_legacy_submitted_items(storage) -> None: + storage.create_run(run_id="legacy_active", config=_config(), items=_items()[:1]) + claimed = storage.claim_items_for_submission(run_id="legacy_active", max_attempts=2) + storage.register_batch( + run_id="legacy_active", + local_batch_id="legacy-local", + provider_batch_id="legacy-provider", + status="submitted", + custom_ids=["row1:a1"], + ) + storage.mark_items_submitted( + run_id="legacy_active", + provider_batch_id="legacy-provider", + submissions=[ + PreparedSubmission( + item_id=claimed[0].item_id, + custom_id="row1:a1", + submission_tokens=10, + ) + ], + ) + storage.create_run(run_id="new_capacity", config=_config(), items=_items()[:1]) + + assert not storage.begin_submission_intent( + run_id="new_capacity", + intent_id="intent-new-capacity", + submissions=[PreparedSubmission(item_id="row1", custom_id="row1:a1", submission_tokens=1)], + quota_scope="openai:gpt-4.1:default", + submission_tokens=1, + capacity_limit=10, + ) + + def test_storage_contract_artifact_pointers_and_summary_rehydration(storage) -> None: storage.create_run(run_id="run_2", config=_config(), items=_items()[:1]) storage.record_request_artifacts( diff --git a/uv.lock b/uv.lock index 241b67e..29c6ceb 100644 --- a/uv.lock +++ b/uv.lock @@ -81,7 +81,7 @@ wheels = [ [[package]] name = "batchor" -version = "0.0.6" +version = "0.0.7" source = { editable = "." } dependencies = [ { name = "openai" }, From 757e9cda7faeaff25f5aa3c8182d22fcf877285d Mon Sep 17 00:00:00 2001 From: YW <58594437+AnsonDev42@users.noreply.github.com> Date: Sat, 18 Jul 2026 22:56:25 +0100 Subject: [PATCH 2/3] fix(runtime): preserve poll deadlines Keep stale terminal-error replays idempotent across storage adapters. --- src/batchor/runtime/ingestion.py | 32 +++++++++++-- src/batchor/runtime/runner.py | 3 ++ src/batchor/storage/memory.py | 7 ++- src/batchor/storage/sqlite_results.py | 19 ++++++-- tests/integration/test_batchor_runner.py | 48 ++++++++++++++++++++ tests/unit/test_batchor_storage_contracts.py | 47 +++++++++++++++++++ 6 files changed, 147 insertions(+), 9 deletions(-) diff --git a/src/batchor/runtime/ingestion.py b/src/batchor/runtime/ingestion.py index 46f5c11..d2a2adc 100644 --- a/src/batchor/runtime/ingestion.py +++ b/src/batchor/runtime/ingestion.py @@ -23,7 +23,12 @@ ) -def _noop_poll_active_batches(_run_id: str, _context: RunContext) -> None: +def _noop_poll_active_batches( + _run_id: str, + _context: RunContext, + *, + deadline: float | None = None, +) -> None: return @@ -47,7 +52,7 @@ class IngestionDeps: emit_event: Callable[..., None] 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 + poll_active_batches: Callable[..., None] = _noop_poll_active_batches monotonic: Callable[[], float] = time.monotonic work_slice_max_items: int = 1000 work_slice_max_seconds: float | None = None @@ -116,7 +121,7 @@ def resume_existing_run( if _deadline_reached(deps, deadline): return if control_state is RunControlState.RUNNING and deps.state.get_active_batches(run_id=run_id): - deps.poll_active_batches(run_id, context) + _poll_active_batches(deps, run_id=run_id, context=context, deadline=deadline) summary = deps.state.get_run_summary(run_id=run_id) if summary.backoff_remaining_sec > 0: return @@ -479,7 +484,7 @@ def _persist_and_advance_boundary( ): if _deadline_reached(deps, deadline): return False, last_poll_at - deps.poll_active_batches(run_id, context) + _poll_active_batches(deps, run_id=run_id, context=context, deadline=deadline) last_poll_at = now if _deadline_reached(deps, deadline): return False, last_poll_at @@ -504,6 +509,25 @@ def _persist_and_advance_boundary( return True, last_poll_at +def _poll_active_batches( + deps: IngestionDeps, + *, + run_id: str, + context: RunContext, + deadline: float | None, +) -> None: + """Reconcile active batches without dropping a wait-cycle deadline. + + The optional deadline keeps older internal test adapters usable for + ordinary ingestion while ensuring deadline-bound refreshes pass the same + absolute limit down to the provider-polling seam. + """ + if deadline is None: + deps.poll_active_batches(run_id, context) + return + deps.poll_active_batches(run_id, context, deadline=deadline) + + def materialize_item_chunks( job: BatchJob[Any, BaseModel], *, diff --git a/src/batchor/runtime/runner.py b/src/batchor/runtime/runner.py index 2a7a55e..b047d10 100644 --- a/src/batchor/runtime/runner.py +++ b/src/batchor/runtime/runner.py @@ -587,11 +587,14 @@ def _poll_active_batches( self, run_id: str, context: RunContext, + *, + deadline: float | None = None, ) -> None: poll_once( self._polling_deps, run_id=run_id, context=context, + deadline=deadline, ) def _results_for_run(self, run_id: str) -> list[BatchResultItem]: diff --git a/src/batchor/storage/memory.py b/src/batchor/storage/memory.py index be8bac2..6b55566 100644 --- a/src/batchor/storage/memory.py +++ b/src/batchor/storage/memory.py @@ -705,7 +705,12 @@ def mark_items_failed( ) -> None: run = self._get_run(run_id) for failure in failures: - item = self._item_for_custom_id(run, failure.custom_id) + try: + item = self._item_for_custom_id(run, failure.custom_id) + except KeyError: + # Replaying a provider-terminal artifact after another + # consumer already cleared the active correlation is safe. + continue if item.status is not ItemStatus.SUBMITTED: continue if failure.count_attempt: diff --git a/src/batchor/storage/sqlite_results.py b/src/batchor/storage/sqlite_results.py index dfa5763..d00fba6 100644 --- a/src/batchor/storage/sqlite_results.py +++ b/src/batchor/storage/sqlite_results.py @@ -111,8 +111,13 @@ def mark_items_failed( for row in rows } payloads: list[dict[str, object | None]] = [] - for index, failure in enumerate(failures): - current = attempts_by_custom_id[failure.custom_id] + for failure in failures: + # Another consumer may have durably consumed this terminal + # row after this caller parsed the same provider artifact. + # Such stale replay is a no-op, not an invalid transition. + current = attempts_by_custom_id.get(failure.custom_id) + if current is None: + continue attempt_count = int(current["attempt_count"]) if failure.count_attempt: attempt_count += 1 @@ -126,10 +131,12 @@ def mark_items_failed( { "b_run_id": run_id, "b_item_id": str(current["item_id"]), + "b_active_custom_id": failure.custom_id, + "b_expected_attempt_count": int(current["attempt_count"]), "b_attempt_count": attempt_count, "b_status": status, "b_terminal_result_sequence": ( - next_sequence + index if status in self.TERMINAL_ITEM_STATUSES else None + next_sequence + len(payloads) if status in self.TERMINAL_ITEM_STATUSES else None ), "b_terminalized_at": (terminalized_at if status in self.TERMINAL_ITEM_STATUSES else None), "b_error_json": _encode_json(serialize_item_failure(failure.error)), @@ -141,6 +148,9 @@ def mark_items_failed( and_( ITEMS_TABLE.c.run_id == bindparam("b_run_id"), ITEMS_TABLE.c.item_id == bindparam("b_item_id"), + ITEMS_TABLE.c.active_custom_id == bindparam("b_active_custom_id"), + ITEMS_TABLE.c.attempt_count == bindparam("b_expected_attempt_count"), + ITEMS_TABLE.c.status == self.ACTIVE_ITEM_STATUS_SUBMITTED, ) ) .values( @@ -154,7 +164,8 @@ def mark_items_failed( active_submission_tokens=0, ) ) - conn.execute(statement, payloads) + if payloads: + conn.execute(statement, payloads) def mark_queued_items_failed( self, diff --git a/tests/integration/test_batchor_runner.py b/tests/integration/test_batchor_runner.py index 48a06b2..4e6b40b 100644 --- a/tests/integration/test_batchor_runner.py +++ b/tests/integration/test_batchor_runner.py @@ -624,6 +624,54 @@ def test_pause_blocks_polling_and_wait_raises_paused_error(tmp_path: Path) -> No assert provider.created_batches == ["batch_0", "batch_1"] +def test_wait_deadline_reaches_resume_boundary_polling(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Attached incomplete ingestion must not begin a provider poll after expiry.""" + path = tmp_path / "items.jsonl" + path.write_text(json.dumps({"id": "row1", "text": "one"}) + "\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 {}, + ) + provider = _ControlledBatchProvider(record_factory=lambda custom_id: _success_record(f"text:{custom_id}")) + runner = BatchRunner( + storage="memory", + provider_factory=lambda _cfg: provider, + temp_root=tmp_path, + ) + run = runner.start( + BatchJob( + items=source, + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), + ) + ) + checkpoint = runner.state.get_ingest_checkpoint(run_id=run.run_id) + assert checkpoint is not None + runner.state.update_ingest_checkpoint( + run_id=run.run_id, + next_item_index=checkpoint.next_item_index, + checkpoint_payload=checkpoint.checkpoint_payload, + ingestion_complete=False, + ) + + calls = 0 + + def wait_clock() -> float: + nonlocal calls + calls += 1 + return 0.0 if calls < 5 else 1.0 + + monkeypatch.setattr("batchor.runtime.run_handle.time.monotonic", wait_clock) + monkeypatch.setattr("batchor.runtime.polling.monotonic", lambda: 1.0) + + with pytest.raises(TimeoutError): + run.wait(timeout=1, poll_interval=0) + + assert provider.polled_batches == [] + assert runner.state.get_active_batches(run_id=run.run_id) + + def test_insufficient_quota_create_failure_auto_pauses_and_resume_continues(tmp_path: Path) -> None: provider = _FakeBatchProvider( record_factory=lambda custom_id: _success_record(f"text:{custom_id}"), diff --git a/tests/unit/test_batchor_storage_contracts.py b/tests/unit/test_batchor_storage_contracts.py index 0a8c307..de6f845 100644 --- a/tests/unit/test_batchor_storage_contracts.py +++ b/tests/unit/test_batchor_storage_contracts.py @@ -424,6 +424,53 @@ def test_storage_contract_completion_counts_consumed_attempts(storage) -> None: assert records[0].attempt_count == 1 +def test_storage_contract_stale_terminal_failure_replay_is_idempotent(storage) -> None: + """A second consumer can safely replay a result parsed before the first committed.""" + storage.create_run(run_id="run_terminal_failure_replay", config=_config(), items=_items()[:1]) + claimed = storage.claim_items_for_submission(run_id="run_terminal_failure_replay", max_attempts=1) + storage.register_batch( + run_id="run_terminal_failure_replay", + local_batch_id="local_1", + provider_batch_id="provider_1", + status="completed", + custom_ids=["row1:a1"], + ) + storage.mark_items_submitted( + run_id="run_terminal_failure_replay", + provider_batch_id="provider_1", + submissions=[PreparedSubmission(item_id=claimed[0].item_id, custom_id="row1:a1", submission_tokens=10)], + ) + failure = ItemFailureRecord( + custom_id="row1:a1", + error=ItemFailure( + error_class="provider_item_error", + message="terminal provider error", + retryable=True, + raw_error={"status": 500}, + ), + count_attempt=True, + ) + + storage.mark_items_failed( + run_id="run_terminal_failure_replay", + failures=[failure], + max_attempts=1, + ) + # This represents a stale concurrent consumer that parsed the same + # terminal artifact before the first consumer committed its transition. + storage.mark_items_failed( + run_id="run_terminal_failure_replay", + failures=[failure], + max_attempts=1, + ) + + record = storage.get_item_records(run_id="run_terminal_failure_replay")[0] + assert record.status is ItemStatus.FAILED_PERMANENT + assert record.attempt_count == 1 + assert record.error is not None + assert record.error.error_class == "provider_item_error" + + def test_storage_contract_retry_then_completion_counts_all_consumed_attempts(storage) -> None: storage.create_run(run_id="run_retry_then_success", config=_config(), items=_items()[:1]) claimed = storage.claim_items_for_submission(run_id="run_retry_then_success", max_attempts=2) From 55aac3af5fd5c1d9cc00f6105bf324bfceff0362 Mon Sep 17 00:00:00 2001 From: YW <58594437+AnsonDev42@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:03:00 +0100 Subject: [PATCH 3/3] fix(runtime): honor poll-overrun deadlines --- src/batchor/runtime/ingestion.py | 27 +++- tests/integration/test_batchor_runner.py | 16 +- tests/unit/test_batchor_runtime_ingestion.py | 153 +++++++++++++++++++ 3 files changed, 187 insertions(+), 9 deletions(-) diff --git a/src/batchor/runtime/ingestion.py b/src/batchor/runtime/ingestion.py index d2a2adc..d655e34 100644 --- a/src/batchor/runtime/ingestion.py +++ b/src/batchor/runtime/ingestion.py @@ -5,7 +5,7 @@ import json import time from dataclasses import dataclass, field -from typing import Any, Callable, Iterator, cast +from typing import Any, Callable, Iterator, Protocol, cast from pydantic import BaseModel @@ -24,14 +24,27 @@ def _noop_poll_active_batches( - _run_id: str, - _context: RunContext, + run_id: str, + context: RunContext, *, deadline: float | None = None, ) -> None: + del run_id, context, deadline return +class ActiveBatchPoller(Protocol): + """Callable seam for reconciling active batches during ingestion.""" + + def __call__( + self, + run_id: str, + context: RunContext, + *, + deadline: float | None = None, + ) -> None: ... + + @dataclass(frozen=True) class IngestionDeps: """Dependencies needed by the ingestion and resume layer. @@ -52,7 +65,7 @@ class IngestionDeps: emit_event: Callable[..., None] submit_pending_items: Callable[[str, RunContext], int] configs_match_for_resume: Callable[[PersistedRunConfig, PersistedRunConfig], bool] - poll_active_batches: Callable[..., None] = _noop_poll_active_batches + poll_active_batches: ActiveBatchPoller = _noop_poll_active_batches monotonic: Callable[[], float] = time.monotonic work_slice_max_items: int = 1000 work_slice_max_seconds: float | None = None @@ -122,6 +135,8 @@ def resume_existing_run( return if control_state is RunControlState.RUNNING and deps.state.get_active_batches(run_id=run_id): _poll_active_batches(deps, run_id=run_id, context=context, deadline=deadline) + if _deadline_reached(deps, deadline): + return summary = deps.state.get_run_summary(run_id=run_id) if summary.backoff_remaining_sec > 0: return @@ -156,6 +171,8 @@ def resume_existing_run( checkpoint_payload=checkpoint.checkpoint_payload, ingestion_complete=True, ) + if _deadline_reached(deps, deadline): + return deps.submit_pending_items(run_id, context) return ingest_job_items( @@ -170,6 +187,8 @@ def resume_existing_run( return if control_state is RunControlState.PAUSED: return + if _deadline_reached(deps, deadline): + return deps.submit_pending_items(run_id, context) diff --git a/tests/integration/test_batchor_runner.py b/tests/integration/test_batchor_runner.py index 4e6b40b..748ef1b 100644 --- a/tests/integration/test_batchor_runner.py +++ b/tests/integration/test_batchor_runner.py @@ -1989,11 +1989,15 @@ def flaky_prompt_builder(item: BatchItem[dict[str, str]]) -> PromptParts: checkpoint = storage.get_ingest_checkpoint(run_id=run_id) assert checkpoint is not None - assert checkpoint.next_item_index == 1000 - assert checkpoint.checkpoint_payload == { - "source_index": 1, - "child_checkpoint": 1, - } + expected_source_items = list(build_source()) + assert 0 < checkpoint.next_item_index <= len(first_records) + 1 + assert checkpoint.checkpoint_payload is not None + resumed_source_items = list(build_source().iter_from_checkpoint(checkpoint.checkpoint_payload)) + assert [source_item.item.item_id for source_item in resumed_source_items] == [ + item.item_id for item in expected_source_items[checkpoint.next_item_index :] + ] + persisted_before_resume = storage.get_item_records(run_id=run_id) + assert [record.item_index for record in persisted_before_resume] == list(range(checkpoint.next_item_index)) assert checkpoint.ingestion_complete is False resumed = runner.start( @@ -2017,6 +2021,7 @@ def flaky_prompt_builder(item: BatchItem[dict[str, str]]) -> PromptParts: assert checkpoint.ingestion_complete is True item_records = storage.get_item_records(run_id=run_id) + assert [record.item_index for record in item_records] == list(range(len(expected_source_items))) last_three = item_records[-3:] assert [record.item_index for record in last_three] == [999, 1000, 1001] assert [record.metadata["batchor_lineage"]["source_ref"] for record in last_three] == [ @@ -2029,6 +2034,7 @@ def flaky_prompt_builder(item: BatchItem[dict[str, str]]) -> PromptParts: ] results = resumed.results() + assert [result.item_id for result in results] == [item.item_id for item in expected_source_items] second_namespace = results[-1].item_id.split("__", maxsplit=1)[0] assert [result.item_id for result in results[-3:]] == [ f"{second_namespace}__row0", diff --git a/tests/unit/test_batchor_runtime_ingestion.py b/tests/unit/test_batchor_runtime_ingestion.py index 3b9bc38..95d2627 100644 --- a/tests/unit/test_batchor_runtime_ingestion.py +++ b/tests/unit/test_batchor_runtime_ingestion.py @@ -40,6 +40,14 @@ def __call__(self) -> float: return next(self._readings) +class _MutableClock: + def __init__(self, value: float) -> None: + self.value = value + + def __call__(self) -> float: + return self.value + + class _CompleteCheckpointSource(CheckpointedItemSource[dict[str, str]]): def source_identity(self) -> SourceIdentity: return SourceIdentity( @@ -339,6 +347,151 @@ def poll_and_backoff(run_id: str, context: object) -> None: assert checkpoint.ingestion_complete is False +def test_resume_existing_run_stops_after_active_poll_reaches_deadline() -> None: + """A poll that overruns the deadline cannot begin source work or submission.""" + source = _ExplodingSource() + job = BatchJob( + items=source, + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + clock = _MutableClock(0.0) + + class _DeadlineAdvancingProvider: + def __init__(self) -> None: + self.polled_batches: list[str] = [] + + def get_batch(self, batch_id: str) -> object: + self.polled_batches.append(batch_id) + clock.value = 1.0 + return object() + + provider = _DeadlineAdvancingProvider() + submitted_runs: list[str] = [] + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda run_id, _context: submitted_runs.append(run_id) or 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=lambda _run_id, _context, *, deadline=None: provider.get_batch("provider_1"), + monotonic=clock, + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + run_id = "resume_deadline_poll_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, + next_item_index=0, + checkpoint_payload=0, + ingestion_complete=False, + ), + ) + storage.register_batch( + run_id=run_id, + local_batch_id="local_1", + provider_batch_id="provider_1", + status="submitted", + custom_ids=[], + ) + + resume_existing_run( + deps, + run_id=run_id, + job=job, + config=config, + context=context, + deadline=1.0, + ) + + assert provider.polled_batches == ["provider_1"] + assert submitted_runs == [] + checkpoint = storage.get_ingest_checkpoint(run_id=run_id) + assert checkpoint is not None + assert checkpoint.next_item_index == 0 + assert checkpoint.ingestion_complete is False + + +def test_resume_existing_run_skips_direct_submit_after_deadline_expiring_poll() -> None: + """The checkpoint-complete resume path observes a poll-overrun deadline.""" + source = _CompleteCheckpointSource() + job = BatchJob( + items=source, + build_prompt=lambda item: PromptParts(prompt=item.payload["text"]), + provider_config=OpenAIProviderConfig(api_key="k", model="gpt-4.1"), + ) + config = build_persisted_config(job) + storage = MemoryStateStore() + clock = _MutableClock(0.0) + + class _DeadlineAdvancingProvider: + def get_batch(self, batch_id: str) -> object: + assert batch_id == "provider_1" + clock.value = 1.0 + return object() + + provider = _DeadlineAdvancingProvider() + submitted_runs: list[str] = [] + deps = IngestionDeps( + state=storage, + emit_event=lambda *args, **kwargs: None, + submit_pending_items=lambda run_id, _context: submitted_runs.append(run_id) or 0, + configs_match_for_resume=lambda stored, supplied: stored == supplied, + poll_active_batches=lambda _run_id, _context, *, deadline=None: provider.get_batch("provider_1"), + monotonic=clock, + ) + context = build_run_context( + config=config, + output_model=None, + create_provider=lambda _cfg: _NoopProvider(), + ) + run_id = "resume_deadline_direct_submit_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, + next_item_index=0, + checkpoint_payload={"done": True}, + ingestion_complete=False, + ), + ) + storage.register_batch( + run_id=run_id, + local_batch_id="local_1", + provider_batch_id="provider_1", + status="submitted", + custom_ids=[], + ) + + resume_existing_run( + deps, + run_id=run_id, + job=job, + config=config, + context=context, + deadline=1.0, + ) + + assert submitted_runs == [] + checkpoint = storage.get_ingest_checkpoint(run_id=run_id) + assert checkpoint is not 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)],