use subprocess timeout to handle modal failures

commit-0 · wenting-zhao · Dec 7, 2024 · Dec 8, 2024 · Dec 8, 2024 · Dec 8, 2024
commit 7576839df79e3e2d8244dcd2b2bd4bedf7749272
diff --git a/examples/star/run.sh b/examples/star/run.sh
@@ -9,5 +9,6 @@ python examples/star/star.py \
   --iteration 5 \
   --learning_rate 1e-6 \
   --per_device_train_batch_size 1 \
-  --gradient_accumulation_steps 8
+  --gradient_accumulation_steps 8 \
+  --max_workers 64
 
diff --git a/examples/star/utils.py b/examples/star/utils.py
@@ -64,6 +64,7 @@ def execute_tests(
                 input=sample,
                 text=True,
                 capture_output=True,
+                timeout=30,
             ): (i, j)
             for (i, j, instance_id, sample) in tasks
         }
@@ -72,9 +73,13 @@ def execute_tests(
             as_completed(futures), total=len(tasks), desc="Executing tests"
         ):
             i, j = futures[future]
-            result = future.result()
-            stdout = result.stdout
-            exit_code = result.returncode
+            try:
+                result = future.result()
+                stdout = result.stdout
+                exit_code = result.returncode
+            except subprocess.TimeoutExpired:
+                stdout = "Timeout"
+                exit_code = -1
             all_traces[i][j] = stdout
             all_execution_results[i][j] = exit_code