From 26aa72db4209dc0d3ce730350fa8d2941fb27487 Mon Sep 17 00:00:00 2001 From: Pierluigi Cosi Date: Tue, 28 Jul 2026 14:59:28 +0200 Subject: [PATCH 1/2] fix active runs --- src/weathergen/utils/plot_training.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/weathergen/utils/plot_training.py b/src/weathergen/utils/plot_training.py index 5e5f9478e..2d149bbc0 100644 --- a/src/weathergen/utils/plot_training.py +++ b/src/weathergen/utils/plot_training.py @@ -18,6 +18,7 @@ import matplotlib.pyplot as plt import numpy as np import yaml +import shutil import weathergen.common.config as config from weathergen.train.utils import TRAIN @@ -908,13 +909,18 @@ 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() - ] + 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() + ] x_scale_log = args.log_x From 93f6e592f42bb5e358dd0107fe028b512c28c100 Mon Sep 17 00:00:00 2001 From: Pierluigi Cosi Date: Tue, 1 Sep 2026 11:30:49 +0200 Subject: [PATCH 2/2] separate function and lint --- src/weathergen/utils/plot_training.py | 46 +++++++++++++++++++-------- 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/src/weathergen/utils/plot_training.py b/src/weathergen/utils/plot_training.py index 2d149bbc0..e43b84d9e 100644 --- a/src/weathergen/utils/plot_training.py +++ b/src/weathergen/utils/plot_training.py @@ -10,6 +10,7 @@ import argparse import logging import pdb +import shutil import subprocess import sys import traceback @@ -18,7 +19,6 @@ import matplotlib.pyplot as plt import numpy as np import yaml -import shutil import weathergen.common.config as config from weathergen.train.utils import TRAIN @@ -31,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, @@ -909,18 +940,7 @@ def plot_train(args=None): ] # determine which runs are still alive (as a process, though they might hang internally) - 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() - ] + runs_active = check_active_runs(runs_ids) x_scale_log = args.log_x