diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 58e885d46..9f10aea0c 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -700,3 +700,232 @@ def test_wire_round_trip_struct_to_dict_contract(): assert set(decoded.keys()) == set(expected.keys()) assert set(decoded["meta"].keys()) == set(expected["meta"].keys()) assert set(decoded["codes"].keys()) == set(expected["codes"].keys()) + + +# --------------------------------------------------------------- +# Deferred finish for upstream-completed requests +# --------------------------------------------------------------- + + +def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs): + """Build a mock scheduler with requests queued for deferred finish.""" + adapter_mock = mocker.MagicMock() + adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs} + cleanup_calls = [] + adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw)) + + scheduler = mocker.MagicMock() + scheduler.chunk_transfer_adapter = adapter_mock + scheduler.connector = None + scheduler.ec_connector = None + scheduler.perf_metrics = None + scheduler.log_stats = False + scheduler.recompute_kv_load_failures = False + scheduler.structured_output_manager = mocker.MagicMock() + scheduler.structured_output_manager.should_advance.return_value = False + scheduler.finished_req_ids_dict = {} + scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_event_publisher = mocker.MagicMock() + scheduler._pending_finish_reqs = list(pending_finish_reqs) + + scheduler._handle_stopped_request = mocker.MagicMock(return_value=True) + scheduler._free_request = mocker.MagicMock(return_value=None) + scheduler._get_routed_experts = mocker.MagicMock(return_value=None) + scheduler.running = list(running) + scheduler.waiting = mocker.MagicMock() + scheduler.waiting.remove_requests = mocker.MagicMock() + scheduler.make_stats = mocker.MagicMock(return_value=None) + scheduler.requests = {r.request_id: r for r in running} + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + model_runner_output = SimpleNamespace( + sampled_token_ids=None, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=None, + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={}, + ) + return scheduler, scheduler_output, model_runner_output, cleanup_calls + + +def test_deferred_finish_emits_finished_output(mocker: MockerFixture): + """A request whose upstream completed with no remaining tokens should + emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state.""" + from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + + request = _HashableRequest( + request_id="req-df1", + external_req_id="ext-df1", + status=RequestStatus.RUNNING, + is_finished=lambda: False, + num_computed_tokens=16, + num_prompt_tokens=16, + prompt_token_ids=list(range(1, 17)), + num_output_placeholders=0, + sampling_params=None, + pooling_params=None, + stop_reason=None, + client_index=0, + take_events=lambda: [], + trace_headers=None, + has_encoder_inputs=False, + take_prefill_stats=lambda: None, + num_nans_in_logits=0, + get_finished_reason=lambda: "stop", + ) + scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler( + mocker, + running=[request], + pending_finish_reqs=[request], + ) + scheduler._free_request.return_value = {"mock": "kv_params"} + + result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out) + + assert request.status == RequestStatus.FINISHED_STOPPED + scheduler._handle_stopped_request.assert_called_once_with(request) + scheduler._free_request.assert_called_once_with(request) + assert len(cleanup_calls) == 1 + assert cleanup_calls[0][0] == ("req-df1", "ext-df1") + + eco = result[0] + assert len(eco.outputs) == 1 + assert eco.outputs[0].request_id == "req-df1" + assert eco.outputs[0].finish_reason == "stop" + assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"} + assert scheduler._pending_finish_reqs == [] + + +def test_deferred_finish_empty_prompt(mocker: MockerFixture): + """A request that never received any tokens (finished immediately upstream) + should still emit a FINISHED output and clean up.""" + from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + + request = _HashableRequest( + request_id="req-df2", + external_req_id="ext-df2", + status=RequestStatus.WAITING, + is_finished=lambda: False, + num_computed_tokens=0, + num_prompt_tokens=0, + prompt_token_ids=[], + num_output_placeholders=0, + sampling_params=None, + pooling_params=None, + stop_reason=None, + client_index=0, + take_events=lambda: [], + trace_headers=None, + has_encoder_inputs=False, + take_prefill_stats=lambda: None, + num_nans_in_logits=0, + get_finished_reason=lambda: "stop", + ) + scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler( + mocker, + running=[], + pending_finish_reqs=[request], + ) + + result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out) + + assert request.status == RequestStatus.FINISHED_STOPPED + scheduler._free_request.assert_called_once_with(request) + assert len(cleanup_calls) == 1 + eco = result[0] + assert len(eco.outputs) == 1 + assert eco.outputs[0].finish_reason == "stop" + assert scheduler._pending_finish_reqs == [] + + +def test_deferred_finish_skips_already_finished(mocker: MockerFixture): + """A request that was aborted between schedule() and update_from_output() + should be skipped without emitting output or freeing resources twice.""" + from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + + request = _HashableRequest( + request_id="req-df3", + external_req_id="ext-df3", + status=RequestStatus.FINISHED_ABORTED, + is_finished=lambda: True, + num_computed_tokens=0, + num_prompt_tokens=0, + prompt_token_ids=[], + num_output_placeholders=0, + sampling_params=None, + pooling_params=None, + stop_reason=None, + client_index=0, + take_events=lambda: [], + trace_headers=None, + has_encoder_inputs=False, + take_prefill_stats=lambda: None, + num_nans_in_logits=0, + get_finished_reason=lambda: "stop", + ) + scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler( + mocker, + running=[], + pending_finish_reqs=[request], + ) + + result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out) + + scheduler._handle_stopped_request.assert_not_called() + scheduler._free_request.assert_not_called() + assert len(cleanup_calls) == 0 + assert 0 not in result + assert scheduler._pending_finish_reqs == [] + + +def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture): + """When _handle_stopped_request returns False (resumable request), the + output must still be emitted so the client stream doesn't hang.""" + from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + + request = _HashableRequest( + request_id="req-df4", + external_req_id="ext-df4", + status=RequestStatus.RUNNING, + is_finished=lambda: False, + num_computed_tokens=16, + num_prompt_tokens=16, + prompt_token_ids=list(range(1, 17)), + num_output_placeholders=0, + sampling_params=None, + pooling_params=None, + stop_reason=None, + client_index=0, + take_events=lambda: [], + trace_headers=None, + has_encoder_inputs=False, + take_prefill_stats=lambda: None, + num_nans_in_logits=0, + get_finished_reason=lambda: "stop", + ) + scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler( + mocker, + running=[request], + pending_finish_reqs=[request], + ) + scheduler._handle_stopped_request.return_value = False + + result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out) + + scheduler._handle_stopped_request.assert_called_once_with(request) + scheduler._free_request.assert_not_called() + assert len(cleanup_calls) == 0 + + eco = result[0] + assert len(eco.outputs) == 1 + assert eco.outputs[0].request_id == "req-df4" + assert eco.outputs[0].finish_reason == "stop" + assert eco.outputs[0].kv_transfer_params is None + assert scheduler._pending_finish_reqs == [] diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index e290d3ef6..1218d416b 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -44,6 +44,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + self._pending_finish_reqs: list[Request] = [] self.input_coordinator: OmniSchedulingCoordinator | None = None if uses_qwen3_omni_full_payload_input_coordinator(model_config): self.input_coordinator = OmniSchedulingCoordinator( @@ -111,23 +112,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): num_computed_tokens = request.num_computed_tokens required_tokens = len(request.prompt_token_ids) - num_computed_tokens - # async_chunk: don't schedule placeholder tokens when no new chunk is available. if required_tokens <= 0: if ( self.chunk_transfer_adapter is not None and request.request_id in self.chunk_transfer_adapter.finished_requests ): - # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED. - if len(request.prompt_token_ids) <= num_computed_tokens: - request.prompt_token_ids.append(0) - try: - request._all_token_ids.append(0) # type: ignore[attr-defined] - except Exception: - pass - required_tokens = len(request.prompt_token_ids) - num_computed_tokens - else: - req_index += 1 - continue + self._pending_finish_reqs.append(request) + req_index += 1 + continue num_new_tokens = min(required_tokens, token_budget) new_blocks = self.kv_cache_manager.allocate_slots( request, @@ -174,11 +166,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # async_chunk: wait for the first upstream chunk (don't start with placeholders). if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0: if request.request_id in self.chunk_transfer_adapter.finished_requests: - request.prompt_token_ids.append(0) - try: - request._all_token_ids.append(0) # type: ignore[attr-defined] - except Exception: - pass + self.waiting.pop_request() + self._pending_finish_reqs.append(request) + continue else: self.waiting.pop_request() skipped_waiting_requests.prepend_request(request) @@ -581,6 +571,36 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors + # Finish async_chunk requests that schedule() collected because their + # upstream completed with no remaining codec tokens. + for request in self._pending_finish_reqs: + if request.is_finished(): + continue + request.status = RequestStatus.FINISHED_STOPPED + finish_reason = request.get_finished_reason() + finished = self._handle_stopped_request(request) + kv_transfer_params = None + if finished: + kv_transfer_params = self._free_request(request) + if self.chunk_transfer_adapter is not None: + self.chunk_transfer_adapter.cleanup( + request.request_id, + getattr(request, "external_req_id", None), + ) + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=finish_reason, + stop_reason=request.stop_reason, + events=request.take_events(), + kv_transfer_params=kv_transfer_params, + trace_headers=request.trace_headers, + ) + ) + stopped_running_reqs.add(request) + self._pending_finish_reqs.clear() + # Remove the stopped requests from the running and waiting queues. if stopped_running_reqs: self.running = remove_all(self.running, stopped_running_reqs)