Scheduler¶

The scheduler (oasr/engine/scheduler.py) is the admission-control and batching policy module of the ASR engine. Each engine step() calls Scheduler.schedule() exactly once; the returned SchedulerOutput specifies which offline requests should be forwarded as a batch this step, which streaming requests should be promoted from waiting to running, and which streaming requests are currently active.

1. Purpose and Responsibilities¶

The scheduler:

Holds two waiting queues — one for streaming, one for offline — and one running pool of admitted streaming requests.
Decides per step which requests to admit, applying length-aware bucketing for offline batches and a cohort-admission policy for streaming.
Bounds starvation via max_wait_time: any request that has waited too long becomes a forced-flush anchor regardless of policy.
Manages priorities so urgent requests bypass FIFO ordering.
Provides O(1) request lookup so the engine's feed_chunk does not pay an O(N) scan per audio chunk on backlog workloads.

The scheduler is purely a policy and bookkeeping object — it never allocates GPU memory, never runs feature extraction, and never touches the model.

2. High-Level Architecture¶

                add_request(req)
                       │
        ┌──────────────┴───────────────┐
   streaming?                       offline?
        ▼                              ▼
  _streaming_waiting (deque)      _offline_waiting (deque)
        │                              │
        │              schedule()      │
        │       ┌──────────┴──────┐    │
        │       │                 │    │
        ▼       ▼                 ▼    ▼
   admit()  build_offline_batch()      ─── forced flush gate
        │       │                 │
        ▼       ▼
  _running   offline_batch
  (OrderedDict<id, Request>)    (List[Request])

  schedule() → SchedulerOutput {running_streams, newly_admitted, offline_batch}

_streaming_waiting, _offline_waiting, and _running are kept consistent via _index: Dict[str, Request] so find_request(rid) and abort_request(rid) are O(1) regardless of which queue the request is in.

3. Internal Structure¶

Data structures¶

Field	Type	Role
`_streaming_waiting`	`Deque[Request]`	Streaming requests not yet admitted. Priority-aware insertion; FIFO within the same priority (§4.4).
`_offline_waiting`	`Deque[Request]`	Offline requests awaiting a batched forward pass.
`_running`	`OrderedDict[str, Request]`	Admitted streaming requests with allocated KV cache. Iteration order = admission order.
`_index`	`Dict[str, Request]`	O(1) lookup across all three queues. Kept in sync on every add / finish / abort.
`_next_stream_id`	`int`	Monotonic stream ID counter assigned during streaming admission.

`SchedulerOutput`¶

@dataclass
class SchedulerOutput:
    running_streams: List[Request]   # all currently-active streaming requests
    newly_admitted:  List[Request]   # promoted to RUNNING this step
    offline_batch:   List[Request]   # one length-bucketed offline batch

The engine's step loop reads:

offline_batch → forwarded through OfflineExecutor.run and finalised.
newly_admitted (streaming subset) → ModelRunner.allocate_stream.
running_streams → drives the per-step fbank-extract / forward / decode passes; the engine, not the scheduler, decides per request whether to run feature extraction, an encoder chunk, or finalisation.

4. Core Algorithms¶

4.1 Offline batching: `_build_offline_batch`¶

q = _offline_waiting
if q empty: return []
cap = max(1, max_batch_size)
policy = config.schedule_policy

if now >= promote_at:                        # §4.5: only when the bound has work
    promote_at = promote_overdue(q)          # overdue requests to the head
force_flush = q[0].waited_for >= max_wait_time

if policy == "fcfs":
    return [popleft() up to cap]

if policy == "window":                       # the default; see §4.1.1
    window = up to cap * length_window_factor requests from the head:
             the overdue ones (any priority), then the best class waiting
    sort window by cost; among the contiguous runs of cap that contain
    q[0] and pass the guards below (and max_batch_frames), take the one
    with the least padding (max_len * cap - sum_len); if none passes,
    cap -= 1 and look again (the anchor alone always ships)
    return that run, in arrival order

if policy == "sjf" and not force_flush:
    sort q by (priority, num_frames)        # ascending length

# "bucket" and "sjf" share length-aware selection:
anchor = q.popleft()                         # oldest (or shortest if SJF)
batch  = [anchor]
min_len, max_len = anchor.num_frames, anchor.num_frames

