diff --git a/lib/ramble/ramble/application.py b/lib/ramble/ramble/application.py index cd96467d9..d71a0f2ee 100644 --- a/lib/ramble/ramble/application.py +++ b/lib/ramble/ramble/application.py @@ -66,7 +66,7 @@ experiment_status = Enum( "experiment_status", - ["UNKNOWN", "SETUP", "RUNNING", "COMPLETE", "SUCCESS", "FAILED", "CANCELLED"], + ["UNKNOWN", "SETUP", "SUBMITTED", "RUNNING", "COMPLETE", "SUCCESS", "FAILED", "CANCELLED"], ) _NULL_CONTEXT = "null" diff --git a/var/ramble/repos/builtin/workflow_managers/slurm/batch_query.tpl b/var/ramble/repos/builtin/workflow_managers/slurm/batch_query.tpl index ff462102b..8995edd75 100644 --- a/var/ramble/repos/builtin/workflow_managers/slurm/batch_query.tpl +++ b/var/ramble/repos/builtin/workflow_managers/slurm/batch_query.tpl @@ -7,19 +7,15 @@ if [ -z "${job_id:-}" ]; then exit 1 fi +# Set up the status_map mapping between +# sacct/squeue status to ramble counterpart +{declare_status_map} + status=$(squeue -h -o "%t" -j "${job_id}" 2>/dev/null) if [ -z "$status" ]; then status=$(sacct -j "${job_id}" -o state -X -n | xargs) fi if [ ! -z "$status" ]; then - # Define a mapping between sacct/squeue status to ramble counterpart - declare -A status_map - status_map["PD"]="SETUP" - status_map["R"]="RUNNING" - status_map["CF"]="SETUP" - status_map["CG"]="COMPLETE" - status_map["COMPLETED"]="COMPLETE" - status_map["CANCELLED+"]="CANCELLED" if [ -v status_map["$status"] ]; then status=${status_map["$status"]} fi diff --git a/var/ramble/repos/builtin/workflow_managers/slurm/workflow_manager.py b/var/ramble/repos/builtin/workflow_managers/slurm/workflow_manager.py index d5cdd5af5..105a1c855 100644 --- a/var/ramble/repos/builtin/workflow_managers/slurm/workflow_manager.py +++ b/var/ramble/repos/builtin/workflow_managers/slurm/workflow_manager.py @@ -16,7 +16,7 @@ # Mapping from squeue/sacct status to Ramble status _STATUS_MAP = { - "PD": "SETUP", + "PD": "SUBMITTED", "R": "RUNNING", "CF": "SETUP", "CG": "COMPLETE", @@ -26,6 +26,14 @@ } +def _declare_status_map(): + """A utility to convert the `_STATUS_MAP` into a bash array""" + entries = ["declare -A status_map"] + for k, v in _STATUS_MAP.items(): + entries.append(f'status_map["{k}"]="{v}"') + return "\n".join(entries) + + class Slurm(WorkflowManagerBase): """Slurm workflow manager""" @@ -77,7 +85,10 @@ def __init__(self, file_path): ) register_template( - name="batch_query", src_name="batch_query.tpl", dest_name="batch_query" + name="batch_query", + src_name="batch_query.tpl", + dest_name="batch_query", + extra_vars={"declare_status_map": _declare_status_map()}, ) register_template(