Skip to content
Open
Changes from 1 commit
Commits
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
30 changes: 16 additions & 14 deletions vllm/v1/engine/async_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -1018,15 +1018,11 @@ async def wait_for_requests_to_drain(self, drain_timeout: int = 300):
)

async def _drain_requests_for_elastic_ep(self, drain_timeout: int) -> None:
try:
logger.info(
"VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
"waiting for requests to drain before scaling"
)
await self.wait_for_requests_to_drain(drain_timeout)
except BaseException:
set_scaling_elastic_ep(False)
raise
logger.info(
"VLLM_ELASTIC_EP_DRAIN_REQUESTS is set, "
"waiting for requests to drain before scaling"
)
await self.wait_for_requests_to_drain(drain_timeout)

async def scale_elastic_ep(
self, new_data_parallel_size: int, drain_timeout: int = 300
Expand Down Expand Up @@ -1065,13 +1061,19 @@ async def _scale_elastic_ep(
self._logger_ref[0] = self.logger_manager
self.logger_manager.log_engine_initialized()

# ScalingMiddleware 503s every request while this is set, including the
# /scale_elastic_ep call that would clear it, so never leave it latched.
set_scaling_elastic_ep(True)
if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
await self._drain_requests_for_elastic_ep(drain_timeout)
try:
if envs.VLLM_ELASTIC_EP_DRAIN_REQUESTS:
await self._drain_requests_for_elastic_ep(drain_timeout)

await self.engine_core.commit_elastic_ep()
self.vllm_config.parallel_config.data_parallel_size = new_data_parallel_size
set_scaling_elastic_ep(False)
await self.engine_core.commit_elastic_ep()
self.vllm_config.parallel_config.data_parallel_size = (
new_data_parallel_size
)
finally:
set_scaling_elastic_ep(False)

async def handle_fault(
self, fault_tolerance_request: FaultToleranceRequest
Expand Down
Loading