Add more...

experiments/eval/scripts/timing.py

@@ -44,12 +44,13 @@ def get_outdirs(jobid: str):
     return sorted(Path("/p/scratch/cellsinsilico/paul/nastja-out").glob(f"*{jobid}*"))
 
 
-def load_array_mean_timings(jobid: str):
-    mts = [
-        load_mean_timings(outdir_path)
-        for outdir_path
-        in get_outdirs(jobid)
-    ]
+def load_array_mean_timings(jobid: str, excluded_array_indices):
+    mts = []
+    for outdir_path in get_outdirs(jobid):
+        if any(str(outdir_path).endswith(str(i)) for i in excluded_array_indices):
+            print(f"Not loading timings for {outdir_path} because it was excluded.", file=sys.stderr)
+            continue
+        mts.append(load_mean_timings(outdir_path))
 
     return pandas.DataFrame(list(mts), columns=mts[0].index)
 
@@ -81,12 +82,28 @@ if __name__ == "__main__":
     for i, jobid in enumerate(args.jobid, 1):
         print(f"({i:2}/{len(args.jobid):2}) Loading accounting data for {jobid}", file=sys.stderr)
         accounting_data = get_accounting_data(jobid)
-        array_mean_timings = load_array_mean_timings(jobid)
+
+        jobs = []
+        excluded_array_indices = []
+        for array_job in accounting_data["jobs"]:
+            # Get metadata related to array
+            array_main_job = array_job["array"]["job_id"]
+            array_index = array_job["array"]["task_id"]
+            # The last step is the actual job we want the data for
+            # The steps before set up cluster etc.
+            last_step = array_job["steps"][-1]
+            if last_step["state"] != "COMPLETED":
+                print(f"WARNING: {array_main_job}.{array_index} has state {last_step['state']}, excluding it from measurements", file=sys.stderr)
+                excluded_array_indices.append(array_index)
+                continue
+            jobs.append(last_step)
+
+        array_mean_timings = load_array_mean_timings(jobid, excluded_array_indices)
 
         results.append({
             "jobid": jobid,
             "means": array_mean_timings.mean().to_dict(),
             "stds": array_mean_timings.std().to_dict(),
-            "accounting": [array_job["steps"][-1] for array_job in accounting_data["jobs"]]
+            "accounting": jobs
         })
     print(json.dumps(results, indent=2 if args.prettify else None))