[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:
co-authored by
aoshen02
Claude Opus 5.5
parent
5e887a078b
commit
a7a9ac2745
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user