Skip to content

Commit 314bd83

Browse files
authored
fix: Move health check off event loop (#740)
1 parent 37ca4b3 commit 314bd83

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

src/aiperf/workers/worker.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
Conversation,
3535
ErrorDetails,
3636
ModelEndpointInfo,
37+
ProcessHealth,
3738
ReasoningResponseData,
3839
RequestInfo,
3940
RequestRecord,
@@ -263,12 +264,13 @@ async def _send_worker_shutdown_message(self) -> None:
263264
)
264265
async def _health_check_task(self) -> None:
265266
"""Task to report the health of the worker to the worker manager."""
266-
await self.publish(self.create_health_message())
267+
health = await asyncio.to_thread(self.get_process_health)
268+
await self.publish(self.create_health_message(health))
267269

268-
def create_health_message(self) -> WorkerHealthMessage:
270+
def create_health_message(self, health: ProcessHealth) -> WorkerHealthMessage:
269271
return WorkerHealthMessage(
270272
service_id=self.service_id,
271-
health=self.get_process_health(),
273+
health=health,
272274
task_stats=self.task_stats,
273275
)
274276

@@ -688,7 +690,8 @@ async def _on_profile_configure_command(self, message: CommandMessage) -> None:
688690
timeout=Environment.DATASET.CONFIGURATION_TIMEOUT,
689691
)
690692
if self.is_debug_enabled:
691-
memory_usage = self.get_process_health().memory_usage / BYTES_PER_MIB
693+
health = await asyncio.to_thread(self.get_process_health)
694+
memory_usage = health.memory_usage / BYTES_PER_MIB
692695
self.memory_usage_before_profiling = memory_usage
693696
self.debug(f"Memory usage before profiling: {memory_usage:.2f} MiB")
694697

0 commit comments

Comments
 (0)