Skip to content

[FIX] Batch-selection protection leaks across batches and the Scheduler boundary #67

Description

@kunkunblueberry

Problem

Scheduler keeps one global selected_request_ids: set[str] describing requests that preemption must not suspend. Its lifetime is implicit and does not match the batch it describes:

  • schedule() clears it at entry; nothing else does.
  • _take() adds to it but never resets it.
  • finish() never removes the ID of a request that has ended.
  • PreemptionManager._is_preemption_candidate() reads e.scheduler.selected_request_ids directly, so the Scheduler's selection state is also a cross-module contract.

Concrete trigger

pd/worker.py::prefill_step() calls engine.scheduler._take(Stage.PREFILL) directly, without going through schedule(). Every such call appends up to max_num_seqs IDs that are never removed, so the set grows for the process lifetime and every previously scheduled request stays ineligible as a preemption victim.

tests/test_prefix_growth.py already has to work around this by clearing scheduler internals by hand:

e.scheduler.selected_request_ids.clear()
assert e.preemption.preempt(e.scheduler.requests["b"])

Observed behavior

preempt() reports no safe victim although one exists, because the only candidate's ID is still in the stale selection set:

e = LLMEngine(model, scheduler_config=SchedulerConfig(max_num_seqs=1, enable_preemption=True))
e.add_request("victim", [1, 2], SamplingParams(max_tokens=5, ignore_eos=True))
e.step()  # "victim" is selected by this batch
e.add_request("requester", [3, 4], SamplingParams(max_tokens=5, ignore_eos=True))
e.preemption.preempt(e.scheduler.requests["requester"])  # returns False

The same happens after abort_request() plus request-ID reuse: the new request inherits the previous request's protection, even though the previous request is no longer registered.

Expected behavior

The exclusion describes only the batch under construction. It must not outlive that batch, reach a direct _take() caller, or survive the request whose ID it names.

Scope

Scheduler-owned selection state and the preemption callback contract. This is not completion of M3: Scheduler/KV private access, engine-owned result progression, and PD private scheduling calls remain future work.

Refs #32.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions