[V1] Fix torch profiling for offline inference (#11125)

Signed-off-by: Roger Wang <ywang@roblox.com>
This commit is contained in:
Roger Wang
2024-12-12 07:51:53 -08:00
committed by GitHub
parent 85362f028c
commit 4816d20aa4
2 changed files with 21 additions and 14 deletions

View File

@@ -105,7 +105,7 @@ class InprocClient(EngineCoreClient):
def __del__(self):
self.shutdown()
async def profile(self, is_start=True) -> None:
def profile(self, is_start=True) -> None:
self.engine_core.profile(is_start)
@@ -212,7 +212,7 @@ class SyncMPClient(MPClient):
def abort_requests(self, request_ids: List[str]) -> None:
self._send_input(EngineCoreRequestType.ABORT, request_ids)
async def profile(self, is_start=True) -> None:
def profile(self, is_start=True) -> None:
self._send_input(EngineCoreRequestType.PROFILE,
EngineCoreProfile(is_start))