Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Core] Priority-based scheduling in async engine #8850

Merged
merged 8 commits into from
Sep 27, 2024
Merged

[Core] Priority-based scheduling in async engine #8850

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
09712dc
Priority in async engine
schoennenbeck Sep 26, 2024
ed32cb9
Formatting
schoennenbeck Sep 26, 2024
57c7b4e
More formatting
schoennenbeck Sep 26, 2024
feb30b9
Rollback of unnecessary formatting
schoennenbeck Sep 26, 2024
c9adc62
Add priority to generate and add_request
schoennenbeck Sep 26, 2024
2b68c58
Fixed error handling
schoennenbeck Sep 27, 2024
ef536f6
Merge branch 'main' into feature/priority-in-async-engine
schoennenbeck Sep 27, 2024
2e3e185
Fixed typo
schoennenbeck Sep 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 23 additions & 2 deletions vllm/engine/async_llm_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,7 @@ async def add_request_async(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
) -> None:
...

Expand All @@ -426,6 +427,7 @@ async def add_request_async(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
) -> None:
...

Expand All @@ -442,6 +444,7 @@ async def add_request_async(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
*,
inputs: Optional[PromptType] = None, # DEPRECATED
) -> None:
Expand All @@ -453,6 +456,9 @@ async def add_request_async(
if lora_request is not None and not self.lora_config:
raise ValueError(f"Got lora_request {lora_request} but LoRA is "
"not enabled!")
if priority != 0 and not self.scheduler_config.policy == "priority":
raise ValueError(f"Got priority {priority} but "
"Priority scheduling is not enabled.")
if arrival_time is None:
arrival_time = time.time()

Expand All @@ -472,6 +478,7 @@ async def add_request_async(
lora_request=lora_request,
prompt_adapter_request=prompt_adapter_request,
trace_headers=trace_headers,
priority=priority,
Copy link
Contributor

@sfc-gh-zhwang sfc-gh-zhwang Sep 27, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i don't think we are passing policy="priority" anywhere, shall we add it to here so that user can actually leverage this. Otherwise, how are folks supposed to use this feature? @njhill

not self.scheduler_config.policy == "priority":

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think currently the only way of actually using the feature is to manually change the policy after creating the engine.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, but shall we allow folks to specify through args? Like other param in the scheduler_config? I feel that's better for wider adoption for this feature.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I misread your comment. Yes I think this needs to go in the EngineArgs or some similar place. However, I think that should be a separate PR.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it's in SchedulerConfig but still needs to be wired in the EngineArgs to be able to enable it externally. Agree that could be done in a separate PR

)

async def check_health_async(self) -> None:
Expand Down Expand Up @@ -822,6 +829,7 @@ def add_request(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
) -> Coroutine[None, None, AsyncGenerator[Union[
RequestOutput, EmbeddingRequestOutput], None]]:
...
Expand All @@ -836,6 +844,7 @@ def add_request(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
) -> Coroutine[None, None, AsyncGenerator[Union[
RequestOutput, EmbeddingRequestOutput], None]]:
...
Expand All @@ -853,6 +862,7 @@ async def add_request(
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
*,
inputs: Optional[PromptType] = None, # DEPRECATED
) -> AsyncGenerator[Union[RequestOutput, EmbeddingRequestOutput], None]:
Expand All @@ -870,6 +880,11 @@ async def add_request(
"error that caused the background loop to stop "
"(AsyncEngineDeadError).")

if (priority != 0
and not self.engine.scheduler_config.policy == "priority"):
raise ValueError(f"Got priority {priority} but "
"Priority scheduling is not enabled.")

stream = self._request_tracker.add_request(
request_id,
verbose=self.log_requests,
Expand All @@ -878,7 +893,9 @@ async def add_request(
arrival_time=arrival_time or time.time(),
lora_request=lora_request,
trace_headers=trace_headers,
prompt_adapter_request=prompt_adapter_request)
prompt_adapter_request=prompt_adapter_request,
priority=priority,
)

return stream.generator()

Expand All @@ -889,7 +906,8 @@ async def generate(
request_id: str,
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
priority: int = 0,
) -> AsyncGenerator[RequestOutput, None]:
"""Generate outputs for a request.

Expand All @@ -906,6 +924,8 @@ async def generate(
trace_headers: OpenTelemetry trace headers.
prompt_adapter_request: Prompt Adapter request to use
for generation, if any.
priority: The priority of the request.
Only applicable with priority scheduling.

Yields:
The output `RequestOutput` objects from the LLMEngine
Expand Down Expand Up @@ -961,6 +981,7 @@ async def generate(
lora_request=lora_request,
trace_headers=trace_headers,
prompt_adapter_request=prompt_adapter_request,
priority=priority,
):
yield LLMEngine.validate_output(output, RequestOutput)

Expand Down
2 changes: 1 addition & 1 deletion vllm/engine/llm_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -786,7 +786,7 @@ def add_request(
raise ValueError(f"Got lora_request {lora_request} but LoRA is "
"not enabled!")

if priority > 0 and not self.scheduler_config.policy == "priority":
if priority != 0 and not self.scheduler_config.policy == "priority":
raise ValueError(f"Got priority {priority} but "
"Priority scheduling is not enabled.")

Expand Down
Loading