Adding handling to say that something takes too long

failed_run.py

@@ -23,27 +23,39 @@ parser.add_argument(
     required=True,
     help="Model to benchmark.",
 )
+parser.add_argument(
+    "--reason",
+    default=None,
+    type=str,
+    required=True,
+    help="Reason for failure -- to update in the requests file",
+)
 
 args = parser.parse_args()
 
 # Updating request
 dataset = load_dataset("AIEnergyScore/requests_debug", split="test", token=TOKEN).to_pandas()
 
-# Set benchmark to failed
-# TODO: This doesn't have to be try-except, we could actually check if the file is there.
-try:
-    # Read error message
-    with open(f"{args.run_dir}/error.log", 'r') as file:
-        for f in file.readlines():
-            if 'Traceback (most recent call last):' in f:
-                error_message = f
-                dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "FAILED"
-                print("Status set to FAILED")
-            else:
-                dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "COMPLETED"
-    # Add a new column for the error message if necessary
-except FileNotFoundError as e:
-    print(f"Could not find {args.run_dir}/error.log")
+## Set benchmark to failed
+
+# If we have a custom reason for failure, add that instead of generic FAILED.
+if args.reason:
+    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = args.reason
+else:
+    # TODO: This doesn't have to be try-except, we could actually check if the file is there...
+    try:
+        # Read error message
+        with open(f"{args.run_dir}/error.log", 'r') as file:
+            for f in file.readlines():
+                if 'Traceback (most recent call last):' in f:
+                    error_message = f
+                    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "FAILED"
+                    print("Status set to FAILED")
+                else:
+                    dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "COMPLETED"
+        # Add a new column for the error message if necessary
+    except FileNotFoundError as e:
+        print(f"Could not find {args.run_dir}/error.log")
     
 updated_dataset = Dataset.from_pandas(dataset)
 updated_dataset.push_to_hub("AIEnergyScore/requests_debug", split="test", token=TOKEN)

parse_requests.py

@@ -1,8 +1,9 @@
 import os
 from datasets import load_dataset, Dataset
 
+skip_statuses = ['COMPLETED', 'FAILED', 'RUNNING']
 TOKEN = os.environ.get("DEBUG")
-requests_dataset = load_dataset("EnergyStarAI/requests_debug", split="test")
+requests_dataset = load_dataset("AIEnergyScore/requests_debug", split="test")
 
 def normalize_task(task):
     # Makes assumption about how the task names are being written, and called.
@@ -13,4 +14,7 @@ requests_dset = requests_dataset.to_pandas()
 
 for model, task in requests_dset[['model','task']].loc[requests_dset['status'] == 'PENDING'].to_dict(orient= 'split', index=False)['data']:
     print("%s,%s" % (model, normalize_task(task)))
-                
+
+# Custom errors we will rerun.
+for model, task in requests_dset[['model','task']].loc[~requests_dset['status'].isin(skip_statuses)].to_dict(orient= 'split', index=False)['data']:
+    print("%s,%s" % (model, normalize_task(task)))