[Core] Async scheduling + structured outputs compatibility (#26866)
Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
@@ -6,6 +6,9 @@ from copy import deepcopy
|
||||
|
||||
from tblib import pickling_support
|
||||
|
||||
# Import fixture
|
||||
from tests.v1.entrypoints.conftest import sample_json_schema # noqa
|
||||
|
||||
# ruff: noqa
|
||||
|
||||
# Install support for pickling exceptions so that we can nicely propagate
|
||||
|
||||
@@ -337,8 +337,6 @@ def test_stop_via_update_from_output():
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_output = ModelRunnerOutput(
|
||||
@@ -385,8 +383,6 @@ def test_stop_via_update_from_output():
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_output = ModelRunnerOutput(
|
||||
@@ -431,8 +427,6 @@ def test_stop_via_update_from_output():
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_output = ModelRunnerOutput(
|
||||
@@ -472,8 +466,6 @@ def test_stop_via_update_from_output():
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_output = ModelRunnerOutput(
|
||||
@@ -1988,7 +1980,6 @@ def test_schedule_skip_tokenizer_init():
|
||||
scheduler.add_request(request)
|
||||
output = scheduler.schedule()
|
||||
assert len(output.scheduled_new_reqs) == len(requests)
|
||||
assert output.grammar_bitmask is None
|
||||
|
||||
|
||||
def test_schedule_skip_tokenizer_init_structured_output_request():
|
||||
|
||||
@@ -7,6 +7,7 @@ import torch._dynamo.config as dynamo_config
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
|
||||
from ...conftest import VllmRunner
|
||||
from ...models.utils import check_outputs_equal
|
||||
@@ -15,9 +16,12 @@ MODEL = "Qwen/Qwen3-0.6B"
|
||||
|
||||
|
||||
@dynamo_config.patch(cache_size_limit=16)
|
||||
def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_preempt_and_async_scheduling_e2e(
|
||||
sample_json_schema, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
"""Test consistency of combos of async scheduling, preemption,
|
||||
uni/multiproc executor, and various sampling parameters."""
|
||||
uni/multiproc executor, and various sampling parameters
|
||||
including structured outputs."""
|
||||
|
||||
first_prompt = (
|
||||
"The following numbers of the sequence "
|
||||
@@ -35,6 +39,12 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
|
||||
dict(bad_words=["the", " the"]),
|
||||
dict(logprobs=2),
|
||||
dict(logprobs=2, presence_penalty=-1.0),
|
||||
dict(structured_outputs=StructuredOutputsParams(json=sample_json_schema)),
|
||||
dict(
|
||||
structured_outputs=StructuredOutputsParams(json=sample_json_schema),
|
||||
logprobs=2,
|
||||
presence_penalty=-1.0,
|
||||
),
|
||||
]
|
||||
|
||||
default_params = dict(
|
||||
@@ -248,7 +248,7 @@ def test_engine_core_concurrent_batches():
|
||||
self,
|
||||
scheduler_output,
|
||||
non_block=False,
|
||||
) -> Future[ModelRunnerOutput]:
|
||||
) -> Future[ModelRunnerOutput | None]:
|
||||
"""Make execute_model non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
@@ -263,6 +263,23 @@ def test_engine_core_concurrent_batches():
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
def sample_tokens(
|
||||
self, grammar_output, non_block=False
|
||||
) -> Future[ModelRunnerOutput]:
|
||||
"""Make sample_tokens non-blocking."""
|
||||
|
||||
# DummyExecutor used only for testing async case.
|
||||
assert non_block
|
||||
|
||||
def _execute():
|
||||
output = self.collective_rpc("sample_tokens", args=(grammar_output,))
|
||||
# Make a copy because output[0] may be reused
|
||||
# by the next batch.
|
||||
return copy.deepcopy(output[0])
|
||||
|
||||
# Use the thread pool instead of creating a new thread
|
||||
return self.thread_pool.submit(_execute)
|
||||
|
||||
@property
|
||||
def max_concurrent_batches(self) -> int:
|
||||
return 2
|
||||
|
||||
@@ -31,7 +31,9 @@ class CustomMultiprocExecutor(MultiprocExecutor):
|
||||
# Drop marker to show that this was run
|
||||
with open(".marker", "w"):
|
||||
...
|
||||
return super().collective_rpc(method, timeout, args, kwargs)
|
||||
return super().collective_rpc(
|
||||
method, timeout, args, kwargs, non_block, unique_reply_rank
|
||||
)
|
||||
|
||||
|
||||
CustomMultiprocExecutorAsync = CustomMultiprocExecutor
|
||||
|
||||
@@ -26,8 +26,6 @@ def _make_empty_scheduler_output():
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
kv_connector_metadata=SharedStorageConnectorMetadata(),
|
||||
)
|
||||
|
||||
|
||||
@@ -981,9 +981,7 @@ def test_scheduler_kv_connector_stats_aggregation():
|
||||
scheduled_encoder_inputs={},
|
||||
num_common_prefix_blocks=[0],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=set(),
|
||||
structured_output_request_ids={},
|
||||
grammar_bitmask=None,
|
||||
free_encoder_mm_hashes=[],
|
||||
)
|
||||
|
||||
engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output)
|
||||
|
||||
@@ -92,8 +92,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -171,8 +169,6 @@ def test_update_states_request_finished(model_runner):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids={req_id},
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
@@ -201,8 +197,6 @@ def test_update_states_request_resumed(model_runner):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
@@ -230,8 +224,6 @@ def test_update_states_request_resumed(model_runner):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
@@ -261,8 +253,6 @@ def test_update_states_no_changes(model_runner):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
@@ -296,8 +286,6 @@ def test_update_states_request_unscheduled(model_runner):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
|
||||
@@ -152,8 +152,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -269,8 +267,6 @@ def test_update_states_request_finished(model_runner, dist_init):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids={req_id},
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
metadata_before = model_runner.input_batch.sampling_metadata
|
||||
@@ -301,8 +297,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
model_runner._update_states(scheduler_output)
|
||||
@@ -330,8 +324,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
metadata_before = model_runner.input_batch.sampling_metadata
|
||||
@@ -423,8 +415,6 @@ def test_update_states_no_changes(model_runner, dist_init):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
metadata_before = model_runner.input_batch.sampling_metadata
|
||||
@@ -460,8 +450,6 @@ def test_update_states_request_unscheduled(model_runner, dist_init):
|
||||
num_common_prefix_blocks=[],
|
||||
finished_req_ids=set(),
|
||||
free_encoder_mm_hashes=[],
|
||||
structured_output_request_ids=[],
|
||||
grammar_bitmask=None,
|
||||
)
|
||||
|
||||
metadata_before = model_runner._update_states(scheduler_output)
|
||||
|
||||
Reference in New Issue
Block a user