feat(api): load-balance background jobs between devices

ssube · Feb 4, 2023 · efee374 · efee374
1 parent 7c96145
commit efee374
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git 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 @@ def done(self, key: str) -> Tuple[bool, int]:
         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
@@ -1,5 +1,5 @@
 {
-  "version": "0.5.0",
+  "version": "0.6.0",
   "bottom": {
     "default": 0,
     "min": 0,