feat(api): split up status endpoint by job status

ssube · Mar 26, 2023 · ccf8d51 · ccf8d51
1 parent ea36082
commit ccf8d51
Showing 1 changed file with 21 additions and 23 deletions.
diff --git a/api/onnx_web/worker/pool.py b/api/onnx_web/worker/pool.py
@@ -372,20 +372,21 @@ def submit(
         job = JobCommand(key, device, fn, args, kwargs)
         self.pending_jobs.append(job)
 
-    def status(self) -> List[Tuple[str, int, bool, bool, bool, bool]]:
-        history = [
-            (
-                name,
-                job.progress,
-                False,
-                job.finished,
-                job.cancelled,
-                job.failed,
-            )
-            for name, job in self.running_jobs.items()
-        ]
-        history.extend(
-            [
+    def status(self) -> Dict[str, List[Tuple[str, int, bool, bool, bool, bool]]]:
+        return {
+            "cancelled": [],
+            "finished": [
+                (
+                    job.job,
+                    job.progress,
+                    False,
+                    job.finished,
+                    job.cancelled,
+                    job.failed,
+                )
+                for job in self.finished_jobs
+            ],
+            "pending": [
                 (
                     job.name,
                     0,
@@ -395,22 +396,19 @@ def status(self) -> List[Tuple[str, int, bool, bool, bool, bool]]:
                     False,
                 )
                 for job in self.pending_jobs
-            ]
-        )
-        history.extend(
-            [
+            ],
+            "running": [
                 (
-                    job.job,
+                    name,
                     job.progress,
                     False,
                     job.finished,
                     job.cancelled,
                     job.failed,
                 )
-                for job in self.finished_jobs
-            ]
-        )
-        return history
+                for name, job in self.running_jobs.items()
+            ],
+        }
 
     def next_job(self, device: str):
         for job in self.pending_jobs: