[Core] Combine per-engine utility results like the Rust client (#59240)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
aoshen02
2026-09-29 22:20:51 +00:00
committed by GitHub
co-authored by aoshen02 Claude Opus 5.5
parent 5e887a078b
commit a7a9ac2745
+25 -13
View File
@@ -1239,6 +1239,23 @@ class AsyncMPClient(MPClient):
async def call_utility_async(self, method: str, *args) -> Any:
return await self._call_utility_async(method, *args, engine=self.core_engine)
async def call_utility_all_async(self, method: str, *args) -> list[Any]:
# Like the Rust client's call_utility: one result per managed engine.
return await asyncio.gather(
*[
self._call_utility_async(method, *args, engine=engine)
for engine in self.core_engines
]
)
async def call_utility_consensus_async(self, method: str, *args) -> Any:
results = await self.call_utility_all_async(method, *args)
if any(result != results[0] for result in results):
raise RuntimeError(
f"Engines returned different {method} results: {results}"
)
return results[0]
async def _call_utility_async(
self, method: str, *args, engine: EngineIdentity
) -> Any:
@@ -1274,7 +1291,7 @@ class AsyncMPClient(MPClient):
await self.call_utility_async("resume_scheduler")
async def is_scheduler_paused_async(self) -> bool:
return await self.call_utility_async("is_scheduler_paused")
return await self.call_utility_consensus_async("is_scheduler_paused")
async def profile_async(
self, is_start: bool = True, profile_prefix: str | None = None
@@ -1287,8 +1304,10 @@ class AsyncMPClient(MPClient):
async def reset_prefix_cache_async(
self, reset_running_requests: bool = False, reset_connector: bool = False
) -> bool:
return await self.call_utility_async(
"reset_prefix_cache", reset_running_requests, reset_connector
return all(
await self.call_utility_all_async(
"reset_prefix_cache", reset_running_requests, reset_connector
)
)
async def reset_encoder_cache_async(self) -> None:
@@ -1304,7 +1323,7 @@ class AsyncMPClient(MPClient):
return await self.call_utility_async("wake_up", tags)
async def is_sleeping_async(self) -> bool:
return await self.call_utility_async("is_sleeping")
return await self.call_utility_consensus_async("is_sleeping")
async def execute_dummy_batch_async(self) -> None:
await self.call_utility_async("execute_dummy_batch")
@@ -1313,7 +1332,7 @@ class AsyncMPClient(MPClient):
await self.call_utility_async("set_weight_version", weight_version)
async def get_weight_version_async(self) -> str:
return await self.call_utility_async("get_weight_version")
return await self.call_utility_consensus_async("get_weight_version")
async def add_lora_async(self, lora_request: LoRARequest) -> bool:
return await self.call_utility_async("add_lora", lora_request)
@@ -1645,14 +1664,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
async def call_utility_async(self, method: str, *args) -> Any:
# Only the result from the first engine is returned.
return (
await asyncio.gather(
*[
self._call_utility_async(method, *args, engine=engine)
for engine in self.core_engines
]
)
)[0]
return (await self.call_utility_all_async(method, *args))[0]
@staticmethod
async def process_engine_outputs(