fix(loads): preserve include filtering after watching mode switch (#22959)
This commit is contained in:
@@ -142,10 +142,11 @@ class _Communicator(Generic[T]):
|
||||
if obj:
|
||||
self._sender.send_pyobj(obj)
|
||||
|
||||
event = self._result_event
|
||||
# NOTE: Capture list ref before await so later awaiters survive clearing.
|
||||
values = self._result_values
|
||||
event = self._result_event
|
||||
await event.wait()
|
||||
# Capture list ref before await so later awaiters survive clearing.
|
||||
|
||||
result_values = copy.deepcopy(values)
|
||||
if self._result_event is event:
|
||||
self._result_event = self._result_values = None
|
||||
@@ -1064,6 +1065,21 @@ class TokenizerCommunicatorMixin:
|
||||
if dp_rank is not None:
|
||||
results = [r for r in results if r.dp_rank == dp_rank]
|
||||
|
||||
# Filter optional sections client-side (scheduler always returns all)
|
||||
if include and "all" not in include:
|
||||
include_set = set(include)
|
||||
_section_attrs = {
|
||||
"memory": "memory",
|
||||
"spec": "speculative",
|
||||
"lora": "lora",
|
||||
"disagg": "disaggregation",
|
||||
"queues": "queues",
|
||||
}
|
||||
for r in results:
|
||||
for key, attr in _section_attrs.items():
|
||||
if key not in include_set:
|
||||
setattr(r, attr, None)
|
||||
|
||||
return results
|
||||
|
||||
async def open_session(
|
||||
|
||||
Reference in New Issue
Block a user