if force_flush:
    return _fill_batch_fifo(batch, q, cap)   # strict FIFO refill

for cand in q:
    new_min = min(min_len, cand.num_frames)
    new_max = max(max_len, cand.num_frames)

    # bucket guard: candidates within length_bucket_ratio of anchor
    if length_bucket_ratio > 0 and new_min/new_max < length_bucket_ratio:
        continue
    # pad-waste guard: reject if padded/useful would exceed max_offline_pad_ratio
    useful = sum(r.num_frames for r in batch) + cand.num_frames
    padded = new_max * (len(batch) + 1)
    if max_offline_pad_ratio > 0 and padded/useful > max_offline_pad_ratio:
        continue

    batch.append(cand)
    min_len, max_len = new_min, new_max
    if len(batch) == cap: break
return batch

Two independent guards bound padded-compute waste, in every length-aware policy ("window", "bucket", "sjf"):

length_bucket_ratio — soft floor on min_len / max_len within a batch. Disabled (0) by default because splitting one bursty transcribe(list_of_N) call into sub-batches multiplies CPU feature-extraction cost while saving only a few percent of GPU compute on real workloads with moderate length spread.
max_offline_pad_ratio — hard cap on (max_len × B) / sum_len. Default 4.0 admits LJSpeech-scale spreads in one batch but rejects pathological mixes of ~1 s and ~30 s clips.

4.1.1 Why `"window"` is the default¶

A padded batch costs its longest row times its width, and once the offline forward is graph-captured that padding is GPU time, not a host-side detail. Arrival-order batching therefore pays for the length spread of whatever arrived together: on LJSpeech (1.5–10 s) FIFO batches of 32 compute ~1.48× the useful frames. Sorting the whole queue ("sjf") removes nearly all of it, but under continuous arrivals a long request loses to every newer short one until max_wait_time forces a FIFO flush — which brings the padding back.

"window" bounds the reordering instead. It sorts only the oldest max_batch_size × length_window_factor requests and always ships the oldest, choosing the least-padded length-run of max_batch_size that contains it:

no starvation by construction — the anchor always ships, so the policy needs no flush to guarantee progress, and a request can be passed over only while it is still inside the window;
a no-op below capacity — with fewer than max_batch_size requests waiting there is nothing to choose, so light load sees FIFO exactly;
one priority class at a time — the window is the best class waiting, plus any request that has waited out max_wait_time, whatever its class.

With the default factor of 4 the padding falls to ~1.11× (simulated on 2048 LJSpeech arrivals; no request delayed more than 8 batches). The accuracy gate's five offline entries are error-for-error identical under it. max_batch_frames, max_offline_pad_ratio and length_bucket_ratio bound its runs as they bound the other policies' batches: a run that breaks one is not a candidate, and the batch narrows until one fits. At their defaults none of them changes a batch on the LJSpeech benchmark workloads.

4.2 Streaming admission: `_admit_streaming`¶

budget = max_batch_size - len(_running)
if budget <= 0: return []

# Length bucketing only — no offset gate (heterogeneous-offset batching).
if streaming_cohort_admit or policy == "sjf":
    # overdue streams first, oldest first, whatever their priority;
    # the rest by priority, then length
    sort _streaming_waiting by
        (0, arrival_time)            if waited_for >= max_wait_time
        (1, priority, num_frames)    otherwise
else:
    promote_overdue(_streaming_waiting)      # §4.5

admitted = []
while _streaming_waiting and len(admitted) < budget:
    req = _streaming_waiting.popleft()
    req.stream_id = _next_stream_id; _next_stream_id += 1
    req.state = RUNNING
    admitted.append(req)
return admitted

4.3 Heterogeneous-offset batching¶

_forward_batched_paged (in oasr/engine/model_runner.py) uses PyTorch FlexAttention over the paged KV pool. FlexAttention accepts a block_mask built from each stream's cache_seqlens and a score_mod that adds the Conformer rel-pos bias matrix_bd / sqrt(d_k). Because both per-stream lengths and per-stream pos_emb are expressed inside the fused kernel, streams with different (offset, cache_seqlens) batch into a single forward — there is no cohort-admit gate.

