Update failed_run.py

failed_run.py

@@ -30,13 +30,18 @@ args = parser.parse_args()
 dataset = load_dataset("EnergyStarAI/requests_debug", split="test", token=TOKEN).to_pandas()
 
 # Set benchmark to failed
-dataset.loc[dataset["model"].isin([args.model_name]), ['status']] = "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:
-       error_message = file.read()
+        for f in file.readlines():
+            if 'RuntimeError' 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"
+                print("Status set to COMPLETE")
     # Add a new column for the error message if necessary
     if "error_message" not in dataset.columns:
         dataset["error_message"] = ""
@@ -46,4 +51,3 @@ except FileNotFoundError as e:
 
 updated_dataset = Dataset.from_pandas(dataset)
 updated_dataset.push_to_hub("EnergyStarAI/requests_debug", split="test", token=TOKEN)
-print("Status set to FAILED")