[Diffusion] Unify diffusion request identity on request_id (#3744)
Signed-off-by: jader <yjader@foxmail.com>
This commit is contained in:
1 parent
beb00fce31
commit
6fa9dbbe42
34 files changed
+421
-461
No files matched your search
@@ -169,15 +169,13 @@ class DiffusionRequestStatus(enum.IntEnum):
|
||||
|
||||
@dataclass
|
||||
class DiffusionRequestState:
|
||||
sched_req_id: str
|
||||
request_id: str
|
||||
req: OmniDiffusionRequest
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
```
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
|
||||
|
||||
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
|
||||
|
||||
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
|
||||
@@ -188,7 +186,6 @@ class DiffusionRequestState:
|
||||
class _BaseScheduler(SchedulerInterface):
|
||||
def __init__(self) -> None:
|
||||
self._request_states = {}
|
||||
self._request_id_to_sched_req_id = {}
|
||||
self._waiting = deque()
|
||||
self._running = []
|
||||
self._finished_req_ids = set()
|
||||
@@ -199,7 +196,8 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
|
||||
duplicated in each policy.
|
||||
|
||||
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
|
||||
|
||||
@@ -219,12 +217,12 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
|
||||
|
||||
#### 2.5 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
sched_req_id = scheduler.add_request(request)
|
||||
request_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
|
||||
@@ -127,7 +127,12 @@ def run_one_path(
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-sample-{index}",
|
||||
)
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
results.append(
|
||||
{
|
||||
@@ -150,6 +155,7 @@ def run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs: dict[str, torch.Tensor],
|
||||
noise: torch.Tensor,
|
||||
request_id: str,
|
||||
) -> torch.Tensor:
|
||||
output = pipeline.forward(
|
||||
OmniDiffusionRequest(
|
||||
@@ -161,6 +167,7 @@ def run_pipeline_forward(
|
||||
"decode_image": False,
|
||||
}
|
||||
),
|
||||
request_id=request_id,
|
||||
)
|
||||
)
|
||||
if output.error:
|
||||
@@ -212,23 +219,38 @@ def benchmark_forward(
|
||||
|
||||
_synchronize(args.device)
|
||||
cold_start_begin = time.perf_counter()
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-cold",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
|
||||
|
||||
warmup_ms: list[float] = []
|
||||
for _ in range(args.warmup_iters):
|
||||
for iter_idx in range(args.warmup_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
warmup_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
benchmark_ms: list[float] = []
|
||||
for _ in range(args.benchmark_iters):
|
||||
for iter_idx in range(args.benchmark_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
|
||||
)
|
||||
_synchronize(args.device)
|
||||
benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
@@ -295,7 +317,12 @@ def main() -> None:
|
||||
dataset=dataset,
|
||||
train_meta=train_meta,
|
||||
collate_samples=collate_open_loop_samples,
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
|
||||
policy,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id="internvla-a1-open-loop",
|
||||
),
|
||||
num_episodes=args.num_episodes,
|
||||
seed=args.seed,
|
||||
device=args.device,
|
||||
|
||||
@@ -60,6 +60,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
|
||||
output_type="pil",
|
||||
generator_device="cpu",
|
||||
),
|
||||
"request_id": "diffusers-backend",
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return OmniDiffusionRequest(**defaults)
|
||||
|
||||
@@ -48,6 +48,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=prompts or [{"prompt": "a cat dancing"}],
|
||||
sampling_params=sp,
|
||||
request_id="dmd2-sanitize",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -48,7 +48,7 @@ def _make_pipeline(cls):
|
||||
|
||||
def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
|
||||
sp = OmniDiffusionSamplingParams(**sp_kwargs)
|
||||
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
|
||||
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
|
||||
@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
|
||||
|
||||
|
||||
def test_should_apply_pe_disables_dummy_warmup_request():
|
||||
req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
|
||||
req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
|
||||
|
||||
assert ErnieImagePipeline._should_apply_pe(req) is False
|
||||
|
||||
|
||||
@@ -167,6 +167,7 @@ def test_basic_generation(ovis_pipeline):
|
||||
# Setup request
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["A photo of a cat"],
|
||||
request_id="ovis-basic",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=256,
|
||||
@@ -206,6 +207,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
|
||||
"negative_prompt": "bad quality",
|
||||
}
|
||||
],
|
||||
request_id="ovis-guidance",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=256,
|
||||
@@ -225,6 +227,7 @@ def test_resolution_check(ovis_pipeline):
|
||||
# Pass odd resolution
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test"],
|
||||
request_id="ovis-resolution",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=250, # Not divisible by 16 (8*2)
|
||||
width=250,
|
||||
|
||||
@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
|
||||
num_waiting_reqs: int = 0
|
||||
|
||||
@property
|
||||
def scheduled_req_ids(self):
|
||||
def scheduled_request_ids(self):
|
||||
ids = [req.request_id for req in self.scheduled_new_reqs]
|
||||
if self.scheduled_cached_reqs:
|
||||
ids.extend(self.scheduled_cached_reqs.sched_req_ids)
|
||||
ids.extend(self.scheduled_cached_reqs.request_ids)
|
||||
return ids
|
||||
|
||||
@property
|
||||
def is_empty(self):
|
||||
return len(self.scheduled_req_ids) == 0
|
||||
return len(self.scheduled_request_ids) == 0
|
||||
|
||||
|
||||
class MockScheduler:
|
||||
@@ -71,31 +71,36 @@ async def test_async_add_req_and_wait_for_response():
|
||||
engine.scheduler = MockScheduler()
|
||||
engine._out_queue = {}
|
||||
engine.abort_queue: queue.Queue[str] = queue.Queue()
|
||||
engine._rpc_lock = threading.Lock()
|
||||
engine._rpc_queue = queue.Queue()
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._init_lock = asyncio.Lock()
|
||||
engine._closed = False
|
||||
engine._loop_started = False
|
||||
engine.main_loop = None
|
||||
|
||||
engine._finalize_finished_request = lambda rid, out, err: out.result
|
||||
|
||||
def mock_execute_batch(sched_output):
|
||||
req_ids = sched_output.scheduled_req_ids
|
||||
request_ids = sched_output.scheduled_request_ids
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
class MockRunnerOutput:
|
||||
def __init__(self, ids):
|
||||
self.req_id = ids
|
||||
self.request_id = ids
|
||||
self.step_index = 0
|
||||
self.finished = True
|
||||
self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
|
||||
|
||||
def get_req_output(self, rid):
|
||||
def get_request_output(self, rid):
|
||||
return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
|
||||
|
||||
return MockRunnerOutput(req_ids)
|
||||
return MockRunnerOutput(request_ids)
|
||||
|
||||
engine.execute_fn = mock_execute_batch
|
||||
|
||||
engine.stop_event = threading.Event()
|
||||
engine.start_background_loop()
|
||||
await engine._check_and_start_background_loop()
|
||||
|
||||
async def run_task(rid):
|
||||
req = SimpleNamespace(request_id=rid)
|
||||
@@ -105,10 +110,13 @@ async def test_async_add_req_and_wait_for_response():
|
||||
|
||||
task_ids = [f"req_{i}" for i in range(5)]
|
||||
tasks = [run_task(rid) for rid in task_ids]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
engine.stop_event.set()
|
||||
engine.worker_thread.join(timeout=5)
|
||||
try:
|
||||
results = await asyncio.gather(*tasks)
|
||||
finally:
|
||||
with engine._cv:
|
||||
engine.stop_event.set()
|
||||
engine._cv.notify_all()
|
||||
engine.worker_thread.join(timeout=5)
|
||||
assert len(results) == 5
|
||||
for rid, res, elapsed in results:
|
||||
assert rid in res.result_data
|
||||
|
||||
@@ -17,15 +17,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _make_request(
|
||||
request_id: str,
|
||||
*,
|
||||
request_ids: list[str] | None = None,
|
||||
) -> OmniDiffusionRequest:
|
||||
def _make_request(request_id: str) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt_{request_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=request_ids or [request_id],
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
@@ -79,15 +74,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
|
||||
engine._finalize_finished_request.assert_not_called()
|
||||
|
||||
|
||||
def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
|
||||
def test_abort_request_id_aborts_scheduler_request() -> None:
|
||||
engine = _make_engine()
|
||||
request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
|
||||
sched_req_id = engine.scheduler.add_request(request)
|
||||
request = _make_request("batch-parent")
|
||||
request_id = engine.scheduler.add_request(request)
|
||||
|
||||
engine.abort("batch-parent")
|
||||
engine._process_aborts_queue()
|
||||
|
||||
state = engine.scheduler.get_request_state(sched_req_id)
|
||||
state = engine.scheduler.get_request_state(request_id)
|
||||
assert state is not None
|
||||
assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
|
||||
|
||||
|
||||
@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
if method in {"execute_model", "execute_stepwise", "generate"}:
|
||||
# args[0] is the request-like object for execute_model.
|
||||
req = args[0] if args else None
|
||||
tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
|
||||
tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
|
||||
return DiffusionOutput(error=f"result_for_{tag}")
|
||||
tag = args[0] if args else method
|
||||
return DiffusionOutput(error=f"rpc_result_for_{tag}")
|
||||
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
return RunnerOutput(
|
||||
req_id=new_req.sched_req_id,
|
||||
request_id=new_req.request_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=result,
|
||||
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
|
||||
def _make_request(tag: str):
|
||||
return SimpleNamespace(
|
||||
request_ids=[tag],
|
||||
request_id=tag,
|
||||
prompts=[f"prompt_{tag}"],
|
||||
sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
|
||||
)
|
||||
@@ -148,6 +148,7 @@ def _make_engine_with_loop(
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._out_queue = {}
|
||||
engine._closed = False
|
||||
engine.abort_queue = queue.Queue()
|
||||
engine._rpc_queue = queue.Queue()
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ import random
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -15,9 +15,53 @@ def _make_request() -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id="request-test",
|
||||
)
|
||||
|
||||
|
||||
def test_request_id_is_required():
|
||||
with pytest.raises(TypeError, match="request_id"):
|
||||
OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
)
|
||||
|
||||
|
||||
def test_request_ids_identity_list_is_removed():
|
||||
req = _make_request()
|
||||
|
||||
assert req.request_id == "request-test"
|
||||
assert not hasattr(req, "request_ids")
|
||||
|
||||
|
||||
def test_request_id_must_be_non_empty():
|
||||
with pytest.raises(ValueError, match="request_id must be a non-empty string"):
|
||||
OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id="",
|
||||
)
|
||||
|
||||
|
||||
def test_dummy_run_request_is_identified_by_reserved_request_id():
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "dummy run"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id=DUMMY_DIFFUSION_REQUEST_ID,
|
||||
)
|
||||
|
||||
assert req.is_dummy_run()
|
||||
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
|
||||
|
||||
|
||||
def test_non_dummy_request_is_not_identified_as_dummy_run():
|
||||
req = _make_request()
|
||||
|
||||
assert not req.is_dummy_run()
|
||||
assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
|
||||
assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
|
||||
|
||||
|
||||
def test_tp_seed_same_across_ranks_and_varies_across_requests():
|
||||
random.seed(0)
|
||||
n_requests = 5
|
||||
|
||||
@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt_{req_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=[req_id],
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
|
||||
def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
|
||||
return RunnerOutput(
|
||||
req_id=req_id,
|
||||
request_id=req_id,
|
||||
step_index=None,
|
||||
finished=finished,
|
||||
result=DiffusionOutput(output=None, error=error),
|
||||
@@ -52,7 +52,7 @@ def _make_step_output(
|
||||
error: str | None = None,
|
||||
):
|
||||
return RunnerOutput(
|
||||
req_id=req_id,
|
||||
request_id=req_id,
|
||||
step_index=step_index,
|
||||
finished=finished,
|
||||
result=DiffusionOutput(output=None, error=error) if error is not None else None,
|
||||
@@ -73,16 +73,16 @@ def _make_step_request(
|
||||
num_inference_steps=num_inference_steps,
|
||||
step_index=step_index,
|
||||
),
|
||||
request_ids=[req_id],
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
|
||||
def _new_ids(sched_output) -> list[str]:
|
||||
return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
|
||||
return [req.request_id for req in sched_output.scheduled_new_reqs]
|
||||
|
||||
|
||||
def _cached_ids(sched_output) -> list[str]:
|
||||
return list(sched_output.scheduled_cached_reqs.sched_req_ids)
|
||||
return list(sched_output.scheduled_cached_reqs.request_ids)
|
||||
|
||||
|
||||
class _StubScheduler(SchedulerInterface):
|
||||
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
|
||||
self._request = request
|
||||
self._output = output
|
||||
self.initialized_with = None
|
||||
self._sched_req_id = request.request_ids[0]
|
||||
self._request_id = request.request_id
|
||||
self._state = None
|
||||
self._scheduled = False
|
||||
self.max_num_running_reqs = 1
|
||||
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
|
||||
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str:
|
||||
assert request is self._request
|
||||
self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
|
||||
return self._sched_req_id
|
||||
self._state = SimpleNamespace(request_id=self._request_id, req=request)
|
||||
return self._request_id
|
||||
|
||||
def schedule(self):
|
||||
if self._scheduled or self._state is None:
|
||||
return SimpleNamespace(
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
scheduled_req_ids=[],
|
||||
scheduled_request_ids=[],
|
||||
is_empty=True,
|
||||
)
|
||||
self._scheduled = True
|
||||
return SimpleNamespace(
|
||||
scheduled_new_reqs=[NewRequestData.from_state(self._state)],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
scheduled_req_ids=[self._state.sched_req_id],
|
||||
scheduled_request_ids=[self._state.request_id],
|
||||
is_empty=False,
|
||||
)
|
||||
|
||||
@@ -123,30 +123,25 @@ class _StubScheduler(SchedulerInterface):
|
||||
del sched_output
|
||||
assert output is self._output
|
||||
self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
return {self._sched_req_id}
|
||||
return {self._request_id}
|
||||
|
||||
def has_requests(self) -> bool:
|
||||
return not self._scheduled
|
||||
|
||||
def get_request_state(self, sched_req_id: str):
|
||||
del sched_req_id
|
||||
def get_request_state(self, request_id: str):
|
||||
del request_id
|
||||
return self._state
|
||||
|
||||
def get_sched_req_id(self, request_id: str) -> str | None:
|
||||
if request_id in self._request.request_ids:
|
||||
return self._sched_req_id
|
||||
return None
|
||||
|
||||
def pop_request_state(self, sched_req_id: str):
|
||||
del sched_req_id
|
||||
def pop_request_state(self, request_id: str):
|
||||
del request_id
|
||||
return self._state
|
||||
|
||||
def preempt_request(self, sched_req_id: str) -> bool:
|
||||
del sched_req_id
|
||||
def preempt_request(self, request_id: str) -> bool:
|
||||
del request_id
|
||||
return False
|
||||
|
||||
def finish_requests(self, sched_req_ids, status) -> None:
|
||||
del sched_req_ids, status
|
||||
def finish_requests(self, request_ids, status) -> None:
|
||||
del request_ids, status
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
@@ -171,7 +166,7 @@ class TestGetSamplingParamsKey:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=["prompt"],
|
||||
sampling_params=sp,
|
||||
request_ids=[f"req-{lora_int_id}-{lora_scale}"],
|
||||
request_id=f"req-{lora_int_id}-{lora_scale}",
|
||||
)
|
||||
|
||||
def test_distinguishes_lora_id(self) -> None:
|
||||
@@ -291,7 +286,7 @@ class TestRequestScheduler:
|
||||
OmniDiffusionRequest(
|
||||
prompts=["prompt_b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(width=768),
|
||||
request_ids=["b"],
|
||||
request_id="b",
|
||||
)
|
||||
)
|
||||
scheduler.add_request(_make_request("c"))
|
||||
@@ -333,7 +328,7 @@ class TestRequestScheduler:
|
||||
assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
|
||||
|
||||
assert self.scheduler.has_requests() is False
|
||||
assert self.scheduler.schedule().scheduled_req_ids == []
|
||||
assert self.scheduler.schedule().scheduled_request_ids == []
|
||||
|
||||
def test_has_requests_state_transition(self) -> None:
|
||||
assert self.scheduler.has_requests() is False
|
||||
@@ -348,60 +343,28 @@ class TestRequestScheduler:
|
||||
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
assert self.scheduler.has_requests() is False
|
||||
|
||||
def test_request_id_mapping_lifecycle(self) -> None:
|
||||
def test_request_id_is_scheduler_key(self) -> None:
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=["prompt_map_a", "prompt_map_b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["map-a", "map-b"],
|
||||
request_id="map-parent",
|
||||
)
|
||||
|
||||
sched_req_id = self.scheduler.add_request(request)
|
||||
request_id = self.scheduler.add_request(request)
|
||||
|
||||
assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
|
||||
assert request_id == "map-parent"
|
||||
state = self.scheduler.get_request_state("map-parent")
|
||||
assert state.request_id == "map-parent"
|
||||
|
||||
self.scheduler.pop_request_state(sched_req_id)
|
||||
self.scheduler.pop_request_state("map-parent")
|
||||
|
||||
assert self.scheduler.get_sched_req_id("map-a") is None
|
||||
assert self.scheduler.get_sched_req_id("map-b") is None
|
||||
assert self.scheduler.get_sched_req_id("map-parent") is None
|
||||
assert self.scheduler.get_request_state("map-parent") is None
|
||||
|
||||
def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
|
||||
self.scheduler.add_request(
|
||||
OmniDiffusionRequest(
|
||||
prompts=["prompt_existing"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["existing-child"],
|
||||
request_id="duplicate-parent",
|
||||
)
|
||||
)
|
||||
def test_duplicate_request_id_is_rejected(self) -> None:
|
||||
self.scheduler.add_request(_make_request("dup"))
|
||||
|
||||
colliding_request = OmniDiffusionRequest(
|
||||
prompts=["prompt_unique"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["unique-child"],
|
||||
request_id="duplicate-parent",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate-parent"):
|
||||
self.scheduler.add_request(colliding_request)
|
||||
|
||||
assert self.scheduler.get_request_state("unique-child") is None
|
||||
assert self.scheduler.get_sched_req_id("unique-child") is None
|
||||
|
||||
valid_request = OmniDiffusionRequest(
|
||||
prompts=["prompt_unique"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["unique-child"],
|
||||
request_id="unique-parent",
|
||||
)
|
||||
sched_req_id = self.scheduler.add_request(valid_request)
|
||||
|
||||
assert sched_req_id == "unique-child"
|
||||
assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
|
||||
with pytest.raises(ValueError, match="request_id 'dup' is already active"):
|
||||
self.scheduler.add_request(_make_request("dup"))
|
||||
|
||||
|
||||
class TestDiffusionEngine:
|
||||
@@ -662,7 +625,7 @@ class TestStepScheduler:
|
||||
finished = self.scheduler.update_from_output(
|
||||
sched_output,
|
||||
RunnerOutput(
|
||||
req_id=req_id,
|
||||
request_id=req_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=None,
|
||||
@@ -934,7 +897,7 @@ class TestStepScheduler:
|
||||
|
||||
for _ in range(expected_steps - 1):
|
||||
sched_output = self.scheduler.schedule()
|
||||
assert sched_output.scheduled_req_ids == [req_id]
|
||||
assert sched_output.scheduled_request_ids == [req_id]
|
||||
next_step = request.sampling_params.step_index + 1
|
||||
assert (
|
||||
self.scheduler.update_from_output(
|
||||
@@ -945,7 +908,7 @@ class TestStepScheduler:
|
||||
)
|
||||
|
||||
final_output = self.scheduler.schedule()
|
||||
assert final_output.scheduled_req_ids == [req_id]
|
||||
assert final_output.scheduled_request_ids == [req_id]
|
||||
assert self.scheduler.update_from_output(
|
||||
final_output,
|
||||
_make_step_output(req_id, step_index=expected_steps, finished=True),
|
||||
|
||||
@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
|
||||
def _make_step_request(num_inference_steps: int = 2):
|
||||
return SimpleNamespace(
|
||||
prompts=["a prompt"],
|
||||
request_ids=["req-1"],
|
||||
request_id="req-1",
|
||||
sampling_params=SimpleNamespace(
|
||||
generator=None,
|
||||
seed=None,
|
||||
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt-{req_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
|
||||
request_ids=[req_id],
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
|
||||
return runner
|
||||
|
||||
|
||||
def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
|
||||
def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
|
||||
scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
||||
num_running_reqs=1,
|
||||
@@ -287,7 +287,7 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
|
||||
|
||||
def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
|
||||
"""Scheduler output for a homogeneous batch (one NewRequestData per req)."""
|
||||
new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
|
||||
new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=new_reqs,
|
||||
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
|
||||
)
|
||||
|
||||
|
||||
def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
|
||||
def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
|
||||
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
||||
num_running_reqs=1,
|
||||
num_waiting_reqs=0,
|
||||
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
|
||||
runner,
|
||||
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
|
||||
)
|
||||
output = result.get_req_output("req-1")
|
||||
output = result.get_request_output("req-1")
|
||||
|
||||
assert output.finished is True
|
||||
assert output.result is not None
|
||||
@@ -384,16 +384,16 @@ class TestRunner:
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
||||
first = result.get_req_output("req-1")
|
||||
assert first.req_id == "req-1"
|
||||
first = result.get_request_output("req-1")
|
||||
assert first.request_id == "req-1"
|
||||
assert first.step_index == 1
|
||||
assert first.finished is False
|
||||
assert first.result is None
|
||||
assert "req-1" in runner.state_cache
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
||||
second = result.get_req_output("req-1")
|
||||
assert second.req_id == "req-1"
|
||||
second = result.get_request_output("req-1")
|
||||
assert second.request_id == "req-1"
|
||||
assert second.step_index == 2
|
||||
assert second.finished is True
|
||||
assert second.result is not None
|
||||
@@ -410,13 +410,13 @@ class TestRunner:
|
||||
runner = _make_runner()
|
||||
req_1 = _make_step_request()
|
||||
req_2 = _make_step_request()
|
||||
req_2.request_ids = ["req-2"]
|
||||
req_2.request_id = "req-2"
|
||||
|
||||
scheduler_output = DiffusionSchedulerOutput(
|
||||
step_id=0,
|
||||
scheduled_new_reqs=[
|
||||
NewRequestData(sched_req_id="req-1", req=req_1),
|
||||
NewRequestData(sched_req_id="req-2", req=req_2),
|
||||
NewRequestData(request_id="req-1", req=req_1),
|
||||
NewRequestData(request_id="req-2", req=req_2),
|
||||
],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
finished_req_ids=set(),
|
||||
@@ -431,7 +431,7 @@ class TestRunner:
|
||||
runner = _make_runner()
|
||||
|
||||
with pytest.raises(ValueError, match="Missing cached state"):
|
||||
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
|
||||
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
|
||||
|
||||
def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
@@ -440,8 +440,8 @@ class TestRunner:
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
||||
output = result.get_req_output("req-1")
|
||||
assert output.req_id == "req-1"
|
||||
output = result.get_request_output("req-1")
|
||||
assert output.request_id == "req-1"
|
||||
assert output.step_index == 0
|
||||
assert output.finished is True
|
||||
assert output.result is not None
|
||||
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.lora_manager = lora_manager
|
||||
worker._step_lora_state = {}
|
||||
output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
|
||||
output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
|
||||
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
|
||||
return worker
|
||||
|
||||
@@ -525,9 +525,9 @@ class TestWorker:
|
||||
"""DiffusionWorker.execute_stepwise"""
|
||||
|
||||
def test_delegates_to_model_runner(self):
|
||||
expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
|
||||
expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
|
||||
worker = _make_step_worker(expected_output=expected)
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
||||
|
||||
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -536,7 +536,7 @@ class TestWorker:
|
||||
def test_deactivates_lora_when_request_has_no_adapter(self):
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -552,7 +552,7 @@ class TestWorker:
|
||||
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(request, request_id="req-1")
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -568,8 +568,8 @@ class TestWorker:
|
||||
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
first = _make_scheduler_output(request, sched_req_id="req-1")
|
||||
second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
|
||||
first = _make_scheduler_output(request, request_id="req-1")
|
||||
second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, first)
|
||||
DiffusionWorker.execute_stepwise(worker, second)
|
||||
@@ -610,10 +610,10 @@ class TestWorker:
|
||||
next_request.sampling_params.lora_request = lora_request
|
||||
|
||||
worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
|
||||
first = _make_scheduler_output(finishing, sched_req_id="req-1")
|
||||
first = _make_scheduler_output(finishing, request_id="req-1")
|
||||
next_batch = _make_scheduler_output(
|
||||
next_request,
|
||||
sched_req_id="req-2",
|
||||
request_id="req-2",
|
||||
step_id=1,
|
||||
finished_req_ids={"req-1"},
|
||||
)
|
||||
@@ -633,11 +633,11 @@ class TestExecutor:
|
||||
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
|
||||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||||
executor._ensure_open = lambda: None
|
||||
expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
|
||||
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
|
||||
executor.collective_rpc = mocker.Mock(return_value=expected)
|
||||
|
||||
request = _make_engine_request("req-step", num_inference_steps=2)
|
||||
scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
|
||||
scheduler_output = _make_scheduler_output(request, request_id="req-step")
|
||||
|
||||
output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
|
||||
|
||||
@@ -653,7 +653,7 @@ class TestEngine:
|
||||
[
|
||||
(
|
||||
lambda _: RunnerOutput(
|
||||
req_id="req-error",
|
||||
request_id="req-error",
|
||||
step_index=1,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error="boom"),
|
||||
@@ -688,7 +688,7 @@ class TestEngine:
|
||||
call_count["n"] += 1
|
||||
finished = call_count["n"] == 2
|
||||
return RunnerOutput(
|
||||
req_id="req-step",
|
||||
request_id="req-step",
|
||||
step_index=call_count["n"],
|
||||
finished=finished,
|
||||
result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
|
||||
@@ -714,7 +714,7 @@ class TestEngine:
|
||||
step["n"] += 1
|
||||
engine.abort("req-stop")
|
||||
return RunnerOutput(
|
||||
req_id="req-stop",
|
||||
request_id="req-stop",
|
||||
step_index=1,
|
||||
finished=False,
|
||||
result=None,
|
||||
@@ -741,7 +741,7 @@ class TestEngine:
|
||||
assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
|
||||
engine.abort("req-mid")
|
||||
return RunnerOutput(
|
||||
req_id="req-mid",
|
||||
request_id="req-mid",
|
||||
step_index=step["n"],
|
||||
finished=False,
|
||||
result=None,
|
||||
@@ -760,7 +760,7 @@ class TestEngine:
|
||||
engine = _make_engine(
|
||||
scheduler,
|
||||
execute_fn=lambda _: RunnerOutput(
|
||||
req_id="req-missing",
|
||||
request_id="req-missing",
|
||||
step_index=1,
|
||||
finished=True,
|
||||
result=None,
|
||||
@@ -777,7 +777,7 @@ class TestEngine:
|
||||
class TestIPC:
|
||||
def test_pack_unpack_runner_output_shm(self):
|
||||
tensor = torch.zeros(300_000, dtype=torch.float32)
|
||||
output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
|
||||
output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
|
||||
|
||||
packed = pack_diffusion_output_shm(output)
|
||||
assert isinstance(packed.result.output, dict)
|
||||
|
||||
@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
|
||||
def _mock_request(tag: str):
|
||||
"""Return a lightweight request object identifiable by *tag*."""
|
||||
return SimpleNamespace(
|
||||
request_ids=[tag],
|
||||
request_id=tag,
|
||||
prompts=[f"prompt_{tag}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
)
|
||||
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
|
||||
req = req_q.get(timeout=10)
|
||||
method = req.get("method", "")
|
||||
args = req.get("args", ())
|
||||
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
|
||||
tag = f"result_for_{args[0].request_ids[0]}"
|
||||
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
|
||||
tag = f"result_for_{args[0].request_id}"
|
||||
elif args:
|
||||
tag = f"result_for_{args[0]}"
|
||||
else:
|
||||
|
||||
@@ -68,7 +68,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
|
||||
|
||||
request = captured["request"]
|
||||
assert request.request_id == "req-parent"
|
||||
assert request.request_ids == ["req-parent-0", "req-parent-1"]
|
||||
assert not hasattr(request, "request_ids")
|
||||
assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
|
||||
assert result.request_id == "req-parent"
|
||||
assert result.images == ["img-1", "img-2"]
|
||||
|
||||
@@ -384,7 +384,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test_recv"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=[req_id],
|
||||
request_id=req_id,
|
||||
)
|
||||
# req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
|
||||
# But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
|
||||
@@ -428,7 +428,6 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["prompt-a", "prompt-b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
|
||||
request_id=parent_req_id,
|
||||
)
|
||||
|
||||
@@ -454,7 +453,6 @@ def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_con
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["prompt-a", "prompt-b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=["req-parent-0", "req-parent-1"],
|
||||
request_id="req-parent",
|
||||
)
|
||||
req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
|
||||
@@ -516,7 +514,7 @@ def test_integration_flow(common_constants):
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test_integ"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_ids=[req_id],
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
# Receive
|
||||
|
||||
@@ -202,6 +202,7 @@ class TeaCacheCoefficientEstimator:
|
||||
self.hook.start_collection()
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=[prompt],
|
||||
request_id="teacache-coefficient-estimator",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_inference_steps=generate_kwargs.get("num_inference_steps", 20),
|
||||
seed=generate_kwargs.get("seed", 42),
|
||||
|
||||
@@ -30,10 +30,10 @@ from vllm_omni.diffusion.registry import (
|
||||
get_diffusion_post_process_func,
|
||||
get_diffusion_pre_process_func,
|
||||
)
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
|
||||
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -217,7 +217,7 @@ class DiffusionEngine:
|
||||
logger.warning("Output is None, returning empty OmniRequestOutput")
|
||||
return [
|
||||
OmniRequestOutput.from_diffusion(
|
||||
request_id=request.request_ids[i] if i < len(request.request_ids) else "",
|
||||
request_id=request.request_id,
|
||||
images=[],
|
||||
prompt=prompt,
|
||||
metrics={},
|
||||
@@ -304,7 +304,7 @@ class DiffusionEngine:
|
||||
if len(request.prompts) == 1:
|
||||
# Single request: return single OmniRequestOutput
|
||||
prompt = request.prompts[0]
|
||||
request_id = request.request_ids[0] if request.request_ids else ""
|
||||
request_id = request.request_id
|
||||
|
||||
if is_text_output:
|
||||
return [
|
||||
@@ -368,10 +368,9 @@ class DiffusionEngine:
|
||||
# Split images based on num_outputs_per_prompt for each request
|
||||
results = []
|
||||
output_idx = 0
|
||||
request_id = request.request_id
|
||||
|
||||
for i, prompt in enumerate(request.prompts):
|
||||
request_id = request.request_ids[i] if i < len(request.request_ids) else ""
|
||||
|
||||
# Get images for this request
|
||||
num_outputs = request.sampling_params.num_outputs_per_prompt
|
||||
start_idx = output_idx
|
||||
@@ -467,17 +466,17 @@ class DiffusionEngine:
|
||||
runner_output = self.execute_fn(sched_output)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
|
||||
"Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
|
||||
)
|
||||
runner_output = BatchRunnerOutput.from_list(
|
||||
[
|
||||
RunnerOutput(
|
||||
req_id=req_id,
|
||||
request_id=request_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error=str(exc)),
|
||||
)
|
||||
for req_id in sched_output.scheduled_req_ids
|
||||
for request_id in sched_output.scheduled_request_ids
|
||||
]
|
||||
)
|
||||
|
||||
@@ -544,7 +543,7 @@ class DiffusionEngine:
|
||||
def _handle_finished_requests(
|
||||
self,
|
||||
finished_ids: set[str],
|
||||
runner_output: RunnerOutput | None = None,
|
||||
runner_output: BaseRunnerOutput | None = None,
|
||||
missing_result_error: str = "Diffusion execution finished without a final output",
|
||||
):
|
||||
for rid in finished_ids:
|
||||
@@ -553,7 +552,7 @@ class DiffusionEngine:
|
||||
if fut is None:
|
||||
continue
|
||||
if runner_output is not None:
|
||||
_output = runner_output.get_req_output(rid)
|
||||
_output = runner_output.get_request_output(rid)
|
||||
else:
|
||||
_output = None
|
||||
out = self._finalize_finished_request(rid, _output, missing_result_error)
|
||||
@@ -579,11 +578,11 @@ class DiffusionEngine:
|
||||
if self._closed:
|
||||
raise RuntimeError("DiffusionEngine is closed.")
|
||||
fut = self.main_loop.create_future()
|
||||
sched_req_id = self.scheduler.add_request(request)
|
||||
self._out_queue[sched_req_id] = fut
|
||||
request_id = self.scheduler.add_request(request)
|
||||
self._out_queue[request_id] = fut
|
||||
self._cv.notify_all()
|
||||
|
||||
return sched_req_id
|
||||
return request_id
|
||||
|
||||
async def get_result(self, request_id: str) -> DiffusionOutput:
|
||||
fut = self._out_queue.get(request_id)
|
||||
@@ -599,37 +598,37 @@ class DiffusionEngine:
|
||||
async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
# No lock needed: add_request is already protected by self._cv, and
|
||||
# all executor calls are serialized inside the busy loop.
|
||||
sched_req_id = self.add_request(request)
|
||||
return await self.get_result(sched_req_id)
|
||||
request_id = self.add_request(request)
|
||||
return await self.get_result(request_id)
|
||||
|
||||
def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
with self._rpc_lock:
|
||||
if self._closed:
|
||||
raise RuntimeError("DiffusionEngine is closed.")
|
||||
target_sched_req_id = self.scheduler.add_request(request)
|
||||
target_request_id = self.scheduler.add_request(request)
|
||||
|
||||
# keep scheduling and executing until the target request is finished
|
||||
while True:
|
||||
self._process_aborts_queue()
|
||||
sched_output = self.scheduler.schedule()
|
||||
if sched_output.is_empty:
|
||||
if target_sched_req_id in sched_output.finished_req_ids:
|
||||
return self._finalize_finished_request(target_sched_req_id)
|
||||
if target_request_id in sched_output.finished_req_ids:
|
||||
return self._finalize_finished_request(target_request_id)
|
||||
if not self.scheduler.has_requests():
|
||||
raise RuntimeError("Diffusion scheduler has no runnable requests.")
|
||||
continue
|
||||
|
||||
# NOTE: add_req_and_wait_for_response() is synchronous, will be only called
|
||||
# within _dummy_run, only one request will be scheduled
|
||||
sched_req_id = sched_output.scheduled_req_ids[0]
|
||||
request_id = sched_output.scheduled_request_ids[0]
|
||||
try:
|
||||
runner_output = self.execute_fn(sched_output)
|
||||
except EngineDeadError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
|
||||
logger.error("Execution failed for diffusion request %s", request_id, exc_info=True)
|
||||
runner_output = RunnerOutput(
|
||||
req_id=sched_req_id,
|
||||
request_id=request_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error=str(exc)),
|
||||
@@ -642,10 +641,10 @@ class DiffusionEngine:
|
||||
# sync func should receive one result
|
||||
if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
|
||||
raise ValueError("Sync func should receive one result at one time")
|
||||
if target_sched_req_id in finished_req_ids:
|
||||
runner_output = runner_output.get_req_output(target_sched_req_id)
|
||||
if target_request_id in finished_req_ids:
|
||||
runner_output = runner_output.get_request_output(target_request_id)
|
||||
return self._finalize_finished_request(
|
||||
target_sched_req_id,
|
||||
target_request_id,
|
||||
runner_output=runner_output,
|
||||
missing_result_error="Diffusion execution finished without a final output.",
|
||||
)
|
||||
@@ -697,7 +696,7 @@ class DiffusionEngine:
|
||||
num_frames = 2 if supports_audio_input or supports_audio_output(self.od_config.model_class_name) else 1
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=[prompt],
|
||||
request_ids=["dummy_req_id"],
|
||||
request_id=DUMMY_DIFFUSION_REQUEST_ID,
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -904,34 +903,27 @@ class DiffusionEngine:
|
||||
def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
|
||||
request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
|
||||
|
||||
sched_req_ids: list[str] = []
|
||||
for request_id in dict.fromkeys(request_ids):
|
||||
sched_req_id = self.scheduler.get_sched_req_id(request_id)
|
||||
if sched_req_id is not None:
|
||||
sched_req_ids.append(sched_req_id)
|
||||
|
||||
for sched_req_id in dict.fromkeys(sched_req_ids):
|
||||
if self.scheduler.get_request_state(sched_req_id) is not None:
|
||||
self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
|
||||
if self.scheduler.get_request_state(request_id) is not None:
|
||||
self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
|
||||
|
||||
def _finalize_finished_request(
|
||||
self,
|
||||
sched_req_id: str,
|
||||
request_id: str,
|
||||
runner_output: RunnerOutput | None = None,
|
||||
missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
|
||||
) -> DiffusionOutput:
|
||||
state = self.scheduler.get_request_state(sched_req_id)
|
||||
popped_state = self.scheduler.pop_request_state(sched_req_id)
|
||||
state = self.scheduler.get_request_state(request_id)
|
||||
popped_state = self.scheduler.pop_request_state(request_id)
|
||||
state = state or popped_state
|
||||
|
||||
if state is None:
|
||||
raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
|
||||
raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")
|
||||
|
||||
if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
|
||||
request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
|
||||
return DiffusionOutput(
|
||||
aborted=True,
|
||||
abort_message=f"Request {request_id} aborted.",
|
||||
abort_message=f"Request {state.req.request_id} aborted.",
|
||||
)
|
||||
|
||||
if runner_output is not None and runner_output.result is not None:
|
||||
|
||||
@@ -311,7 +311,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
raise RuntimeError(f"Unexpected response type for execute_request: {type(result)!r}")
|
||||
|
||||
return RunnerOutput(
|
||||
req_id=new_req.sched_req_id,
|
||||
request_id=new_req.request_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=result,
|
||||
@@ -334,9 +334,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
# TODO: Remove this fallback; DiffusionOutput cannot faithfully represent
|
||||
# failed multi-request step batches.
|
||||
if isinstance(result, DiffusionOutput):
|
||||
req_id = scheduler_output.scheduled_req_ids[0] if scheduler_output.scheduled_req_ids else ""
|
||||
request_id = scheduler_output.scheduled_request_ids[0] if scheduler_output.scheduled_request_ids else ""
|
||||
return RunnerOutput(
|
||||
req_id=req_id,
|
||||
request_id=request_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=result,
|
||||
|
||||
@@ -128,7 +128,6 @@ class InlineStageDiffusionClient(StageClientBase):
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=[prompt],
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[request_id],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
@@ -187,7 +186,6 @@ class InlineStageDiffusionClient(StageClientBase):
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=prompts,
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
|
||||
@@ -247,8 +247,10 @@ class ErnieImagePipeline(
|
||||
|
||||
@staticmethod
|
||||
def _is_warmup_request(req: OmniDiffusionRequest) -> bool:
|
||||
request_ids = getattr(req, "request_ids", None) or ()
|
||||
return len(request_ids) == 1 and request_ids[0] == "dummy_req_id"
|
||||
is_dummy_run = getattr(req, "is_dummy_run", None)
|
||||
if callable(is_dummy_run):
|
||||
return bool(is_dummy_run())
|
||||
return OmniDiffusionRequest.is_dummy_run_request_id(getattr(req, "request_id", None))
|
||||
|
||||
@staticmethod
|
||||
def _should_apply_pe(req: OmniDiffusionRequest) -> bool:
|
||||
|
||||
@@ -689,15 +689,11 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
|
||||
|
||||
# NOTE: DiffusionEngine does an internal warmup "dummy run" during
|
||||
# initialization. That request has no request_id and does not carry
|
||||
# initialization. That request carries the fixed dummy request_id and does not carry
|
||||
# Stage-0 (AR) outputs via req.extra. For GLM-Image, we allow that
|
||||
# specific warmup request to proceed by synthesizing minimal prior
|
||||
# tokens, while still raising a clear error for real requests.
|
||||
is_dummy_warmup = (
|
||||
not getattr(req, "request_id", None)
|
||||
and prompt == "dummy run"
|
||||
and (req.sampling_params.num_inference_steps == 1)
|
||||
)
|
||||
is_dummy_warmup = req.is_dummy_run()
|
||||
|
||||
# Get pre-computed prompt embeddings if provided
|
||||
if isinstance(first_prompt, str):
|
||||
|
||||
@@ -3,10 +3,12 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
||||
|
||||
DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id"
|
||||
|
||||
|
||||
@dataclass
|
||||
class OmniDiffusionRequest:
|
||||
@@ -24,13 +26,14 @@ class OmniDiffusionRequest:
|
||||
|
||||
prompts: list[OmniPromptType] # Actually supporting str-based prompts
|
||||
sampling_params: OmniDiffusionSamplingParams
|
||||
|
||||
request_ids: list[str] = field(default_factory=list)
|
||||
request_id: str | None = None
|
||||
request_id: str
|
||||
kv_sender_info: dict | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize dependent fields after dataclass initialization."""
|
||||
if not isinstance(self.request_id, str) or not self.request_id:
|
||||
raise ValueError("OmniDiffusionRequest.request_id must be a non-empty string.")
|
||||
|
||||
# When neither a generator nor a seed is provided, assign a random seed
|
||||
# so that all ranks derive the same generator state.
|
||||
if self.sampling_params.generator is None and self.sampling_params.seed is None:
|
||||
@@ -56,3 +59,10 @@ class OmniDiffusionRequest:
|
||||
# so downstream code always has a valid value.
|
||||
if self.sampling_params.guidance_scale_2 is None:
|
||||
self.sampling_params.guidance_scale_2 = self.sampling_params.guidance_scale
|
||||
|
||||
def is_dummy_run(self) -> bool:
|
||||
return self.is_dummy_run_request_id(self.request_id)
|
||||
|
||||
@classmethod
|
||||
def is_dummy_run_request_id(cls, request_id: str | None) -> bool:
|
||||
return request_id == DUMMY_DIFFUSION_REQUEST_ID
|
||||
@@ -46,7 +46,6 @@ class _BaseScheduler(SchedulerInterface):
|
||||
def __init__(self) -> None:
|
||||
self.od_config: OmniDiffusionConfig | None = None
|
||||
self._request_states: dict[str, DiffusionRequestState] = {}
|
||||
self._request_id_to_sched_req_id: dict[str, str] = {}
|
||||
self._step_id: int = 0
|
||||
self._waiting: deque[str] = deque()
|
||||
self._running: list[str] = []
|
||||
@@ -57,7 +56,6 @@ class _BaseScheduler(SchedulerInterface):
|
||||
def initialize(self, od_config: OmniDiffusionConfig) -> None:
|
||||
self.od_config = od_config
|
||||
self._request_states.clear()
|
||||
self._request_id_to_sched_req_id.clear()
|
||||
self._step_id = 0
|
||||
self._waiting.clear()
|
||||
self._running.clear()
|
||||
@@ -71,32 +69,31 @@ class _BaseScheduler(SchedulerInterface):
|
||||
self._reset_scheduler_state()
|
||||
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str:
|
||||
sched_req_id = self._make_sched_req_id(request)
|
||||
return self._add_request_with_sched_req_id(sched_req_id, request)
|
||||
return self._add_request_with_request_id(request.request_id, request)
|
||||
|
||||
def _add_request_with_sched_req_id(self, sched_req_id: str, request: OmniDiffusionRequest) -> str:
|
||||
state = self._make_request_state(sched_req_id, request)
|
||||
request_ids = self._request_ids_for_mapping(request)
|
||||
self._register_request_ids(request_ids, sched_req_id)
|
||||
self._request_states[sched_req_id] = state
|
||||
self._waiting.append(sched_req_id)
|
||||
logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, sched_req_id, len(self._waiting))
|
||||
return sched_req_id
|
||||
def _add_request_with_request_id(self, request_id: str, request: OmniDiffusionRequest) -> str:
|
||||
if request_id in self._request_states:
|
||||
raise ValueError(f"request_id {request_id!r} is already active.")
|
||||
state = self._make_request_state(request_id, request)
|
||||
self._request_states[request_id] = state
|
||||
self._waiting.append(request_id)
|
||||
logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting))
|
||||
return request_id
|
||||
|
||||
def schedule(self) -> DiffusionSchedulerOutput:
|
||||
scheduled_new_reqs: list[NewRequestData] = []
|
||||
scheduled_cached_req_ids: list[str] = []
|
||||
scheduled_cached_request_ids: list[str] = []
|
||||
|
||||
# First, schedule the RUNNING request(s)
|
||||
for sched_req_id in self._running:
|
||||
state = self._request_states.get(sched_req_id)
|
||||
for request_id in self._running:
|
||||
state = self._request_states.get(request_id)
|
||||
if state is not None:
|
||||
scheduled_cached_req_ids.append(sched_req_id)
|
||||
scheduled_cached_request_ids.append(request_id)
|
||||
|
||||
# Second, schedule WAITING requests while capacity remains.
|
||||
while self._waiting and len(self._running) < self.max_num_running_reqs:
|
||||
sched_req_id = self._waiting[0]
|
||||
state = self._request_states.get(sched_req_id)
|
||||
request_id = self._waiting[0]
|
||||
state = self._request_states.get(request_id)
|
||||
if state is None:
|
||||
self._waiting.popleft()
|
||||
continue
|
||||
@@ -108,16 +105,16 @@ class _BaseScheduler(SchedulerInterface):
|
||||
if not self._running:
|
||||
self._running_sampling_params_key = state.sampling_params_key
|
||||
state.status = DiffusionRequestStatus.RUNNING
|
||||
self._running.append(sched_req_id)
|
||||
self._running.append(request_id)
|
||||
if was_new_request:
|
||||
scheduled_new_reqs.append(NewRequestData.from_state(state))
|
||||
else:
|
||||
scheduled_cached_req_ids.append(sched_req_id)
|
||||
scheduled_cached_request_ids.append(request_id)
|
||||
|
||||
scheduler_output = DiffusionSchedulerOutput(
|
||||
step_id=self._step_id,
|
||||
scheduled_new_reqs=scheduled_new_reqs,
|
||||
scheduled_cached_reqs=CachedRequestData(sched_req_ids=scheduled_cached_req_ids),
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=scheduled_cached_request_ids),
|
||||
finished_req_ids=set(self._finished_req_ids),
|
||||
num_running_reqs=len(self._running),
|
||||
num_waiting_reqs=len(self._waiting),
|
||||
@@ -131,40 +128,33 @@ class _BaseScheduler(SchedulerInterface):
|
||||
def has_requests(self) -> bool:
|
||||
return bool(self._waiting or self._running)
|
||||
|
||||
def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
|
||||
return self._request_states.get(sched_req_id)
|
||||
def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
|
||||
return self._request_states.get(request_id)
|
||||
|
||||
def get_sched_req_id(self, request_id: str) -> str | None:
|
||||
return self._request_id_to_sched_req_id.get(request_id)
|
||||
def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
|
||||
self._pop_extra_request_state(request_id)
|
||||
return self._request_states.pop(request_id, None)
|
||||
|
||||
def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
|
||||
self._pop_extra_request_state(sched_req_id)
|
||||
state = self._request_states.pop(sched_req_id, None)
|
||||
if state is not None:
|
||||
self._unregister_request_ids(self._request_ids_for_mapping(state.req), sched_req_id)
|
||||
return state
|
||||
|
||||
def preempt_request(self, sched_req_id: str) -> bool:
|
||||
if sched_req_id not in self._request_states:
|
||||
def preempt_request(self, request_id: str) -> bool:
|
||||
if request_id not in self._request_states:
|
||||
return False
|
||||
if sched_req_id in self._running:
|
||||
self._running.remove(sched_req_id)
|
||||
if request_id in self._running:
|
||||
self._running.remove(request_id)
|
||||
if not self._running:
|
||||
self._running_sampling_params_key = None
|
||||
self._waiting.appendleft(sched_req_id)
|
||||
self._request_states[sched_req_id].status = DiffusionRequestStatus.PREEMPTED
|
||||
self._waiting.appendleft(request_id)
|
||||
self._request_states[request_id].status = DiffusionRequestStatus.PREEMPTED
|
||||
return True
|
||||
return False
|
||||
|
||||
def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
|
||||
def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
|
||||
assert DiffusionRequestStatus.is_finished(status)
|
||||
if isinstance(sched_req_ids, str):
|
||||
sched_req_ids = [sched_req_ids]
|
||||
self._finish_requests({sched_req_id: status for sched_req_id in sched_req_ids})
|
||||
if isinstance(request_ids, str):
|
||||
request_ids = [request_ids]
|
||||
self._finish_requests({request_id: status for request_id in request_ids})
|
||||
|
||||
def close(self) -> None:
|
||||
self._request_states.clear()
|
||||
self._request_id_to_sched_req_id.clear()
|
||||
self._waiting.clear()
|
||||
self._running.clear()
|
||||
self._running_sampling_params_key = None
|
||||
@@ -183,33 +173,31 @@ class _BaseScheduler(SchedulerInterface):
|
||||
running_to_remove: set[str] = set()
|
||||
waiting_to_remove: set[str] = set()
|
||||
|
||||
for sched_req_id, status in statuses.items():
|
||||
for request_id, status in statuses.items():
|
||||
assert DiffusionRequestStatus.is_finished(status)
|
||||
state = self._request_states.get(sched_req_id)
|
||||
state = self._request_states.get(request_id)
|
||||
if state is None or state.is_finished():
|
||||
continue
|
||||
|
||||
finished_req_ids.add(sched_req_id)
|
||||
if sched_req_id in self._running:
|
||||
running_to_remove.add(sched_req_id)
|
||||
if sched_req_id in self._waiting:
|
||||
waiting_to_remove.add(sched_req_id)
|
||||
finished_req_ids.add(request_id)
|
||||
if request_id in self._running:
|
||||
running_to_remove.add(request_id)
|
||||
if request_id in self._waiting:
|
||||
waiting_to_remove.add(request_id)
|
||||
|
||||
if running_to_remove:
|
||||
self._running = [sched_req_id for sched_req_id in self._running if sched_req_id not in running_to_remove]
|
||||
self._running = [request_id for request_id in self._running if request_id not in running_to_remove]
|
||||
if not self._running:
|
||||
self._running_sampling_params_key = None
|
||||
if waiting_to_remove:
|
||||
self._waiting = deque(
|
||||
sched_req_id for sched_req_id in self._waiting if sched_req_id not in waiting_to_remove
|
||||
)
|
||||
self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove)
|
||||
|
||||
for sched_req_id in finished_req_ids:
|
||||
state = self._request_states[sched_req_id]
|
||||
status = statuses[sched_req_id]
|
||||
for request_id in finished_req_ids:
|
||||
state = self._request_states[request_id]
|
||||
status = statuses[request_id]
|
||||
state.status = status
|
||||
if status == DiffusionRequestStatus.FINISHED_ERROR:
|
||||
state.error = None if errors is None else errors.get(sched_req_id)
|
||||
state.error = None if errors is None else errors.get(request_id)
|
||||
else:
|
||||
state.error = None
|
||||
|
||||
@@ -227,7 +215,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
# marked finished at that point, but we still need to surface its id
|
||||
# in this update so the engine can observe the terminal state.
|
||||
finished_req_ids = {
|
||||
sched_req_id for sched_req_id in sched_output.scheduled_req_ids if sched_req_id in self._finished_req_ids
|
||||
request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids
|
||||
}
|
||||
finished_req_ids |= self._finish_requests(statuses, errors)
|
||||
return finished_req_ids
|
||||
@@ -235,12 +223,12 @@ class _BaseScheduler(SchedulerInterface):
|
||||
def _reset_scheduler_state(self) -> None:
|
||||
"""Reset subclass-owned state during initialize()/close()."""
|
||||
|
||||
def _pop_extra_request_state(self, sched_req_id: str) -> None:
|
||||
def _pop_extra_request_state(self, request_id: str) -> None:
|
||||
"""Remove subclass-owned per-request state before popping request state."""
|
||||
|
||||
def _make_request_state(self, sched_req_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
|
||||
def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
|
||||
return DiffusionRequestState(
|
||||
sched_req_id=sched_req_id,
|
||||
request_id=request_id,
|
||||
req=request,
|
||||
sampling_params_key=get_sampling_params_key(request),
|
||||
)
|
||||
@@ -258,24 +246,3 @@ class _BaseScheduler(SchedulerInterface):
|
||||
state = self._request_states.get(self._running[0])
|
||||
self._running_sampling_params_key = None if state is None else state.sampling_params_key
|
||||
return self._running_sampling_params_key
|
||||
|
||||
def _register_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
|
||||
unique_request_ids = list(dict.fromkeys(request_ids))
|
||||
for request_id in unique_request_ids:
|
||||
existing = self._request_id_to_sched_req_id.get(request_id)
|
||||
if existing is not None and existing != sched_req_id:
|
||||
raise ValueError(f"request_id {request_id!r} is already mapped to active sched_req_id {existing!r}.")
|
||||
for request_id in unique_request_ids:
|
||||
self._request_id_to_sched_req_id[request_id] = sched_req_id
|
||||
|
||||
def _unregister_request_ids(self, request_ids: list[str], sched_req_id: str) -> None:
|
||||
for request_id in request_ids:
|
||||
if self._request_id_to_sched_req_id.get(request_id) == sched_req_id:
|
||||
self._request_id_to_sched_req_id.pop(request_id, None)
|
||||
|
||||
@staticmethod
|
||||
def _request_ids_for_mapping(request: OmniDiffusionRequest) -> list[str]:
|
||||
request_ids = list(request.request_ids)
|
||||
if request_id := getattr(request, "request_id", None):
|
||||
request_ids.append(request_id)
|
||||
return list(dict.fromkeys(request_ids))
|
||||
@@ -4,22 +4,17 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.worker.utils import RunnerOutput
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DiffusionRequestStatus(enum.IntEnum):
|
||||
"""Request status tracked by diffusion scheduler."""
|
||||
@@ -75,10 +70,7 @@ class SamplingParamsKey:
|
||||
class DiffusionRequestState:
|
||||
"""Scheduler-owned state for one queued OmniDiffusionRequest."""
|
||||
|
||||
# Unique scheduler-owned request ID.
|
||||
# NOTE: This identifies one OmniDiffusionRequest, which may contain multiple request_ids.
|
||||
# TODO: Align this with OmniDiffusionRequest.request_ids once scheduler batching is supported.
|
||||
sched_req_id: str
|
||||
request_id: str
|
||||
req: OmniDiffusionRequest
|
||||
sampling_params_key: SamplingParamsKey | None = None
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
@@ -92,23 +84,23 @@ class DiffusionRequestState:
|
||||
class NewRequestData:
|
||||
"""Full request payload for a newly scheduled diffusion request."""
|
||||
|
||||
sched_req_id: str
|
||||
request_id: str
|
||||
req: OmniDiffusionRequest
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state: DiffusionRequestState) -> NewRequestData:
|
||||
return cls(sched_req_id=state.sched_req_id, req=state.req)
|
||||
return cls(request_id=state.request_id, req=state.req)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CachedRequestData:
|
||||
"""Cached diffusion requests that only need their scheduler ids resent."""
|
||||
"""Cached diffusion requests that only need their request ids resent."""
|
||||
|
||||
sched_req_ids: list[str]
|
||||
request_ids: list[str]
|
||||
|
||||
@classmethod
|
||||
def make_empty(cls) -> CachedRequestData:
|
||||
return cls(sched_req_ids=[])
|
||||
return cls(request_ids=[])
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -123,21 +115,18 @@ class DiffusionSchedulerOutput:
|
||||
num_waiting_reqs: int
|
||||
|
||||
@cached_property
|
||||
def scheduled_req_ids(self) -> list[str]:
|
||||
def scheduled_request_ids(self) -> list[str]:
|
||||
"""
|
||||
All scheduled request ids in this cycle, including both new and cached ones.
|
||||
NOTE:
|
||||
This id is generated and owned by the scheduler,
|
||||
and may be different from the OmniDiffusionRequest.request_ids.
|
||||
"""
|
||||
return [
|
||||
*(req.sched_req_id for req in self.scheduled_new_reqs),
|
||||
*self.scheduled_cached_reqs.sched_req_ids,
|
||||
*(req.request_id for req in self.scheduled_new_reqs),
|
||||
*self.scheduled_cached_reqs.request_ids,
|
||||
]
|
||||
|
||||
@property
|
||||
def num_scheduled_reqs(self) -> int:
|
||||
return len(self.scheduled_req_ids)
|
||||
return len(self.scheduled_request_ids)
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
@@ -147,25 +136,6 @@ class DiffusionSchedulerOutput:
|
||||
class SchedulerInterface(ABC):
|
||||
"""Abstract lifecycle contract for diffusion schedulers."""
|
||||
|
||||
def _make_sched_req_id(self, request: OmniDiffusionRequest) -> str:
|
||||
"""
|
||||
Generate a unique scheduler request ID for the given request.
|
||||
The default implementation uses the first request_id from the request if available,
|
||||
otherwise generates a random one.
|
||||
"""
|
||||
if request.request_ids:
|
||||
base = request.request_ids[0]
|
||||
else:
|
||||
logger.warning("Request has no request_ids, generating a random one. Request: %s", request)
|
||||
base = f"req_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
sched_req_id = base
|
||||
suffix = 1
|
||||
while self.get_request_state(sched_req_id) is not None:
|
||||
sched_req_id = f"{base}#{suffix}"
|
||||
suffix += 1
|
||||
return sched_req_id
|
||||
|
||||
@abstractmethod
|
||||
def initialize(self, od_config: OmniDiffusionConfig) -> None:
|
||||
"""Initialize or reset scheduler state."""
|
||||
@@ -183,7 +153,7 @@ class SchedulerInterface(ABC):
|
||||
"""Update scheduler state from executor output."""
|
||||
|
||||
@abstractmethod
|
||||
def get_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
|
||||
def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
|
||||
"""Return request state if present."""
|
||||
|
||||
@abstractmethod
|
||||
@@ -191,19 +161,15 @@ class SchedulerInterface(ABC):
|
||||
"""Return whether the scheduler still owns runnable requests."""
|
||||
|
||||
@abstractmethod
|
||||
def get_sched_req_id(self, request_id: str) -> str | None:
|
||||
"""Resolve a public request_id to the active scheduler request id."""
|
||||
|
||||
@abstractmethod
|
||||
def pop_request_state(self, sched_req_id: str) -> DiffusionRequestState | None:
|
||||
def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
|
||||
"""Remove and return request state if present."""
|
||||
|
||||
@abstractmethod
|
||||
def preempt_request(self, sched_req_id: str) -> bool:
|
||||
def preempt_request(self, request_id: str) -> bool:
|
||||
"""Preempt a running request back to waiting."""
|
||||
|
||||
@abstractmethod
|
||||
def finish_requests(self, sched_req_ids: str | list[str], status: DiffusionRequestStatus) -> None:
|
||||
def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
|
||||
"""Mark one or more requests finished."""
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -26,26 +26,26 @@ class RequestScheduler(_BaseScheduler):
|
||||
return super().schedule()
|
||||
|
||||
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
|
||||
scheduled_req_ids = sched_output.scheduled_req_ids
|
||||
if not scheduled_req_ids:
|
||||
scheduled_request_ids = sched_output.scheduled_request_ids
|
||||
if not scheduled_request_ids:
|
||||
return set()
|
||||
|
||||
terminal_statuses: dict[str, DiffusionRequestStatus] = {}
|
||||
terminal_errors: dict[str, str | None] = {}
|
||||
for sched_req_id in scheduled_req_ids:
|
||||
state = self._request_states.get(sched_req_id)
|
||||
for request_id in scheduled_request_ids:
|
||||
state = self._request_states.get(request_id)
|
||||
if state is None or state.is_finished():
|
||||
continue
|
||||
req_output = output.get_req_output(sched_req_id)
|
||||
req_output = output.get_request_output(request_id)
|
||||
result = req_output.result if req_output is not None else None
|
||||
if result is None:
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[sched_req_id] = "No output result"
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[request_id] = "No output result"
|
||||
elif result.error:
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[sched_req_id] = result.error
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[request_id] = result.error
|
||||
else:
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
terminal_errors[sched_req_id] = None
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
terminal_errors[request_id] = None
|
||||
|
||||
return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
|
||||
@@ -38,84 +38,84 @@ class StepScheduler(_BaseScheduler):
|
||||
self._request_progress.clear()
|
||||
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str:
|
||||
sched_req_id = self._make_sched_req_id(request)
|
||||
request_id = request.request_id
|
||||
total_steps = self._get_total_steps(request)
|
||||
if total_steps <= 0:
|
||||
raise ValueError(f"Diffusion request {sched_req_id} must have positive total_steps, got {total_steps}")
|
||||
raise ValueError(f"Diffusion request {request_id} must have positive total_steps, got {total_steps}")
|
||||
|
||||
current_step = request.sampling_params.step_index or 0
|
||||
if current_step < 0 or current_step >= total_steps:
|
||||
raise ValueError(
|
||||
f"Diffusion request {sched_req_id} has invalid initial step_index {current_step} "
|
||||
f"Diffusion request {request_id} has invalid initial step_index {current_step} "
|
||||
f"for total_steps={total_steps}"
|
||||
)
|
||||
|
||||
request.sampling_params.step_index = current_step
|
||||
sched_req_id = self._add_request_with_sched_req_id(sched_req_id, request)
|
||||
self._request_progress[sched_req_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
|
||||
request_id = self._add_request_with_request_id(request_id, request)
|
||||
self._request_progress[request_id] = _StepProgress(current_step=current_step, total_steps=total_steps)
|
||||
logger.debug(
|
||||
"StepScheduler add_request: %s (step=%d/%d, waiting=%d)",
|
||||
sched_req_id,
|
||||
request_id,
|
||||
current_step,
|
||||
total_steps,
|
||||
len(self._waiting),
|
||||
)
|
||||
return sched_req_id
|
||||
return request_id
|
||||
|
||||
def schedule(self) -> DiffusionSchedulerOutput:
|
||||
return super().schedule()
|
||||
|
||||
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
|
||||
scheduled_req_ids = sched_output.scheduled_req_ids
|
||||
if not scheduled_req_ids:
|
||||
scheduled_request_ids = sched_output.scheduled_request_ids
|
||||
if not scheduled_request_ids:
|
||||
return set()
|
||||
|
||||
terminal_statuses: dict[str, DiffusionRequestStatus] = {}
|
||||
terminal_errors: dict[str, str | None] = {}
|
||||
for sched_req_id in scheduled_req_ids:
|
||||
state = self._request_states.get(sched_req_id)
|
||||
progress = self._request_progress.get(sched_req_id)
|
||||
for request_id in scheduled_request_ids:
|
||||
state = self._request_states.get(request_id)
|
||||
progress = self._request_progress.get(request_id)
|
||||
if state is None or progress is None or state.is_finished():
|
||||
continue
|
||||
req_output = output.get_req_output(sched_req_id)
|
||||
req_output = output.get_request_output(request_id)
|
||||
if req_output is None:
|
||||
logger.warning(
|
||||
"No RunnerOutput for request %s, treating as error",
|
||||
sched_req_id,
|
||||
request_id,
|
||||
)
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[sched_req_id] = "No output for request"
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[request_id] = "No output for request"
|
||||
continue
|
||||
|
||||
req_result = req_output.result
|
||||
output_error = req_result.error if req_result is not None else None
|
||||
if output_error is not None:
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[sched_req_id] = output_error
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[request_id] = output_error
|
||||
continue
|
||||
|
||||
if req_output.step_index is None:
|
||||
logger.warning(
|
||||
"Received RunnerOutput with no step_index for request %s, treating as error",
|
||||
sched_req_id,
|
||||
request_id,
|
||||
)
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[sched_req_id] = "Missing step_index in RunnerOutput"
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR
|
||||
terminal_errors[request_id] = "Missing step_index in RunnerOutput"
|
||||
continue
|
||||
|
||||
# We assume that the decoding stage is executed immediately after the denoising stage completes.
|
||||
progress.current_step = req_output.step_index
|
||||
state.req.sampling_params.step_index = req_output.step_index
|
||||
if req_output.finished:
|
||||
terminal_statuses[sched_req_id] = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
terminal_errors[sched_req_id] = None
|
||||
terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
terminal_errors[request_id] = None
|
||||
else:
|
||||
state.error = None
|
||||
|
||||
return self._finalize_update_from_output(sched_output, terminal_statuses, terminal_errors)
|
||||
|
||||
def _pop_extra_request_state(self, sched_req_id: str) -> None:
|
||||
self._request_progress.pop(sched_req_id, None)
|
||||
def _pop_extra_request_state(self, request_id: str) -> None:
|
||||
self._request_progress.pop(request_id, None)
|
||||
|
||||
def _get_total_steps(self, request: OmniDiffusionRequest) -> int:
|
||||
sampling = request.sampling_params
|
||||
|
||||
@@ -159,7 +159,6 @@ class StageDiffusionProc:
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=[prompt],
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[request_id],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
@@ -189,7 +188,6 @@ class StageDiffusionProc:
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=prompts,
|
||||
sampling_params=sampling_params,
|
||||
request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
|
||||
request_id=request_id,
|
||||
kv_sender_info=kv_sender_info,
|
||||
)
|
||||
|
||||
@@ -322,50 +322,44 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
self, scheduler_output: DiffusionSchedulerOutput
|
||||
) -> tuple[list[DiffusionRequestState], list[str]]:
|
||||
"""Step-before update: cleanup finished requests and get/create one running state."""
|
||||
for req_id in scheduler_output.finished_req_ids:
|
||||
self.state_cache.pop(req_id, None)
|
||||
for request_id in scheduler_output.finished_req_ids:
|
||||
self.state_cache.pop(request_id, None)
|
||||
|
||||
resolved: list[DiffusionRequestState] = []
|
||||
new_req_id: list[str] = []
|
||||
new_request_ids: list[str] = []
|
||||
try:
|
||||
# process new requests
|
||||
for sched_new_req in scheduler_output.scheduled_new_reqs:
|
||||
# new_req_data = scheduler_output.scheduled_new_reqs[0]
|
||||
req_id = sched_new_req.sched_req_id
|
||||
request_id = sched_new_req.request_id
|
||||
req = sched_new_req.req
|
||||
new_req_id.append(req_id)
|
||||
if req_id in self.state_cache:
|
||||
raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
|
||||
request_ids = req.request_ids or [req_id]
|
||||
if len(request_ids) != len(req.prompts):
|
||||
raise ValueError(
|
||||
f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
|
||||
)
|
||||
new_request_ids.append(request_id)
|
||||
if request_id in self.state_cache:
|
||||
raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.")
|
||||
new_state = DiffusionRequestState(
|
||||
req_id=req_id,
|
||||
request_id=request_id,
|
||||
sampling=copy.deepcopy(req.sampling_params),
|
||||
prompts=req.prompts,
|
||||
)
|
||||
self.state_cache[req_id] = new_state
|
||||
self.state_cache[request_id] = new_state
|
||||
resolved.append(new_state)
|
||||
|
||||
# process cached requests
|
||||
for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
|
||||
state = self.state_cache.get(req_id)
|
||||
for request_id in scheduler_output.scheduled_cached_reqs.request_ids:
|
||||
state = self.state_cache.get(request_id)
|
||||
if state is None:
|
||||
raise ValueError(f"Missing cached state for request {req_id}.")
|
||||
raise ValueError(f"Missing cached state for request {request_id}.")
|
||||
resolved.append(state)
|
||||
except Exception:
|
||||
for req_id in new_req_id:
|
||||
self.state_cache.pop(req_id, None)
|
||||
for request_id in new_request_ids:
|
||||
self.state_cache.pop(request_id, None)
|
||||
raise
|
||||
|
||||
return resolved, new_req_id
|
||||
return resolved, new_request_ids
|
||||
|
||||
def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
|
||||
# process new reqs
|
||||
for state in states:
|
||||
if state.req_id in new_request_ids:
|
||||
if state.request_id in new_request_ids:
|
||||
# set generator
|
||||
if state.sampling.generator is None and state.sampling.seed is not None:
|
||||
if state.sampling.generator_device is not None:
|
||||
@@ -407,7 +401,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
|
||||
for state in states:
|
||||
if interrupted or state.denoise_completed:
|
||||
self.state_cache.pop(state.req_id, None)
|
||||
self.state_cache.pop(state.request_id, None)
|
||||
|
||||
def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
|
||||
model_state = getattr(self, "model_state", None)
|
||||
@@ -449,7 +443,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
for state in states:
|
||||
runner_output_list.append(
|
||||
RunnerOutput(
|
||||
req_id=state.req_id,
|
||||
request_id=state.request_id,
|
||||
step_index=state.step_index,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error="stepwise denoise interrupted"),
|
||||
@@ -470,7 +464,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
result = None
|
||||
runner_output_list.append(
|
||||
RunnerOutput(
|
||||
req_id=req.req_id,
|
||||
request_id=req.request_id,
|
||||
step_index=req.step_index,
|
||||
finished=req.denoise_completed,
|
||||
result=result,
|
||||
|
||||
@@ -189,7 +189,7 @@ class DiffusionWorker:
|
||||
self.lora_manager: DiffusionLoRAManager | None = None
|
||||
# Worker-side cache of (lora_request, lora_scale) per scheduled
|
||||
# request id. Used by step mode to recover LoRA identity for cached
|
||||
# requests, which only carry their sched_req_id in subsequent ticks.
|
||||
# requests, which only carry their request_id in subsequent ticks.
|
||||
self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
|
||||
self.stage_id = getattr(od_config, "stage_id", 0)
|
||||
self.init_device()
|
||||
@@ -391,12 +391,12 @@ class DiffusionWorker:
|
||||
homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
|
||||
so any scheduled request id resolves to the active LoRA identity.
|
||||
"""
|
||||
for sched_req_id in scheduler_output.finished_req_ids:
|
||||
self._step_lora_state.pop(sched_req_id, None)
|
||||
for request_id in scheduler_output.finished_req_ids:
|
||||
self._step_lora_state.pop(request_id, None)
|
||||
|
||||
for new_req in scheduler_output.scheduled_new_reqs:
|
||||
sampling = new_req.req.sampling_params
|
||||
self._step_lora_state[new_req.sched_req_id] = (
|
||||
self._step_lora_state[new_req.request_id] = (
|
||||
sampling.lora_request,
|
||||
sampling.lora_scale,
|
||||
)
|
||||
@@ -406,8 +406,8 @@ class DiffusionWorker:
|
||||
|
||||
lora_request: LoRARequest | None = None
|
||||
lora_scale = 1.0
|
||||
for sched_req_id in scheduler_output.scheduled_req_ids:
|
||||
entry = self._step_lora_state.get(sched_req_id)
|
||||
for request_id in scheduler_output.scheduled_request_ids:
|
||||
entry = self._step_lora_state.get(request_id)
|
||||
if entry is not None:
|
||||
lora_request, lora_scale = entry
|
||||
break
|
||||
|
||||
@@ -81,8 +81,8 @@ def _select_states(
|
||||
return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
|
||||
|
||||
|
||||
def _prepare_req_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
|
||||
return [state.req_id for state in states]
|
||||
def _prepare_request_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
|
||||
return [state.request_id for state in states]
|
||||
|
||||
|
||||
def _prepare_prompt_field_on_state(
|
||||
@@ -192,11 +192,13 @@ def _get_request_prompt_seq_lens(
|
||||
mask_attr=mask_attr,
|
||||
)
|
||||
if embeds is None:
|
||||
raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
|
||||
raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
|
||||
|
||||
if mask is not None:
|
||||
if mask.shape[0] != embeds.shape[0]:
|
||||
raise ValueError(f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.req_id}.")
|
||||
raise ValueError(
|
||||
f"{mask_attr} batch dimension does not match {embeds_attr} for request {state.request_id}."
|
||||
)
|
||||
return _get_seq_lens_from_mask(mask)
|
||||
|
||||
seq_lens = getattr(state, seq_lens_attr)
|
||||
@@ -220,7 +222,7 @@ def _prepare_request_prompt_field(
|
||||
mask_attr=mask_attr,
|
||||
)
|
||||
if embeds is None:
|
||||
raise ValueError(f"{embeds_attr} is not initialized on request {state.req_id}.")
|
||||
raise ValueError(f"{embeds_attr} is not initialized on request {state.request_id}.")
|
||||
|
||||
actual_seq_lens = _get_request_prompt_seq_lens(
|
||||
state,
|
||||
@@ -248,7 +250,7 @@ def _prepare_request_prompt_field(
|
||||
if current_seq_len > target_seq_len:
|
||||
if max_actual_seq_len > target_seq_len:
|
||||
raise ValueError(
|
||||
f"{embeds_attr} for request {state.req_id} requires seq_len "
|
||||
f"{embeds_attr} for request {state.request_id} requires seq_len "
|
||||
f"{max_actual_seq_len}, got target {target_seq_len}."
|
||||
)
|
||||
return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len]
|
||||
@@ -349,7 +351,7 @@ def _require_state_latents(
|
||||
) -> torch.Tensor:
|
||||
latents = state.latents
|
||||
if latents is None:
|
||||
raise ValueError(f"Request {state.req_id} has no latents while preparing {for_field}.")
|
||||
raise ValueError(f"Request {state.request_id} has no latents while preparing {for_field}.")
|
||||
return latents
|
||||
|
||||
|
||||
@@ -526,12 +528,12 @@ def _prepare_negative_prompt_embeds(
|
||||
|
||||
def _same_composition(
|
||||
cached_batch: InputBatch | None,
|
||||
req_ids: list[str],
|
||||
request_ids: list[str],
|
||||
idx_mapping_np: np.ndarray,
|
||||
) -> bool:
|
||||
if cached_batch is None:
|
||||
return False
|
||||
if cached_batch.req_ids != req_ids:
|
||||
if cached_batch.request_ids != request_ids:
|
||||
return False
|
||||
return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
|
||||
|
||||
@@ -581,7 +583,7 @@ class InputBatch:
|
||||
unchanged.
|
||||
"""
|
||||
|
||||
req_ids: list[str]
|
||||
request_ids: list[str]
|
||||
num_reqs: int
|
||||
num_reqs_after_padding: int
|
||||
idx_mapping: torch.Tensor
|
||||
@@ -604,9 +606,9 @@ class InputBatch:
|
||||
negative_txt_seq_lens: list[int] | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if len(self.req_ids) != int(self.idx_mapping.numel()):
|
||||
raise ValueError("`req_ids` and `idx_mapping` must have the same length.")
|
||||
if self.num_reqs != len(self.req_ids):
|
||||
if len(self.request_ids) != int(self.idx_mapping.numel()):
|
||||
raise ValueError("`request_ids` and `idx_mapping` must have the same length.")
|
||||
if self.num_reqs != len(self.request_ids):
|
||||
raise ValueError("`num_reqs` must match the number of request ids.")
|
||||
if self.num_reqs_after_padding < self.num_reqs:
|
||||
raise ValueError("`num_reqs_after_padding` must be >= `num_reqs`.")
|
||||
@@ -653,11 +655,11 @@ class InputBatch:
|
||||
selected_states: Sequence[DiffusionRequestState],
|
||||
idx_mapping: torch.Tensor,
|
||||
idx_mapping_np: np.ndarray,
|
||||
req_ids: list[str],
|
||||
request_ids: list[str],
|
||||
) -> InputBatch:
|
||||
self.req_ids = req_ids
|
||||
self.num_reqs = len(req_ids)
|
||||
self.num_reqs_after_padding = len(req_ids)
|
||||
self.request_ids = request_ids
|
||||
self.num_reqs = len(request_ids)
|
||||
self.num_reqs_after_padding = len(request_ids)
|
||||
self.idx_mapping = idx_mapping
|
||||
self.idx_mapping_np = idx_mapping_np
|
||||
self.latents = _prepare_latents(selected_states, out=self.latents)
|
||||
@@ -675,9 +677,9 @@ class InputBatch:
|
||||
) -> InputBatch:
|
||||
"""Build a temporary step-local batch view from request states."""
|
||||
selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
|
||||
req_ids = _prepare_req_ids(selected_states)
|
||||
request_ids = _prepare_request_ids(selected_states)
|
||||
|
||||
if _same_composition(cached_batch, req_ids, idx_mapping_np):
|
||||
if _same_composition(cached_batch, request_ids, idx_mapping_np):
|
||||
assert cached_batch is not None
|
||||
cached_batch._repack_dynamic_fields(selected_states)
|
||||
return cached_batch
|
||||
@@ -687,14 +689,14 @@ class InputBatch:
|
||||
selected_states,
|
||||
idx_mapping,
|
||||
idx_mapping_np,
|
||||
req_ids,
|
||||
request_ids,
|
||||
)
|
||||
|
||||
prompt_embeds, prompt_embeds_mask = _prepare_prompt_embeds(selected_states)
|
||||
negative_prompt_embeds, negative_prompt_embeds_mask = _prepare_negative_prompt_embeds(selected_states)
|
||||
do_true_cfg, true_cfg_scale, cfg_normalize = _prepare_cfg_scalars(selected_states)
|
||||
return cls(
|
||||
req_ids=req_ids,
|
||||
request_ids=request_ids,
|
||||
num_reqs=len(selected_states),
|
||||
num_reqs_after_padding=len(selected_states),
|
||||
idx_mapping=idx_mapping,
|
||||
|
||||
@@ -38,7 +38,7 @@ class DiffusionRequestState:
|
||||
"""
|
||||
|
||||
# ── Identity / request-level inputs ──
|
||||
req_id: str
|
||||
request_id: str
|
||||
sampling: OmniDiffusionSamplingParams
|
||||
prompts: list[OmniPromptType] | None = None
|
||||
|
||||
@@ -112,7 +112,7 @@ class DiffusionRequestState:
|
||||
|
||||
class BaseRunnerOutput(ABC):
|
||||
@abstractmethod
|
||||
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
|
||||
def get_request_output(self, request_id: str) -> RunnerOutput | None:
|
||||
pass
|
||||
|
||||
|
||||
@@ -125,13 +125,13 @@ class RunnerOutput(BaseRunnerOutput):
|
||||
_request_state_cache.
|
||||
"""
|
||||
|
||||
req_id: str
|
||||
request_id: str
|
||||
step_index: int | None = None
|
||||
finished: bool = False
|
||||
result: DiffusionOutput | None = None
|
||||
|
||||
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
|
||||
return self if self.req_id == sched_req_id else None
|
||||
def get_request_output(self, request_id: str) -> RunnerOutput | None:
|
||||
return self if self.request_id == request_id else None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -140,18 +140,18 @@ class BatchRunnerOutput(BaseRunnerOutput):
|
||||
_id_to_idx: dict[str, int] = field(init=False, repr=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._id_to_idx = {out.req_id: i for i, out in enumerate(self.runner_outputs)}
|
||||
self._id_to_idx = {out.request_id: i for i, out in enumerate(self.runner_outputs)}
|
||||
|
||||
def __getitem__(self, req_id: str) -> RunnerOutput | None:
|
||||
"""access single RunnerOutput by req_id"""
|
||||
idx = self._id_to_idx.get(req_id)
|
||||
def __getitem__(self, request_id: str) -> RunnerOutput | None:
|
||||
"""access single RunnerOutput by request_id"""
|
||||
idx = self._id_to_idx.get(request_id)
|
||||
return self.runner_outputs[idx] if idx is not None else None
|
||||
|
||||
def get_req_output(self, sched_req_id: str) -> RunnerOutput | None:
|
||||
return self[sched_req_id]
|
||||
def get_request_output(self, request_id: str) -> RunnerOutput | None:
|
||||
return self[request_id]
|
||||
|
||||
@property
|
||||
def req_ids(self) -> list[str]:
|
||||
def request_ids(self) -> list[str]:
|
||||
return list(self._id_to_idx.keys())
|
||||
|
||||
def __len__(self) -> int:
|
||||
|
||||
@@ -12,6 +12,8 @@ from typing import Any
|
||||
import torch
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
from .factory import OmniConnectorFactory
|
||||
from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec
|
||||
from .utils.env import expand_env_int
|
||||
@@ -1014,7 +1016,7 @@ class OmniKVTransferManager:
|
||||
return None, 0
|
||||
|
||||
# Skip during warmup dummy run — no sender is available.
|
||||
if request_id == "dummy_req_id":
|
||||
if OmniDiffusionRequest.is_dummy_run_request_id(request_id):
|
||||
logger.info("Skip receiving KV cache for dummy warmup request")
|
||||
return None, 0
|
||||
|
||||
@@ -1171,12 +1173,7 @@ class OmniKVTransferManager:
|
||||
@staticmethod
|
||||
def _resolve_request_id(req: Any) -> str | None:
|
||||
"""Resolve the logical request ID used for KV transfer lookups."""
|
||||
request_id = getattr(req, "request_id", None)
|
||||
if request_id:
|
||||
return request_id
|
||||
if hasattr(req, "request_ids") and req.request_ids:
|
||||
return req.request_ids[0]
|
||||
return None
|
||||
return getattr(req, "request_id", None)
|
||||
|
||||
# Legacy compatibility method
|
||||
def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
|
||||
|
||||
@@ -1245,9 +1245,7 @@ class OmniConnectorModelRunnerMixin:
|
||||
if self._kv_transfer_manager is None:
|
||||
return False
|
||||
|
||||
request_id = getattr(req, "request_id", None) or (
|
||||
req.request_ids[0] if hasattr(req, "request_ids") and req.request_ids else None
|
||||
)
|
||||
request_id = getattr(req, "request_id", None)
|
||||
if not request_id:
|
||||
logger.warning("Request has no ID, cannot receive KV cache")
|
||||
return False
|
||||
|
||||
Reference in new issue
Block a user