vllm-project / vllm-project/aibrix

Batch: resume a truncated job from partial results instead of rerunning the whole input

Open
#2,535 1 comment 0 reactions 1 assignee Claimed by @scarlet25151 View on GitHub
area/batch kind/enhancement kind/feature
Dominant language
Go
Stars
5.1k
Forks
697
Avg merge
1d 19h
Merged PRs (30d)
104

Description

## Problem

A batch job's lifecycle is bound to a single resource allocation. On constrained or preemptible capacity that allocation can end well before the completion window; the job then lands in a terminal state (`expired` / `resource_failed`) and **all completed work is discarded** — the only recovery is resubmitting the entire input file.

For multi-hour or multi-day workloads this is the dominant cost, and it compounds: shorter leases mean higher truncation probability, which pushes you toward smaller input files, which means more allocation rounds. There is no way out of the loop without resume at request granularity.

## Current state

Durable per-request done-markers already exist, but they can't be used across attempts:

- `storage/adapter.py:495` — key is `batch:{job.job_id}:done/{idx}`, value `output:{etag}` / `error:{etag}`
- `job_driver/base.py:475` — `_get_next_pass_start` scans them to find a resume point

Two limits:

1. **Keyed by `job_id`** — a resubmitted job gets a new id and cannot see the previous attempt's markers. This only rescues a worker restart *within* the same job.
2. **Keyed by line index** — incompatible with resubmitting a subset of the input, since indices shift. `custom_id` is the only identity stable across attempts.

`JobProgressTracker` (`state/job_progress_tracker.py:44-49`) is in-memory only; just the aggregate `request_counts` is persisted, so we know *how many* finished, not *which*.

## Proposal

Rather than extend the marker store, treat the artifacts themselves as the ledger, keyed on `custom_id`:

```
delta = input − (output ∪ non-retryable errors)
```

Resubmit `delta` as a new job; the final result is the union of all rounds' outputs, aligned against the original input. The delta shrinks geometrically (usually a few rounds), and shard size adapts to whatever lease was actually granted instead of being guessed up front.

This needs two things from the platform:

**1. Partial finalize on truncation.** `output_dataset` / `error_dataset` are only assembled in `finalize_job` (`base.py:1555`). `_finish_stopped_job` (`base.py:592`) covers a graceful stop, but a hard teardown at lease expiry can skip it — the per-request results and etags survive in storage yet are unreachable to the client. Guaranteeing a partial flush on truncation is the single blocking requirement.

**2. Distinguishable failure causes in the error dataset.** Requests killed by teardown surface as 5xx/connection errors next to genuine 4xx/validation failures. A client computing `input − (output ∪ error)` would silently drop the former. Error records already carry `BatchJobErrorCode` and `status_code`; we should document and guarantee which codes are retryable so clients can classify reliably.

Known and acceptable: requests in flight at truncation may be re-executed, bounded by the client concurrency setting.

## Follow-ups (not in scope here)

- Elastic `resource_request` (min/max replicas) so a partial allocation still starts, instead of all-or-nothing
- A logical workload resource above `Job` owning dataset split, aggregate progress and deadline — today `Job` conflates "the user's job" with "one physical execution attempt"

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.