[2685] Plot training for LSF

src/weathergen/utils/plot_training.py

@@ -10,6 +10,7 @@
 import argparse
 import logging
 import pdb
+import shutil
 import subprocess
 import sys
 import traceback
@@ -30,6 +31,37 @@
 PLOT_DPI_VALUE = 150
 
 
+def check_active_runs(runs_ids):
+    """
+    Check if the specified runs are active.
+
+    Parameters
+    ----------
+    runs_ids : dict
+        Dictionary of run IDs to check.
+
+    Returns
+    -------
+    list
+        List of booleans indicating if each run is active.
+    """
+    if shutil.which("squeue"):
+        sq_arg = "--format='%.18i %.9P %.30j %.8u %.8T %.10M %.9l %.6D %R' --me"
+        ret = subprocess.run(["squeue", sq_arg], capture_output=True)
+        running_state = "RUNNING"
+    else:
+        ret = subprocess.run(["bjobs", "-o", "jobid stat job_name"], capture_output=True)
+        running_state = "RUN"
+
+    lines = str(ret.stdout).split("\\n")
+    runs_active = [
+        any([run_id in line and running_state in line for line in lines[1:]])
+        for run_id in runs_ids.keys()
+    ]
+    return runs_active
+
+
+####################################################################################################
 def _add_legend(
     labels,
     outside: bool,
@@ -908,13 +940,7 @@ def plot_train(args=None):
     ]
 
     # determine which runs are still alive (as a process, though they might hang internally)
-    sq_arg = "--format='%.18i %.9P %.30j %.8u %.8T %.10M %.9l %.6D %R' --me"
-    ret = subprocess.run(["squeue", sq_arg], capture_output=True)
-    lines = str(ret.stdout).split("\\n")
-    runs_active = [
-        any([run_id in line and "RUNNING" in line for line in lines[1:]])
-        for run_id in runs_ids.keys()
-    ]
+    runs_active = check_active_runs(runs_ids)
 
     x_scale_log = args.log_x