fix(api): include worker totals in status endpoint

ssube · Mar 26, 2023 · 36bfcca · 36bfcca
1 parent 2d2283e
commit 36bfcca
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/api/onnx_web/worker/pool.py b/api/onnx_web/worker/pool.py
@@ -369,6 +369,9 @@ def submit(
         self.pending_jobs.append(job)
 
     def status(self) -> Dict[str, List[Tuple[str, int, bool, bool, bool, bool]]]:
+        """
+        Returns a tuple of: job/device, progress, progress, finished, cancelled, failed
+        """
         return {
             "cancelled": [],
             "finished": [
@@ -404,6 +407,17 @@ def status(self) -> Dict[str, List[Tuple[str, int, bool, bool, bool, bool]]]:
                 )
                 for name, job in self.running_jobs.items()
             ],
+            "total": [
+                (
+                    device,
+                    total,
+                    self.workers[device].is_alive(),
+                    False,
+                    False,
+                    False,
+                )
+                for device, total in self.total_jobs.items()
+            ],
         }
 
     def next_job(self, device: str):