Practical consequences:

New requests join an in-flight pool at any time. The "freed slot is filled immediately" behaviour — previously available only with streaming_cohort_admit=False — is now the default for all policies.
streaming_cohort_admit=True no longer gates admission on offset == 0; it only sorts the waiting queue by length so each batched forward sees length-similar utterances (less padded-compute waste from a single very-long stream stretching the chunk-attention context for the rest of the cohort).
That sort is bounded by max_wait_time, like every offline policy. A waiting stream's num_frames is the audio it has buffered so far, so it grows for as long as the stream waits. Sorted by length alone, the stream that had waited longest came last: once more streams waited than there were free slots, every newer arrival overtook it on every tick and its backlog grew for as long as the load lasted.
The encoder's positional embedding is built per-stream when offsets differ (RelPositionalEncoding.position_encoding(offsets_tensor, …)), collapsing to broadcast when they happen to match.

4.4 Priority handling¶

Lower Request.priority is sooner. 0 is the default and the most urgent class: Request clamps a priority to at least 0, so a caller can defer its own work but cannot move it ahead of everyone else's. The front-ends pass the wire's priority (an i32) straight through, so the clamp is what stops one client with a negative value from overtaking all the others.

Both queues use _insert_ordered on add_request, which files a request behind every request of the same or a better class:

When every request has the same priority (the default) this is an append (O(1)).
Otherwise it scans from the back for the last request whose priority is not worse, O(N). From the back, because the head may hold overdue requests (§4.5) in arrival order rather than priority order; a scan from the front could stop at a deferred one there and file the request ahead of better ones behind it.

Priority is the leading key of the length sorts: SJF orders by (priority, num_frames), and cohort admission sorts the streams that are not overdue by (priority, num_frames).

4.5 Forced flush (starvation bound)¶

promote_overdue(q)        # every request with waited_for >= max_wait_time,
                          # oldest first, moves to the head
force_flush = q[0].waited_for >= cfg.max_wait_time

The bound outranks priority. The queue is otherwise in priority order, so its head is the best class's oldest request, and a flush read off the head alone never fired for a deferred request while better work kept arriving: it waited for as long as the better load lasted. promote_overdue runs before every offline selection that can need it and before every unsorted streaming admission (cohort admission's sort puts overdue streams first itself), so the head is the request that has waited longest whenever any is overdue. A consequence: under sustained overload every waiting request is soon overdue, and service falls back to arrival order across classes. Priority orders work within max_wait_time; it cannot defer a request past it.

An offline selection skips the bound until it can have work, since re-reading every waiting request's arrival time was most of a selection's cost on a deep queue. promote_overdue returns the instant the next waiting request turns overdue, and add_request brings that forward to the new request's own deadline, or to now when the request was filed ahead of others (it may be ahead of an overdue one). Nothing else can put an overdue request out of place: removing a request never does, and while the head is overdue every policy leaves what it does not take in order (the BatchingPolicy contract; sjf sorts only when nothing is overdue).

When the oldest offline request has been waiting too long:

Skip the SJF re-sort (it would push this request back).
Skip bucket and pad-waste guards — fill the batch via FIFO from the current head.

This guarantees that under sustained heavy load no request waits longer than ~max_wait_time + step_duration. The default max_wait_time=0.2s keeps p99 latency bounded while still letting normal-load batches collect ideal length-similar peers.

4.6 Preferred batch sizes (Triton-style)¶

preferred_batch_size: Optional[List[int]] constrains both queues to admit batches only at the listed B values. The smallest preferred value is the minimum batch size the scheduler will emit; the largest is the effective ceiling (must be <= max_batch_size, which remains the hard cap on KV / slot-pool resources).

# Streaming admission:
n_avail   = min(budget, len(_streaming_waiting))
candidate = len(_running) + n_avail
target    = largest preferred <= candidate          # 0 if candidate < min preferred
n_admit   = max(0, target - len(_running))
if n_admit == 0:
    if oldest waiting >= max_wait_time: n_admit = n_avail   # force-flush
    else: n_admit = 0

# Offline batch (applied after _build_offline_batch finishes):
target = largest preferred <= len(batch)
if target == 0 and not force_flush: requeue everything, return []
if target < len(batch):           requeue overflow to head of queue

