Skip to content

Commit

Permalink
fix(api): log number of background workers on startup
Browse files Browse the repository at this point in the history
  • Loading branch information
ssube committed Feb 4, 2023
1 parent 4e7bfd7 commit 732aa27
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions api/onnx_web/device_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ def on_progress(step: int, timestep: int, latents: Any):
if self.is_cancelled():
raise Exception('job has been cancelled')
else:
logger.debug('setting progress for job %s to %s', self.key, step)
self.set_progress(step)

return on_progress
Expand All @@ -56,7 +57,6 @@ def set_cancel(self, cancel: bool = True) -> None:
def set_progress(self, progress: int) -> None:
with self.progress.get_lock():
self.progress.value = progress
logger.debug('setting progress for job %s to %s', self.key, progress)


class Job:
Expand Down Expand Up @@ -91,7 +91,14 @@ class DevicePoolExecutor:
def __init__(self, devices: List[str], pool: Optional[Union[ProcessPoolExecutor, ThreadPoolExecutor]] = None):
self.devices = devices
self.jobs = []
self.pool = pool or ThreadPoolExecutor(len(devices))

device_count = len(devices)
if pool is None:
logger.info('creating thread pool executor for %s devices: %s', device_count, devices)
self.pool = ThreadPoolExecutor(device_count)
else:
logger.info('using existing pool for %s devices: %s', device_count, devices)
self.pool = pool

def cancel(self, key: str) -> bool:
'''
Expand Down

0 comments on commit 732aa27

Please sign in to comment.