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:
|
if obj:
|
||||||
self._sender.send_pyobj(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
|
values = self._result_values
|
||||||
|
event = self._result_event
|
||||||
await event.wait()
|
await event.wait()
|
||||||
# Capture list ref before await so later awaiters survive clearing.
|
|
||||||
result_values = copy.deepcopy(values)
|
result_values = copy.deepcopy(values)
|
||||||
if self._result_event is event:
|
if self._result_event is event:
|
||||||
self._result_event = self._result_values = None
|
self._result_event = self._result_values = None
|
||||||
@@ -1064,6 +1065,21 @@ class TokenizerCommunicatorMixin:
|
|||||||
if dp_rank is not None:
|
if dp_rank is not None:
|
||||||
results = [r for r in results if r.dp_rank == dp_rank]
|
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
|
return results
|
||||||
|
|
||||||
async def open_session(
|
async def open_session(
|
||||||
|
|||||||
Reference in New Issue
Block a user