From 5e0231c01b0008c7a0394e94a55a55601124437d Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sat, 4 Feb 2023 15:49:05 -0600 Subject: [PATCH] feat(api): distribute jobs to devices using round-robin (#38) --- api/onnx_web/device_pool.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/api/onnx_web/device_pool.py b/api/onnx_web/device_pool.py index 936065eb..fdeda89e 100644 --- a/api/onnx_web/device_pool.py +++ b/api/onnx_web/device_pool.py @@ -92,11 +92,13 @@ class Job: class DevicePoolExecutor: devices: List[DeviceParams] = None jobs: List[Job] = None + next_device: int = 0 pool: Union[ProcessPoolExecutor, ThreadPoolExecutor] = None def __init__(self, devices: List[DeviceParams], pool: Optional[Union[ProcessPoolExecutor, ThreadPoolExecutor]] = None): self.devices = devices self.jobs = [] + self.next_device = 0 device_count = len(devices) if pool is None: @@ -131,11 +133,16 @@ class DevicePoolExecutor: logger.warn('checking status for unknown key: %s', key) return (None, 0) + def get_next_device(self): + device = self.next_device + self.next_device = (self.next_device + 1) % len(self.devices) + return device + def prune(self): self.jobs[:] = [job for job in self.jobs if job.future.done()] def submit(self, key: str, fn: Callable[..., None], /, *args, **kwargs) -> None: - context = JobContext(key, self.devices, device_index=0) + context = JobContext(key, self.devices, device_index=self.get_next_device()) future = self.pool.submit(fn, context, *args, **kwargs) job = Job(key, future, context) self.jobs.append(job)