max_wait_time is the escape valve: when no preferred grouping is reachable but the oldest request has waited too long, the scheduler ships whatever is available even if the resulting size is not preferred. Force-flushed batches bypass the preferred-size trim too — the wait deadline outranks the size constraint.

When set, preferred_batch_size also:

Drives the encoder CUDA-Graph pre-warm in ASREngine.__init__ (ModelRunner.prewarm_encoder_graphs(...)) so the first real chunk at each preferred B replays a captured graph instead of paying capture latency on the request path.
Defaults feature_graph_batch_buckets to the same list, so the encoder graph cache and the feature graph cache share one B ladder.
Snaps Scheduler.split_offline_batch to the preferred sizes — each micro-batch is the largest preferred value <= remaining, falling back to the remainder for the (rare) trailing odd chunk.

5. Data Flow¶

                add_request                 step()
                     │                        │
                     ▼                        ▼
              ┌──────────────┐         ┌─────────────┐
   streaming → _streaming_w │         │ schedule()  │
   offline   → _offline_w   │ ──────▶ │             │
              │ _running    │         │             │
              │ _index      │         └──────┬──────┘
              └──────────────┘                │
                                              ▼
                                      SchedulerOutput
                                              │
                                              ▼
                                  ┌─────────────────────┐
                                  │  Engine step loop   │
                                  │  (allocate / fbank / │
                                  │   forward / decode) │
                                  └──────────┬──────────┘
                                             │
                                             ▼
                              finish_request | abort_request
                                             │
                                             ▼
                                   _running.pop, _index.pop

6. Configuration Options¶

All scheduler-relevant knobs live on EngineConfig:

Option	Default	Effect
`max_batch_size`	32	Encoder forward batch size — caps the running streaming pool and is the GPU forward width of the offline batch. Offline admission per step is capped at `max_batch_size` (one length-bucketed batch per `step()`). Remains the hard cap on resources when `preferred_batch_size` is set.
`preferred_batch_size`	`None`	Triton-style preferred B values. When set, streaming admission and offline batch construction snap to one of these sizes; `max_wait_time` is the escape valve. Also drives encoder CUDA-Graph pre-warm and defaults `feature_graph_batch_buckets`. All values must be `<= max_batch_size`. `None` keeps the legacy "admit greedily" behaviour.
`length_bucket_ratio`	0.0	Soft floor on `min_len/max_len` inside a bucket. `0` disables.
`max_offline_pad_ratio`	4.0	Hard cap on `(max_len × B) / sum_len`. `0` disables.
`max_wait_time`	0.2 s	Starvation bound: oldest offline request triggers forced flush, and a streaming request waiting this long is admitted ahead of the length-sorted rest.
`schedule_policy`	`"window"`	One of `"fcfs"`, `"window"`, `"bucket"`, `"sjf"`.
`length_window_factor`	4	`"window"`: batches are chosen from the oldest `max_batch_size × factor` requests. Larger removes more padding and reorders more.
`streaming_cohort_admit`	`True`	Length-bucket the waiting queue before admission so each batched paged forward sees length-similar utterances; streams past `max_wait_time` go first, oldest first. No longer gates on offset alignment (FlexAttention handles per-stream `cache_seqlens` natively).

Choosing a policy¶

Workload	Recommended policy
Ordering matters (live transcription, real-time pipelines)	`fcfs`
Mixed length, throughput-sensitive (default)	`window`
Anchor plus first-fit peers from the whole queue (the previous default)	`bucket`
Heavy backlog of mixed lengths, latency-tolerant for outliers	`sjf`

sjf can starve long requests indefinitely without max_wait_time — keep it ≤ a few hundred ms unless you explicitly want long-tail latency.

7. Usage Examples¶

The scheduler is owned by ASREngine; users do not construct it directly. Direct interaction is mostly useful in tests or in a custom engine.

7.1 Inside the engine step loop¶

sched: SchedulerOutput = self._scheduler.schedule()

for req in sched.newly_admitted:
    if req.streaming:
        self._model_runner.allocate_stream(req)

if sched.offline_batch:
    outputs.extend(self._run_offline_batch(sched.offline_batch))

for req in sched.running_streams:
    # extract, forward, decode, finalise...

