From fd78d61bcc781bd30312fd93907b9a4fe154e410 Mon Sep 17 00:00:00 2001
From: "Bin Bao (Meta Employee)" <binbao@meta.com>
Date: Sun, 10 Dec 2023 08:34:35 -0800
Subject: [PATCH] Fix a missing write_csv_when_exception problem (#115370)

Summary:
Fix a problem shown in https://github.com/pytorch/pytorch/actions/runs/7124839624/job/19400589129 when a model times out.

X-link: https://github.com/pytorch/pytorch/pull/115370
Approved by: https://github.com/eellison

Reviewed By: atalman

Differential Revision: D51999570

Pulled By: desertfire

fbshipit-source-id: 77d4cc1ccce63f8d3409914094478e93964e61d4
---
 userbenchmark/dynamo/dynamobench/common.py | 43 +++++++++-------------
 1 file changed, 18 insertions(+), 25 deletions(-)

diff --git a/userbenchmark/dynamo/dynamobench/common.py b/userbenchmark/dynamo/dynamobench/common.py
index fdeb4e218b..c26d67fe77 100644
--- a/userbenchmark/dynamo/dynamobench/common.py
+++ b/userbenchmark/dynamo/dynamobench/common.py
@@ -3145,6 +3145,24 @@ def main(runner, original_dir=None, args=None):
         process_entry(0, runner, original_dir, args)
 
 
+def write_csv_when_exception(name: str, status: str, device=None):
+    print(status)
+    placeholder_batch_size = 0
+    devices = [device] if device is not None else args.devices
+    if args.accuracy:
+        headers = ["dev", "name", "batch_size", "accuracy"]
+        rows = [[device, name, placeholder_batch_size, status] for device in devices]
+    elif args.performance:
+        headers = ["dev", "name", "batch_size", "speedup", "abs_latency"]
+        rows = [[device, name, placeholder_batch_size, 0.0, 0.0] for device in devices]
+    else:
+        headers = []
+        rows = [[device, name, placeholder_batch_size, 0.0] for device in devices]
+
+    for row in rows:
+        output_csv(output_filename, headers, row)
+
+
 def run(runner, args, original_dir=None):
     # Pass the parsed args object to benchmark runner object
     runner.args = args
@@ -3500,31 +3518,6 @@ def run(runner, args, original_dir=None):
             # Go back to main branch
             repo.git.checkout(main_branch)
     elif args.only:
-
-        def write_csv_when_exception(name: str, status: str, device=None):
-            print(status)
-            placeholder_batch_size = 0
-            devices = [device] if device is not None else args.devices
-            if args.accuracy:
-                headers = ["dev", "name", "batch_size", "accuracy"]
-                rows = [
-                    [device, name, placeholder_batch_size, status] for device in devices
-                ]
-            elif args.performance:
-                headers = ["dev", "name", "batch_size", "speedup", "abs_latency"]
-                rows = [
-                    [device, name, placeholder_batch_size, 0.0, 0.0]
-                    for device in devices
-                ]
-            else:
-                headers = []
-                rows = [
-                    [device, name, placeholder_batch_size, 0.0] for device in devices
-                ]
-
-            for row in rows:
-                output_csv(output_filename, headers, row)
-
         model_name = args.only
         for device in args.devices:
             batch_size = args.batch_size