From efee374c164f6db7ef8a8b109e8382f3ab5b7aee Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sat, 4 Feb 2023 16:22:50 -0600 Subject: [PATCH] feat(api): load-balance background jobs between devices --- api/onnx_web/device_pool.py | 9 ++++++--- api/params.json | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/api/onnx_web/device_pool.py b/api/onnx_web/device_pool.py index 47fb8383..3568407d 100644 --- a/api/onnx_web/device_pool.py +++ b/api/onnx_web/device_pool.py @@ -1,3 +1,4 @@ +from collections import Counter from concurrent.futures import Future, ThreadPoolExecutor, ProcessPoolExecutor from logging import getLogger from multiprocessing import Value @@ -134,9 +135,11 @@ class DevicePoolExecutor: return (None, 0) def get_next_device(self): - device = self.next_device - self.next_device = (self.next_device + 1) % len(self.devices) - return device + job_devices = [job.context.device_index.value for job in self.jobs] + queued = Counter(job_devices).most_common() + logger.debug('jobs queued by device: %s', queued) + + return queued[-1] def prune(self): self.jobs[:] = [job for job in self.jobs if job.future.done()] diff --git a/api/params.json b/api/params.json index e5761331..08796608 100644 --- a/api/params.json +++ b/api/params.json @@ -1,5 +1,5 @@ { - "version": "0.5.0", + "version": "0.6.0", "bottom": { "default": 0, "min": 0,