7.2 Aborting a request¶

req = engine._scheduler.abort_request(rid)
if req is not None and req.stream_context is not None:
    engine._model_runner.free_stream(req)

7.3 Looking up a request by id (e.g. for `feed_chunk`)¶

req = self._scheduler.find_request(request_id)   # O(1)
if req is None:
    raise KeyError(f"unknown or finished request_id {request_id!r}")
self._input_processor.append_streaming_chunk(req, chunk, is_last=is_last)

7.4 Custom direct usage (test scenarios)¶

sched = Scheduler(EngineConfig(max_batch_size=4, schedule_policy="sjf"))
for r in requests:
    sched.add_request(r)

while sched.has_pending():
    out = sched.schedule()
    # ... drive forward, then mark finished:
    for req in finished:
        sched.finish_request(req.request_id)

8. Error Handling and Edge Cases¶

Condition	Behaviour
`finish_request(rid)` for an id not in `_running`	`KeyError` from `OrderedDict.pop`. The engine ensures finishes only target running streams.
`abort_request(rid)` for an unknown id	Returns `None` — caller decides whether to raise.
Empty waiting queues	`schedule()` returns an empty `SchedulerOutput` (no offline batch, no admissions). `running_streams` may still be non-empty.
`max_batch_size == len(_running)`	No streaming admission this step (`budget <= 0`).
Cohort-admit gate blocking	Empty `newly_admitted` list with non-zero waiting queue — recover when running pool drains.
`length_bucket_ratio=0` and `max_offline_pad_ratio=0`	Effectively disables bucketing; behaves like FCFS within each step but still respects cap.
Forced flush fires	All bucket guards skipped; the resulting batch may be highly padded.
Priority insert into already long queue	O(N) walk from the back, but only when priorities differ.
Negative `priority`	Clamped to `0` by `Request`.

9. Performance Considerations¶

add_request is O(1) when the new request's priority is no better than the queue's last, which includes the all-default case; O(N) otherwise.
The offline wait bound is O(N), and runs rarely. When it runs it reads every waiting request's arrival time, sorts those times if any request is overdue (linear on a queue already in arrival order), and rebuilds the queue only when the overdue requests are not already at the head. It runs only when a request has turned overdue or was filed ahead of others (§4.5): once per burst of one class, not once per selection. Unsorted streaming admission runs it every time; SJF and cohort admission add one O(N log N) sort.
find_request / abort_request are O(1) thanks to _index.
No GPU work. The scheduler is pure Python on dataclasses; it is rarely the bottleneck. Profile with NVTX (engine.step → schedule) to confirm.
Cohort admission trades latency for throughput. On backlog workloads it is a large throughput win, because _forward_batched_paged can take the full B path on every step. On low-concurrency or interactive workloads it adds idle time at cohort boundaries — measure both on your traffic before deciding.
Offline admission is one max_batch_size batch per step. The scheduler admits a single length-bucketed batch of up to max_batch_size requests, which OfflineExecutor runs as one offline forward (GPU fbank → forward → decode). For larger batches, raise max_batch_size.

10. Extension Points¶

Custom policy. Add a new schedule_policy value and branch in _build_offline_batch / _admit_streaming. Keep the force_flush invariant so starvation stays bounded.
Custom priorities. Replace _insert_ordered with a heap if you need many priority levels and large queues. Keep the _index invariant — feed_chunk depends on it.
Multi-tenant fairness. Add a per-tenant Request field and group the offline batch by tenant before length bucketing. Cohort admission could be scoped per tenant similarly.
Preemption. The current scheduler has no preempt path: if the block pool runs out, the engine fails. A real preemption policy would swap a low-priority running stream out of _running, free its blocks, and reinsert it into _streaming_waiting — implementable here without touching the engine if you also expose a hook for releasing the cache.

11. Quick Reference¶

add_request(req)        # route to streaming/offline waiting queue
schedule() → output     # one-step decision
finish_request(rid)     # running → finished, drop from index
abort_request(rid)      # remove from any queue, return req or None
find_request(rid)       # O(1) lookup across all queues
has_pending()           # any waiting or running? (engine loop guard)
num_waiting / num_waiting_streaming / num_waiting_offline / num_running