From 9ea25780c66dc91a30d358dbaad3d16965702d87 Mon Sep 17 00:00:00 2001
From: Tara Ding <38710454+windtara0619@users.noreply.github.com>
Date: Mon, 27 Apr 2026 22:44:00 -0700
Subject: [PATCH] fix benchmark

---
 main.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/main.py b/main.py
index 8be3036ab..bab80379a 100644
--- a/main.py
+++ b/main.py
@@ -315,6 +315,12 @@ def prompt_worker(q, server_instance):
 
             need_gc = True
 
+            if benchmark_mode:
+                e.history_result["benchmark"] = {
+                    "execution_ms": (time.perf_counter() - execution_start_time) * 1000.0,
+                    "nodes": e.node_timing_ms,
+                }
+
             remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
             q.task_done(item_id,
                         e.history_result,
@@ -335,12 +341,6 @@ def prompt_worker(q, server_instance):
             else:
                 logging.info("Prompt executed in {:.2f} seconds".format(execution_time))
 
-            if benchmark_mode:
-                e.history_result["benchmark"] = {
-                    "execution_ms": execution_time * 1000.0,
-                    "nodes": e.node_timing_ms,
-                }
-
             if not asset_seeder.is_disabled():
                 paths = _collect_output_absolute_paths(e.history_result)
                 register_output_files(paths, job_id=prompt_id)