bugfix: status was passed wrongly if result is submitted (as draft is…

… not updated)
mlcommons · Oct 22, 2024 · 2a1d55d · 2a1d55d
1 parent 7d20e31
commit 2a1d55d
Showing 1 changed file with 23 additions and 23 deletions.
diff --git a/cli/medperf/web_ui/datasets/routes_run.py b/cli/medperf/web_ui/datasets/routes_run.py
@@ -116,6 +116,7 @@ def run_with_logs():
                     draft.logs.append(log)
                 result, status = future.result()
                 draft.status = status
+                results[result.local_id] = result
         finally:
             _task_queue.task_done()
 
@@ -146,7 +147,7 @@ async def run_benchmark(dataset_id: int, benchmark_id: int, model_id: int):
     return RunStatus(
         model_id=draft.model_id,
         status=draft.status,
-        result=results.get(draft.result_id)
+        result=draft.get_result()
     )
 
 
@@ -155,31 +156,30 @@ async def get_run_status(dataset_id: int, benchmark_id: int, model_id: int):
     result_id = f"b{benchmark_id}m{model_id}d{dataset_id}"
     draft = _drafts.get(result_id)
 
-    if not draft:
-        result = _load_result_if_exists(result_id)
-        if result:
-            if str(result.id).isdigit():
-                status = DraftStatus.submitted
-            else:
-                status = DraftStatus.executed
-            return RunStatus(
-                model_id=model_id,
-                status=status,
-                result=result
-            )
+    result = _load_result_if_exists(result_id)
+    if result:
+        if str(result.id).isdigit():
+            status = DraftStatus.submitted
         else:
-            return RunStatus(
-                model_id=model_id,
-                status=DraftStatus.n_a,
-                result=None
+            status = DraftStatus.executed
+        return RunStatus(
+            model_id=model_id,
+            status=status,
+            result=result
+        )
+    elif draft:
+        return RunStatus(
+            model_id=draft.model_id,
+            status=draft.status,
+            result=draft.get_result()
+        )
+    else:
+        return RunStatus(
+            model_id=model_id,
+            status=DraftStatus.n_a,
+            result=None
             )
 
-    return RunStatus(
-        model_id=draft.model_id,
-        status=draft.status,
-        result=draft.result
-    )
-
 
 @router.get("/run_draft/logs", response_class=StreamingResponse)
 async def get_run_logs(dataset_id: int, benchmark_id: int, model_id: int):