[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
@@ -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"]
|
||||
|
||||
Reference in new issue
Block a user