987 lines
38 KiB
Python
987 lines
38 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import copy
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from collections import deque
|
|
from contextlib import nullcontext
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Deque,
|
|
Dict,
|
|
Generic,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
TypeVar,
|
|
)
|
|
|
|
import fastapi
|
|
import zmq
|
|
|
|
from sglang.srt.managers.io_struct import (
|
|
AttachHiCacheStorageReqInput,
|
|
AttachHiCacheStorageReqOutput,
|
|
CheckWeightsReqInput,
|
|
CheckWeightsReqOutput,
|
|
ClearHiCacheReqInput,
|
|
ClearHiCacheReqOutput,
|
|
CloseSessionReqInput,
|
|
DestroyWeightsUpdateGroupReqInput,
|
|
DestroyWeightsUpdateGroupReqOutput,
|
|
DetachHiCacheStorageReqInput,
|
|
DetachHiCacheStorageReqOutput,
|
|
DumperControlReqInput,
|
|
DumperControlReqOutput,
|
|
ExpertDistributionReq,
|
|
ExpertDistributionReqOutput,
|
|
ExpertDistributionReqType,
|
|
FlushCacheReqInput,
|
|
FlushCacheReqOutput,
|
|
GetInternalStateReq,
|
|
GetInternalStateReqOutput,
|
|
GetLoadReqInput,
|
|
GetLoadReqOutput,
|
|
GetLoadsReqInput,
|
|
GetLoadsReqOutput,
|
|
GetWeightsByNameReqInput,
|
|
GetWeightsByNameReqOutput,
|
|
InitWeightsSendGroupForRemoteInstanceReqInput,
|
|
InitWeightsSendGroupForRemoteInstanceReqOutput,
|
|
InitWeightsUpdateGroupReqInput,
|
|
InitWeightsUpdateGroupReqOutput,
|
|
LoadLoRAAdapterFromTensorsReqInput,
|
|
LoadLoRAAdapterFromTensorsReqOutput,
|
|
LoadLoRAAdapterReqInput,
|
|
LoadLoRAAdapterReqOutput,
|
|
LoRAUpdateOutput,
|
|
OpenSessionReqInput,
|
|
PinPrefixReqInput,
|
|
PinPrefixReqOutput,
|
|
ProfileReq,
|
|
ProfileReqOutput,
|
|
ProfileReqType,
|
|
ReleaseMemoryOccupationReqInput,
|
|
ReleaseMemoryOccupationReqOutput,
|
|
ResumeMemoryOccupationReqInput,
|
|
ResumeMemoryOccupationReqOutput,
|
|
SendWeightsToRemoteInstanceReqInput,
|
|
SendWeightsToRemoteInstanceReqOutput,
|
|
SetInternalStateReq,
|
|
SetInternalStateReqOutput,
|
|
SlowDownReqInput,
|
|
SlowDownReqOutput,
|
|
UnloadLoRAAdapterReqInput,
|
|
UnloadLoRAAdapterReqOutput,
|
|
UpdateWeightsFromDistributedReqInput,
|
|
UpdateWeightsFromDistributedReqOutput,
|
|
UpdateWeightsFromIPCReqInput,
|
|
UpdateWeightsFromIPCReqOutput,
|
|
UpdateWeightsFromTensorReqInput,
|
|
UpdateWeightsFromTensorReqOutput,
|
|
)
|
|
from sglang.srt.server_args import LoRARef, ServerArgs
|
|
from sglang.srt.utils import get_bool_env_var
|
|
from sglang.utils import TypeBasedDispatcher
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
|
|
|
T = TypeVar("T")
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _Communicator(Generic[T]):
|
|
"""Note: The communicator now only run up to 1 in-flight request at any time."""
|
|
|
|
def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"):
|
|
self._sender = sender
|
|
self._fan_out = fan_out
|
|
self._mode = mode
|
|
self._result_event: Optional[asyncio.Event] = None
|
|
self._result_values: Optional[List[T]] = None
|
|
self._ready_queue: Deque[asyncio.Future] = deque()
|
|
|
|
assert mode in ["queueing", "watching"]
|
|
|
|
async def queueing_call(self, obj: T):
|
|
ready_event = asyncio.Event()
|
|
if self._result_event is not None or len(self._ready_queue) > 0:
|
|
self._ready_queue.append(ready_event)
|
|
await ready_event.wait()
|
|
assert self._result_event is None
|
|
assert self._result_values is None
|
|
|
|
if obj:
|
|
self._sender.send_pyobj(obj)
|
|
|
|
self._result_event = asyncio.Event()
|
|
self._result_values = []
|
|
await self._result_event.wait()
|
|
result_values = self._result_values
|
|
self._result_event = self._result_values = None
|
|
|
|
if len(self._ready_queue) > 0:
|
|
self._ready_queue.popleft().set()
|
|
|
|
return result_values
|
|
|
|
async def watching_call(self, obj):
|
|
if self._result_event is None:
|
|
assert self._result_values is None
|
|
self._result_values = []
|
|
self._result_event = asyncio.Event()
|
|
|
|
if obj:
|
|
self._sender.send_pyobj(obj)
|
|
|
|
await self._result_event.wait()
|
|
result_values = copy.deepcopy(self._result_values)
|
|
self._result_event = self._result_values = None
|
|
return result_values
|
|
|
|
async def __call__(self, obj):
|
|
if self._mode == "queueing":
|
|
return await self.queueing_call(obj)
|
|
else:
|
|
return await self.watching_call(obj)
|
|
|
|
def handle_recv(self, recv_obj: T):
|
|
self._result_values.append(recv_obj)
|
|
if len(self._result_values) == self._fan_out:
|
|
self._result_event.set()
|
|
|
|
@staticmethod
|
|
def merge_results(results):
|
|
all_success = all([r.success for r in results])
|
|
all_message = [r.message for r in results]
|
|
all_message = " | ".join(all_message)
|
|
return all_success, all_message
|
|
|
|
|
|
class TokenizerCommunicatorMixin:
|
|
"""Mixin class for TokenizerManager to handle communication with the scheduler."""
|
|
|
|
def init_communicators(self: TokenizerManager, server_args: ServerArgs):
|
|
# Communicators
|
|
self.init_weights_update_group_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.destroy_weights_update_group_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.update_weights_from_distributed_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.init_weights_send_group_for_remote_instance_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.send_weights_to_remote_instance_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.update_weights_from_tensor_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.update_weights_from_ipc_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.get_weights_by_name_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.release_memory_occupation_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.resume_memory_occupation_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.check_weights_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.slow_down_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.flush_cache_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.clear_hicache_storage_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.attach_hicache_storage_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.detach_hicache_storage_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.pin_prefix_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.profile_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.get_internal_state_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.set_internal_state_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.expert_distribution_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.update_lora_adapter_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.get_load_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size, mode="watching"
|
|
)
|
|
self.get_loads_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
self.dumper_control_communicator = _Communicator(
|
|
self.send_to_scheduler, server_args.dp_size
|
|
)
|
|
|
|
self._result_dispatcher += self._get_communicator_dispatcher()
|
|
|
|
def _get_communicator_dispatcher(self: TokenizerManager):
|
|
return TypeBasedDispatcher(
|
|
[
|
|
(
|
|
InitWeightsUpdateGroupReqOutput,
|
|
self.init_weights_update_group_communicator.handle_recv,
|
|
),
|
|
(
|
|
DestroyWeightsUpdateGroupReqOutput,
|
|
self.destroy_weights_update_group_communicator.handle_recv,
|
|
),
|
|
(
|
|
UpdateWeightsFromDistributedReqOutput,
|
|
self.update_weights_from_distributed_communicator.handle_recv,
|
|
),
|
|
(
|
|
InitWeightsSendGroupForRemoteInstanceReqOutput,
|
|
self.init_weights_send_group_for_remote_instance_communicator.handle_recv,
|
|
),
|
|
(
|
|
SendWeightsToRemoteInstanceReqOutput,
|
|
self.send_weights_to_remote_instance_communicator.handle_recv,
|
|
),
|
|
(
|
|
UpdateWeightsFromTensorReqOutput,
|
|
self.update_weights_from_tensor_communicator.handle_recv,
|
|
),
|
|
(
|
|
UpdateWeightsFromIPCReqOutput,
|
|
self.update_weights_from_ipc_communicator.handle_recv,
|
|
),
|
|
(
|
|
GetWeightsByNameReqOutput,
|
|
self.get_weights_by_name_communicator.handle_recv,
|
|
),
|
|
(
|
|
ReleaseMemoryOccupationReqOutput,
|
|
self.release_memory_occupation_communicator.handle_recv,
|
|
),
|
|
(
|
|
ResumeMemoryOccupationReqOutput,
|
|
self.resume_memory_occupation_communicator.handle_recv,
|
|
),
|
|
(
|
|
CheckWeightsReqOutput,
|
|
self.check_weights_communicator.handle_recv,
|
|
),
|
|
(
|
|
SlowDownReqOutput,
|
|
self.slow_down_communicator.handle_recv,
|
|
),
|
|
(
|
|
ClearHiCacheReqOutput,
|
|
self.clear_hicache_storage_communicator.handle_recv,
|
|
),
|
|
(
|
|
AttachHiCacheStorageReqOutput,
|
|
self.attach_hicache_storage_communicator.handle_recv,
|
|
),
|
|
(
|
|
DetachHiCacheStorageReqOutput,
|
|
self.detach_hicache_storage_communicator.handle_recv,
|
|
),
|
|
(
|
|
PinPrefixReqOutput,
|
|
self.pin_prefix_communicator.handle_recv,
|
|
),
|
|
(
|
|
FlushCacheReqOutput,
|
|
self.flush_cache_communicator.handle_recv,
|
|
),
|
|
(
|
|
ProfileReqOutput,
|
|
self.profile_communicator.handle_recv,
|
|
),
|
|
(
|
|
GetInternalStateReqOutput,
|
|
self.get_internal_state_communicator.handle_recv,
|
|
),
|
|
(
|
|
SetInternalStateReqOutput,
|
|
self.set_internal_state_communicator.handle_recv,
|
|
),
|
|
(
|
|
ExpertDistributionReqOutput,
|
|
self.expert_distribution_communicator.handle_recv,
|
|
),
|
|
(
|
|
LoRAUpdateOutput,
|
|
self.update_lora_adapter_communicator.handle_recv,
|
|
),
|
|
(
|
|
GetLoadReqOutput,
|
|
self.get_load_communicator.handle_recv,
|
|
),
|
|
(
|
|
GetLoadsReqOutput,
|
|
self.get_loads_communicator.handle_recv,
|
|
),
|
|
(
|
|
DumperControlReqOutput,
|
|
self.dumper_control_communicator.handle_recv,
|
|
),
|
|
]
|
|
)
|
|
|
|
async def flush_cache(self: TokenizerManager) -> FlushCacheReqOutput:
|
|
self.auto_create_handle_loop()
|
|
return (await self.flush_cache_communicator(FlushCacheReqInput()))[0]
|
|
|
|
async def clear_hicache_storage(self: TokenizerManager) -> ClearHiCacheReqOutput:
|
|
"""Clear the hierarchical cache storage."""
|
|
self.auto_create_handle_loop()
|
|
# Delegate to the scheduler to handle HiCacheStorage clearing
|
|
return (await self.clear_hicache_storage_communicator(ClearHiCacheReqInput()))[
|
|
0
|
|
]
|
|
|
|
async def attach_hicache_storage(
|
|
self: TokenizerManager,
|
|
hicache_storage_backend: str,
|
|
hicache_storage_backend_extra_config_json: Optional[str] = None,
|
|
hicache_storage_prefetch_policy: Optional[str] = None,
|
|
hicache_write_policy: Optional[str] = None,
|
|
) -> AttachHiCacheStorageReqOutput:
|
|
"""Attach (enable) HiCache storage backend at runtime."""
|
|
self.auto_create_handle_loop()
|
|
results = await self.attach_hicache_storage_communicator(
|
|
AttachHiCacheStorageReqInput(
|
|
hicache_storage_backend=hicache_storage_backend,
|
|
hicache_storage_backend_extra_config_json=hicache_storage_backend_extra_config_json,
|
|
hicache_storage_prefetch_policy=hicache_storage_prefetch_policy,
|
|
hicache_write_policy=hicache_write_policy,
|
|
)
|
|
)
|
|
|
|
all_success, all_message = _Communicator.merge_results(results)
|
|
out = AttachHiCacheStorageReqOutput(success=all_success, message=all_message)
|
|
# TODO: partial rollback if failed
|
|
if all_success:
|
|
# Keep tokenizer side server_info consistent with scheduler side.
|
|
self.server_args.hicache_storage_backend = hicache_storage_backend
|
|
if hicache_storage_backend_extra_config_json is not None:
|
|
self.server_args.hicache_storage_backend_extra_config = (
|
|
hicache_storage_backend_extra_config_json
|
|
)
|
|
if hicache_storage_prefetch_policy is not None:
|
|
self.server_args.hicache_storage_prefetch_policy = (
|
|
hicache_storage_prefetch_policy
|
|
)
|
|
if hicache_write_policy is not None:
|
|
self.server_args.hicache_write_policy = hicache_write_policy
|
|
return out
|
|
|
|
async def detach_hicache_storage(
|
|
self: TokenizerManager,
|
|
) -> DetachHiCacheStorageReqOutput:
|
|
"""Detach (disable) HiCache storage backend at runtime."""
|
|
self.auto_create_handle_loop()
|
|
results = await self.detach_hicache_storage_communicator(
|
|
DetachHiCacheStorageReqInput()
|
|
)
|
|
|
|
all_success, all_message = _Communicator.merge_results(results)
|
|
out = DetachHiCacheStorageReqOutput(success=all_success, message=all_message)
|
|
# TODO: partial rollback if failed
|
|
if all_success:
|
|
self.server_args.hicache_storage_backend = None
|
|
self.server_args.hicache_storage_backend_extra_config = None
|
|
return out
|
|
|
|
async def pin_prefix(
|
|
self: TokenizerManager, token_ids: List[int], ttl_seconds: int = 300
|
|
) -> PinPrefixReqOutput:
|
|
"""Pin a prefix by token_ids to resist eviction."""
|
|
results = await self.pin_prefix_communicator(
|
|
PinPrefixReqInput(token_ids=token_ids, ttl_seconds=ttl_seconds)
|
|
)
|
|
all_success, all_message = _Communicator.merge_results(results)
|
|
total = sum(r.nodes_pinned for r in results)
|
|
return PinPrefixReqOutput(
|
|
success=all_success, nodes_pinned=total, message=all_message
|
|
)
|
|
|
|
async def start_profile(
|
|
self: TokenizerManager,
|
|
output_dir: Optional[str] = None,
|
|
start_step: Optional[int] = None,
|
|
num_steps: Optional[int] = None,
|
|
activities: Optional[List[str]] = None,
|
|
with_stack: Optional[bool] = None,
|
|
record_shapes: Optional[bool] = None,
|
|
profile_by_stage: bool = False,
|
|
merge_profiles: bool = False,
|
|
profile_prefix: Optional[str] = None,
|
|
profile_stages: Optional[List[str]] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
|
|
with_stack = False if with_stack is False or env_with_stack is False else True
|
|
env_record_shapes: bool = get_bool_env_var(
|
|
"SGLANG_PROFILE_RECORD_SHAPES", "true"
|
|
)
|
|
record_shapes = (record_shapes is not False) and env_record_shapes
|
|
req = ProfileReq(
|
|
type=ProfileReqType.START_PROFILE,
|
|
output_dir=output_dir,
|
|
start_step=start_step,
|
|
num_steps=num_steps,
|
|
activities=activities,
|
|
with_stack=with_stack,
|
|
record_shapes=record_shapes,
|
|
profile_by_stage=profile_by_stage,
|
|
profile_id=str(time.time()),
|
|
merge_profiles=merge_profiles,
|
|
profile_prefix=profile_prefix,
|
|
profile_stages=profile_stages,
|
|
)
|
|
return await self._execute_profile(req)
|
|
|
|
async def stop_profile(self: TokenizerManager):
|
|
self.auto_create_handle_loop()
|
|
req = ProfileReq(type=ProfileReqType.STOP_PROFILE)
|
|
return await self._execute_profile(req)
|
|
|
|
async def _execute_profile(self: TokenizerManager, req: ProfileReq):
|
|
result = (await self.profile_communicator(req))[0]
|
|
if not result.success:
|
|
raise RuntimeError(result.message)
|
|
return result
|
|
|
|
async def start_expert_distribution_record(self: TokenizerManager):
|
|
self.auto_create_handle_loop()
|
|
req = ExpertDistributionReq(action=ExpertDistributionReqType.START_RECORD)
|
|
await self.expert_distribution_communicator(req)
|
|
|
|
async def stop_expert_distribution_record(self: TokenizerManager):
|
|
self.auto_create_handle_loop()
|
|
req = ExpertDistributionReq(action=ExpertDistributionReqType.STOP_RECORD)
|
|
await self.expert_distribution_communicator(req)
|
|
|
|
async def dump_expert_distribution_record(self: TokenizerManager):
|
|
self.auto_create_handle_loop()
|
|
req = ExpertDistributionReq(action=ExpertDistributionReqType.DUMP_RECORD)
|
|
await self.expert_distribution_communicator(req)
|
|
|
|
async def init_weights_update_group(
|
|
self: TokenizerManager,
|
|
obj: InitWeightsUpdateGroupReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
assert (
|
|
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
|
|
), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
|
|
|
|
results = await self.init_weights_update_group_communicator(obj)
|
|
return _Communicator.merge_results(results)
|
|
|
|
async def destroy_weights_update_group(
|
|
self: TokenizerManager,
|
|
obj: DestroyWeightsUpdateGroupReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
assert (
|
|
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
|
|
), "dp_size must be 1 or dp attention must be enabled for destroy parameter update group"
|
|
|
|
results = await self.destroy_weights_update_group_communicator(obj)
|
|
return _Communicator.merge_results(results)
|
|
|
|
async def update_weights_from_distributed(
|
|
self: TokenizerManager,
|
|
obj: UpdateWeightsFromDistributedReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
assert (
|
|
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
|
|
), "dp_size must be 1 or dp attention must be enabled for update weights from distributed"
|
|
|
|
if obj.abort_all_requests:
|
|
self.abort_request(abort_all=True)
|
|
|
|
# Immediately update the weights if the engine is in paused state
|
|
async with self.is_pause_cond:
|
|
is_paused = self.is_pause
|
|
|
|
lock_context = (
|
|
self.model_update_lock.writer_lock if not is_paused else nullcontext()
|
|
)
|
|
async with lock_context:
|
|
results = await self.update_weights_from_distributed_communicator(obj)
|
|
|
|
success, message = _Communicator.merge_results(results)
|
|
if success and obj.weight_version is not None:
|
|
self._update_weight_version_if_provided(obj.weight_version)
|
|
message += f" Weight version updated to {obj.weight_version}."
|
|
|
|
return success, message
|
|
|
|
async def init_weights_send_group_for_remote_instance(
|
|
self: TokenizerManager,
|
|
obj: InitWeightsSendGroupForRemoteInstanceReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
# TODO: support DP
|
|
assert (
|
|
self.server_args.dp_size == 1
|
|
), "dp_size must be 1 for init_weights_send_group_for_remote_instance"
|
|
result = (
|
|
await self.init_weights_send_group_for_remote_instance_communicator(obj)
|
|
)[0]
|
|
return result.success, result.message
|
|
|
|
async def send_weights_to_remote_instance(
|
|
self: TokenizerManager,
|
|
obj: SendWeightsToRemoteInstanceReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
# TODO: support DP
|
|
assert (
|
|
self.server_args.dp_size == 1
|
|
), "dp_size must be 1 for send_weights_to_remote_instance"
|
|
result = (await self.send_weights_to_remote_instance_communicator(obj))[0]
|
|
return result.success, result.message
|
|
|
|
async def update_weights_from_tensor(
|
|
self: TokenizerManager,
|
|
obj: UpdateWeightsFromTensorReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
self.auto_create_handle_loop()
|
|
assert (
|
|
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
|
|
), "dp_size must be 1 or dp attention must be enabled for update weights from tensor"
|
|
|
|
if obj.abort_all_requests:
|
|
self.abort_request(abort_all=True)
|
|
|
|
# Immediately update the weights if the engine is in paused state
|
|
async with self.is_pause_cond:
|
|
is_paused = self.is_pause
|
|
|
|
lock_context = (
|
|
self.model_update_lock.writer_lock if not is_paused else nullcontext()
|
|
)
|
|
async with lock_context:
|
|
results = await self.update_weights_from_tensor_communicator(obj)
|
|
|
|
success, message = _Communicator.merge_results(results)
|
|
if success and obj.weight_version is not None:
|
|
self._update_weight_version_if_provided(obj.weight_version)
|
|
message += f" Weight version updated to {obj.weight_version}."
|
|
|
|
return success, message
|
|
|
|
async def update_weights_from_ipc(
|
|
self: TokenizerManager,
|
|
obj: UpdateWeightsFromIPCReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> Tuple[bool, str]:
|
|
"""Update weights via IPC for checkpoint-engine integration."""
|
|
self.auto_create_handle_loop()
|
|
try:
|
|
# For now, we only support single data parallel instance
|
|
assert (
|
|
self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
|
|
), "dp_size must be 1 or dp attention must be enabled for update weights from IPC"
|
|
logger.info("Starting IPC weight update")
|
|
# This means that weight sync cannot run while requests are in progress.
|
|
async with self.model_update_lock.writer_lock:
|
|
result = (await self.update_weights_from_ipc_communicator(obj))[0]
|
|
success, message = result.success, result.message
|
|
except Exception as e:
|
|
error_msg = f"IPC weight update failed: {str(e)}"
|
|
logger.error(error_msg)
|
|
success, message = False, error_msg
|
|
|
|
if success and obj.weight_version is not None:
|
|
self._update_weight_version_if_provided(obj.weight_version)
|
|
message += f" Weight version updated to {obj.weight_version}."
|
|
|
|
return success, message
|
|
|
|
async def _unload_lora_adapter_locked(
|
|
self: TokenizerManager,
|
|
obj: UnloadLoRAAdapterReqInput,
|
|
) -> UnloadLoRAAdapterReqOutput:
|
|
assert (
|
|
self.lora_update_lock.locked()
|
|
), "self.lora_update_lock must be locked in order for self._unload_lora_adapter_locked() to be called"
|
|
|
|
# Unregister the LoRA adapter from the registry to stop new requests for this adapter
|
|
# from being started.
|
|
lora_id = await self.lora_registry.unregister(obj.lora_name)
|
|
obj.lora_id = lora_id
|
|
|
|
# Initiate the actual unloading operation at the backend processes only after all
|
|
# ongoing requests using this LoRA adapter are finished.
|
|
await self.lora_registry.wait_for_unload(lora_id)
|
|
result = (await self.update_lora_adapter_communicator(obj))[0]
|
|
|
|
return result
|
|
|
|
async def load_lora_adapter(
|
|
self: TokenizerManager,
|
|
obj: LoadLoRAAdapterReqInput,
|
|
_: Optional[fastapi.Request] = None,
|
|
) -> LoadLoRAAdapterReqOutput:
|
|
self.auto_create_handle_loop()
|
|
|
|
try:
|
|
if not self.server_args.enable_lora:
|
|
raise ValueError(
|
|
"LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
|
|
)
|
|
|
|
# TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
|
|
# with dp_size > 1.
|
|
assert (
|
|
self.server_args.dp_size == 1
|
|
), "dp_size must be 1 for dynamic lora loading"
|
|
logger.info(
|
|
"Start load Lora adapter. Lora name=%s, path=%s",
|
|
obj.lora_name,
|
|
obj.lora_path,
|
|
)
|
|
|
|
async with self.lora_update_lock:
|
|
# Generate new uniquely identifiable LoRARef object.
|
|
new_adapter = LoRARef(
|
|
lora_name=obj.lora_name,
|
|
lora_path=obj.lora_path,
|
|
pinned=obj.pinned,
|
|
)
|
|
|
|
# Trigger the actual loading operation at the backend processes.
|
|
obj.lora_id = new_adapter.lora_id
|
|
result = (await self.update_lora_adapter_communicator(obj))[0]
|
|
|
|
# Register the LoRA adapter only after loading is successful.
|
|
if result.success:
|
|
await self.lora_registry.register(new_adapter)
|
|
self.lora_ref_cache[obj.lora_name] = new_adapter
|
|
|
|
if self.server_args.max_loaded_loras is not None:
|
|
while (
|
|
self.lora_registry.num_registered_loras
|
|
> self.server_args.max_loaded_loras
|
|
):
|
|
lru_lora_name = await self.lora_registry.lru_lora_name(
|
|
exclude_pinned=True
|
|
)
|
|
if lru_lora_name is None:
|
|
raise ValueError(
|
|
"Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. "
|
|
f"LoRA registry is: {self.lora_registry._registry}"
|
|
)
|
|
|
|
logger.info(
|
|
f"Unloading least recently used LoRA adapter '{lru_lora_name}' "
|
|
f"(current number of adapters: {self.lora_registry.num_registered_loras}, "
|
|
f"max allowed: {self.server_args.max_loaded_loras})"
|
|
)
|
|
|
|
unload_result = await self._unload_lora_adapter_locked(
|
|
UnloadLoRAAdapterReqInput(lora_name=lru_lora_name)
|
|
)
|
|
if not unload_result.success:
|
|
raise ValueError(
|
|
f"Error while unloading LRU LoRA adapter '{lru_lora_name}': "
|
|
f"{unload_result.error_message}"
|
|
)
|
|
del result.loaded_adapters[lru_lora_name]
|
|
|
|
return result
|
|
except ValueError as e:
|
|
return LoadLoRAAdapterReqOutput(
|
|
success=False,
|
|
error_message=str(e),
|
|
)
|
|
|
|
async def load_lora_adapter_from_tensors(
|
|
self: TokenizerManager,
|
|
obj: LoadLoRAAdapterFromTensorsReqInput,
|
|
_: Optional[fastapi.Request] = None,
|
|
) -> LoadLoRAAdapterFromTensorsReqOutput:
|
|
self.auto_create_handle_loop()
|
|
|
|
try:
|
|
if not self.server_args.enable_lora:
|
|
raise ValueError(
|
|
"LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
|
|
)
|
|
|
|
assert (
|
|
self.server_args.dp_size == 1
|
|
), "dp_size must be 1 for dynamic lora loading"
|
|
logger.info(
|
|
"Start load Lora adapter from tensors. Lora name=%s",
|
|
obj.lora_name,
|
|
)
|
|
|
|
async with self.lora_update_lock:
|
|
new_adapter = LoRARef(
|
|
lora_name=obj.lora_name,
|
|
lora_path="__tensor__",
|
|
pinned=obj.pinned,
|
|
)
|
|
obj.lora_id = new_adapter.lora_id
|
|
result = (await self.update_lora_adapter_communicator(obj))[0]
|
|
|
|
if result.success:
|
|
await self.lora_registry.register(new_adapter)
|
|
self.lora_ref_cache[obj.lora_name] = new_adapter
|
|
if self.server_args.max_loaded_loras is not None:
|
|
while (
|
|
self.lora_registry.num_registered_loras
|
|
> self.server_args.max_loaded_loras
|
|
):
|
|
lru_lora_name = await self.lora_registry.lru_lora_name(
|
|
exclude_pinned=True
|
|
)
|
|
if lru_lora_name is None:
|
|
raise ValueError(
|
|
"Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. "
|
|
f"LoRA registry is: {self.lora_registry._registry}"
|
|
)
|
|
|
|
logger.info(
|
|
f"Unloading least recently used LoRA adapter '{lru_lora_name}' "
|
|
f"(current number of adapters: {self.lora_registry.num_registered_loras}, "
|
|
f"max allowed: {self.server_args.max_loaded_loras})"
|
|
)
|
|
|
|
unload_result = await self._unload_lora_adapter_locked(
|
|
UnloadLoRAAdapterReqInput(lora_name=lru_lora_name)
|
|
)
|
|
if not unload_result.success:
|
|
raise ValueError(
|
|
f"Error while unloading LRU LoRA adapter '{lru_lora_name}': "
|
|
f"{unload_result.error_message}"
|
|
)
|
|
del result.loaded_adapters[lru_lora_name]
|
|
|
|
return result
|
|
except ValueError as e:
|
|
return LoadLoRAAdapterFromTensorsReqOutput(
|
|
success=False,
|
|
error_message=str(e),
|
|
)
|
|
|
|
async def unload_lora_adapter(
|
|
self: TokenizerManager,
|
|
obj: UnloadLoRAAdapterReqInput,
|
|
_: Optional[fastapi.Request] = None,
|
|
) -> UnloadLoRAAdapterReqOutput:
|
|
self.auto_create_handle_loop()
|
|
|
|
try:
|
|
if not self.server_args.enable_lora:
|
|
raise ValueError(
|
|
"LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
|
|
)
|
|
|
|
assert (
|
|
obj.lora_name is not None
|
|
), "lora_name must be provided to unload LoRA adapter"
|
|
|
|
# TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
|
|
# with dp_size > 1.
|
|
assert (
|
|
self.server_args.dp_size == 1
|
|
), "dp_size must be 1 for dynamic lora loading"
|
|
logger.info(
|
|
"Start unload Lora adapter. Lora name=%s",
|
|
obj.lora_name,
|
|
)
|
|
|
|
async with self.lora_update_lock:
|
|
return await self._unload_lora_adapter_locked(obj)
|
|
except ValueError as e:
|
|
return UnloadLoRAAdapterReqOutput(success=False, error_message=str(e))
|
|
|
|
async def get_weights_by_name(
|
|
self: TokenizerManager,
|
|
obj: GetWeightsByNameReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
results = await self.get_weights_by_name_communicator(obj)
|
|
all_parameters = [r.parameter for r in results]
|
|
if self.server_args.dp_size == 1:
|
|
return all_parameters[0]
|
|
else:
|
|
return all_parameters
|
|
|
|
async def release_memory_occupation(
|
|
self: TokenizerManager,
|
|
obj: ReleaseMemoryOccupationReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
await self.release_memory_occupation_communicator(obj)
|
|
|
|
async def resume_memory_occupation(
|
|
self: TokenizerManager,
|
|
obj: ResumeMemoryOccupationReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
await self.resume_memory_occupation_communicator(obj)
|
|
|
|
async def check_weights(
|
|
self: TokenizerManager,
|
|
obj: CheckWeightsReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
) -> CheckWeightsReqOutput:
|
|
self.auto_create_handle_loop()
|
|
results = await self.check_weights_communicator(obj)
|
|
return _Communicator.merge_results(results)
|
|
|
|
async def slow_down(
|
|
self: TokenizerManager,
|
|
obj: SlowDownReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
await self.slow_down_communicator(obj)
|
|
|
|
async def get_internal_state(self: TokenizerManager) -> List[Dict[Any, Any]]:
|
|
self.auto_create_handle_loop()
|
|
req = GetInternalStateReq()
|
|
responses: List[GetInternalStateReqOutput] = (
|
|
await self.get_internal_state_communicator(req)
|
|
)
|
|
# Many DP ranks
|
|
return [res.internal_state for res in responses]
|
|
|
|
async def set_internal_state(
|
|
self: TokenizerManager, obj: SetInternalStateReq
|
|
) -> List[bool]:
|
|
self.auto_create_handle_loop()
|
|
responses: List[SetInternalStateReqOutput] = (
|
|
await self.set_internal_state_communicator(obj)
|
|
)
|
|
return [res.updated for res in responses]
|
|
|
|
async def dumper_control(
|
|
self: TokenizerManager, obj: DumperControlReqInput
|
|
) -> List[DumperControlReqOutput]:
|
|
self.auto_create_handle_loop()
|
|
return await self.dumper_control_communicator(obj)
|
|
|
|
async def get_load(self: TokenizerManager) -> List[GetLoadReqOutput]:
|
|
self.auto_create_handle_loop()
|
|
req = GetLoadReqInput()
|
|
return await self.get_load_communicator(req)
|
|
|
|
async def get_loads(
|
|
self: TokenizerManager,
|
|
include: Optional[List[str]] = None,
|
|
dp_rank: Optional[int] = None,
|
|
) -> List[GetLoadsReqOutput]:
|
|
"""
|
|
Get comprehensive load metrics for /v1/loads endpoint.
|
|
|
|
Args:
|
|
include: List of sections to include. Options: core, memory, spec, lora, disagg, queues, all
|
|
dp_rank: Optional filter for specific DP rank
|
|
|
|
Returns:
|
|
List of GetLoadsReqOutput, one per scheduler (filtered by dp_rank if specified)
|
|
"""
|
|
self.auto_create_handle_loop()
|
|
req = GetLoadsReqInput(
|
|
include=include if include else ["all"],
|
|
dp_rank=dp_rank,
|
|
)
|
|
results = await self.get_loads_communicator(req)
|
|
|
|
# Filter by dp_rank if specified
|
|
if dp_rank is not None:
|
|
results = [r for r in results if r.dp_rank == dp_rank]
|
|
|
|
return results
|
|
|
|
async def open_session(
|
|
self: TokenizerManager,
|
|
obj: OpenSessionReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
self.auto_create_handle_loop()
|
|
if obj.streaming:
|
|
if not self.server_args.enable_streaming_session:
|
|
raise ValueError(
|
|
"Streaming sessions are disabled. "
|
|
"Please relaunch with --enable-streaming-session."
|
|
)
|
|
if (
|
|
self.server_args.speculative_algorithm is not None
|
|
and not self.server_args.disable_overlap_schedule
|
|
):
|
|
raise ValueError(
|
|
"Streaming sessions are incompatible with speculative decoding v2 "
|
|
"(overlap + speculative). Use --disable-overlap-schedule or "
|
|
"disable speculative decoding."
|
|
)
|
|
|
|
if obj.session_id is None:
|
|
obj.session_id = uuid.uuid4().hex
|
|
elif obj.session_id in self.session_futures:
|
|
return None
|
|
|
|
self.send_to_scheduler.send_pyobj(obj)
|
|
|
|
self.session_futures[obj.session_id] = asyncio.Future()
|
|
session_id = await self.session_futures[obj.session_id]
|
|
del self.session_futures[obj.session_id]
|
|
return session_id
|
|
|
|
async def close_session(
|
|
self: TokenizerManager,
|
|
obj: CloseSessionReqInput,
|
|
request: Optional[fastapi.Request] = None,
|
|
):
|
|
await self.send_to_scheduler.send_pyobj(obj)
|
|
|
|
def _update_weight_version_if_provided(
|
|
self: TokenizerManager, weight_version: Optional[str]
|
|
) -> None:
|
|
"""Update weight version if provided."""
|
|
if weight_version is not None:
|
|
self.server_args.weight_version = weight_version
|