Update slurm features

.github/ci/example_1b_apptainer_config.yaml

@@ -11,6 +11,11 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
   copy_to_rundir:

.github/ci/example_1b_conda_config.yaml

@@ -11,6 +11,11 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
     copy_to_rundir:

.github/ci/example_2a_apptainer_config.yaml

@@ -11,6 +11,11 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
   copy_to_rundir:

.github/ci/example_2a_conda_config.yaml

@@ -11,6 +11,11 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
   copy_to_rundir:

.github/ci/example_3_apptainer_config.yaml

@@ -12,6 +12,23 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  convert-clim:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-events:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-xml:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
   copy_to_rundir:

.github/ci/example_3_conda_config.yaml

@@ -12,6 +12,23 @@ pecan_parallelism_mode: slurm-sbatch
 sbatch_additional_arguments: "-p node"
 workflow_parallelism_mode: slurm-srun
 srun_additional_arguments: "-p node --mem=0"
+step_dispatch_configuration:
+  convert-clim:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-events:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
+  build-xml:
+    slurm: true
+    method: srun
+    additional_arguments: "-p node --mem=0"
 aws_profile: magic
 external_paths:
   copy_to_rundir:
@@ -29,3 +46,5 @@ paths:
   irrigation_info_path: "data_raw/management/irrigation/v1.1/"
   fertilization_info_path: "data_raw/management/fertilization/v1.0/"
   ncc_info_path: "data_raw/management/ncc/v1.0/"
+  site_era5_path: "data_raw/ERA5_CA_nc"
+  landtrendr_agb_file: "data_raw/IC/landtrendr_2016_biomass_by_dwr_parcel.parquet"

examples/1b_statewide_woody/example_user_config.yaml

@@ -30,6 +30,14 @@ workflow_parallelism_mode: slurm-srun
 # separate slurm_partition key -- it's just another srun flag).
 srun_additional_arguments: "--mem=0"
 
+# Steps the manifest marks slurm: true (here: build-ic) must have a matching
+# entry below or the CLI fails fast.
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
+
 # AWS CLI named profile for S3 access (get-demo-data command).
 aws_profile: magic
 

examples/2a_grass/example_user_config.yaml

@@ -56,5 +56,13 @@ workflow_parallelism_mode: slurm-srun
 # separate slurm_partition key -- it's just another srun flag).
 srun_additional_arguments: "--mem=0"
 
+# Steps the manifest marks slurm: true (here: build-ic) must have a matching
+# entry below or the CLI fails fast.
+step_dispatch_configuration:
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
+
 # AWS CLI named profile for S3 access (get-demo-data command).
 aws_profile: magic

examples/3_rowcrop/example_user_config.yaml

@@ -36,6 +36,26 @@ workflow_parallelism_mode: slurm-srun
 # separate slurm_partition key -- pass -p <partition> here if your cluster needs one.
 srun_additional_arguments: "--mem=0"
 
+# Steps the manifest marks slurm: true for prepare-example-3 (convert-clim,
+# build-ic, build-events, build-xml) must each have a matching entry below
+# or the CLI fails fast.
+step_dispatch_configuration:
+  convert-clim:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
+  build-ic:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
+  build-events:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
+  build-xml:
+    slurm: true
+    method: srun
+    additional_arguments: "--mem=0"
 
 # Pecan parallelism
 pecan_parallelism_mode: slurm-sbatch
@@ -50,8 +70,8 @@ external_paths:
     template_file: "examples/3_rowcrop/template.xml"
     pft_map: "examples/3_rowcrop/crop2pft.csv"
   link_in_place:
-    site_era5_path: "data_raw/ERA5_CA_nc"
-    landtrendr_agb_file: "data_raw/IC/landtrendr_2016_biomass_by_dwr_parcel.parquet"
+  #  site_era5_path: "data_raw/ERA5_CA_nc"
+  #  landtrendr_agb_file: "data_raw/IC/landtrendr_2016_biomass_by_dwr_parcel.parquet"
 
 paths:
   phenology_info_path: "data_raw/management/phenology/v1.0/"
@@ -61,3 +81,5 @@ paths:
   irrigation_info_path: "data_raw/management/irrigation/v1.1/"
   fertilization_info_path: "data_raw/management/fertilization/v1.0/"
   ncc_info_path: "data_raw/management/ncc/v1.0/"
+  site_era5_path: "data_raw/ERA5_CA_nc"
+  landtrendr_agb_file: "data_raw/IC/landtrendr_2016_biomass_by_dwr_parcel.parquet"

magic-ensemble

@@ -296,13 +296,14 @@ check_slurm_available() {
 # for srun). ---
 build_srun_prefix() {
   local job_name="$1"
+  local extra_args="$2"
   SRUN_PREFIX=(srun)
   if [[ -n "$job_name" ]]; then
     SRUN_PREFIX+=(--job-name "Magic-${job_name}")
   fi
-  if [[ -n "$srun_additional_arguments" && "$srun_additional_arguments" != "null" ]]; then
+  if [[ -n "$extra_args" && "$extra_args" != "null" ]]; then
     local -a extra
-    read -r -a extra <<< "$srun_additional_arguments"
+    read -r -a extra <<< "$extra_args"
     SRUN_PREFIX+=("${extra[@]}")
   fi
 }
@@ -393,6 +394,69 @@ validate_slurm_steps() {
   done
 }
 
+# --- Validate and resolve the user config's step_dispatch_configuration map
+# against the manifest's per-step 'slurm' flags. A step marked slurm: true
+# in the manifest requires a matching entry here, or the CLI fails fast
+# (before any step runs); a step marked slurm: false ignores any entry
+# present, but warns, since it indicates user intent that won't take
+# effect. Populates STEP_SLURM_USER_ENABLED[] (the entry's own opt-out) and
+# STEP_SLURM_PARAMS[] (resolved additional_arguments, falling back to the
+# global srun_additional_arguments -- with a warning -- when the entry
+# doesn't supply its own), both index-aligned with STEPS/STEP_NAMES/STEP_SLURM. ---
+validate_step_dispatch_configuration() {
+  STEP_SLURM_USER_ENABLED=()
+  STEP_SLURM_PARAMS=()
+  local i name entry enabled method params
+  for i in "${!STEPS[@]}"; do
+    name="${STEP_NAMES[i]}"
+    entry=$(yq eval ".step_dispatch_configuration[\"${name}\"]" "$CONFIG_FILE" 2>/dev/null)
+
+    if [[ "${STEP_SLURM[i]}" != "true" ]]; then
+      if [[ -n "$entry" && "$entry" != "null" ]]; then
+        echo "magic-ensemble: WARNING: step '${name}' has a step_dispatch_configuration entry in your config, but the manifest marks it slurm: false -- running it without Slurm, as the manifest requires." >&2
+      fi
+      STEP_SLURM_USER_ENABLED+=("false")
+      STEP_SLURM_PARAMS+=("")
+      continue
+    fi
+
+    if [[ -z "$entry" || "$entry" == "null" ]]; then
+      echo "magic-ensemble: step '${name}' is marked slurm: true in the manifest but has no step_dispatch_configuration entry in your config ($CONFIG_FILE). Add one (see workflow_manifest.yaml / magic-ensemble-DEVELOPERS.md) or the manifest cannot dispatch this step." >&2
+      exit 1
+    fi
+
+    enabled=$(yq eval ".step_dispatch_configuration[\"${name}\"].slurm" "$CONFIG_FILE" 2>/dev/null)
+    [[ -z "$enabled" || "$enabled" == "null" ]] && enabled="true"
+    STEP_SLURM_USER_ENABLED+=("$enabled")
+
+    if [[ "$enabled" != "true" ]]; then
+      STEP_SLURM_PARAMS+=("")
+      continue
+    fi
+
+    method=$(yq eval ".step_dispatch_configuration[\"${name}\"].method" "$CONFIG_FILE" 2>/dev/null)
+    [[ -z "$method" || "$method" == "null" ]] && method="srun"
+    case "$method" in
+      srun) ;;
+      sbatch)
+        echo "magic-ensemble: step '${name}' requests step_dispatch_configuration method 'sbatch', but sbatch step-wrapping is not yet implemented -- only 'srun' is supported. Use method: srun or remove the method key." >&2
+        exit 1
+        ;;
+      *)
+        echo "magic-ensemble: step '${name}' has unknown step_dispatch_configuration method '${method}'. Valid options: srun" >&2
+        exit 1
+        ;;
+    esac
+
+    params=$(yq eval ".step_dispatch_configuration[\"${name}\"].additional_arguments" "$CONFIG_FILE" 2>/dev/null)
+    if [[ -z "$params" || "$params" == "null" ]]; then
+      echo "magic-ensemble: WARNING: step '${name}' has no additional_arguments in its step_dispatch_configuration entry -- falling back to the global srun_additional_arguments." >&2
+      params="$srun_additional_arguments"
+    fi
+    STEP_SLURM_PARAMS+=("$params")
+  done
+}
+
 # --- Decide whether step i of the current command should be srun-wrapped:
 # manifest opts in, user config hasn't force-disabled it, and srun is
 # actually available. Sets STEP_SLURM_ARGS to () or ("--slurm" "<step name>")
@@ -412,6 +476,11 @@ compute_slurm_arg() {
 
   [[ "${STEP_SLURM[i]}" != "true" ]] && return
 
+  if [[ "${STEP_SLURM_USER_ENABLED[i]}" != "true" ]]; then
+    log_info "step '${STEP_NAMES[i]}': slurm disabled via step_dispatch_configuration in your config -- running it unwrapped, as configured."
+    return
+  fi
+
   if [[ "$workflow_parallelism_mode" == "local" ]]; then
     echo "magic-ensemble: WARNING: step '${STEP_NAMES[i]}' is marked slurm: true in the manifest, but workflow_parallelism_mode is set to 'local' in your config -- running it unwrapped, as configured." >&2
     return
@@ -422,7 +491,7 @@ compute_slurm_arg() {
     return
   fi
 
-  STEP_SLURM_ARGS=(--slurm "${STEP_NAMES[i]}")
+  STEP_SLURM_ARGS=(--slurm "${STEP_NAMES[i]}" "${STEP_SLURM_PARAMS[i]}")
 }
 
 # --- R library check for step at index i (reads r_libraries from manifest step; skip if empty or .sh) ---
@@ -463,19 +532,21 @@ run_script() {
   local use_apptainer=0
   local use_slurm=0
   local slurm_job_name=""
+  local slurm_params=""
   local script_cwd="$REPO_ROOT"
   while [[ $# -gt 0 ]]; do
     case "${1:-}" in
       '') shift ;;
       --apptainer) use_apptainer=1; shift ;;
       --slurm)
-        if [[ $# -lt 2 ]]; then
-          echo "magic-ensemble: run_script --slurm requires STEP_NAME." >&2
+        if [[ $# -lt 3 ]]; then
+          echo "magic-ensemble: run_script --slurm requires STEP_NAME PARAMS." >&2
           exit 1
         fi
         use_slurm=1
         slurm_job_name="$2"
-        shift 2
+        slurm_params="$3"
+        shift 3
         ;;
       --cwd)
         if [[ $# -lt 2 ]]; then
@@ -498,7 +569,7 @@ run_script() {
 
   local -a SRUN_PREFIX=()
   if [[ $use_slurm -eq 1 ]]; then
-    build_srun_prefix "$slurm_job_name"
+    build_srun_prefix "$slurm_job_name" "$slurm_params"
   fi
 
   if [[ $use_apptainer -eq 1 ]]; then
@@ -535,9 +606,11 @@ run_shell_script() {
   local step_index="${2:-0}"
   local use_slurm=0
   local slurm_job_name=""
+  local slurm_params=""
   if [[ "${3:-}" == "--slurm" ]]; then
     use_slurm=1
     slurm_job_name="${4:-}"
+    slurm_params="${5:-}"
   fi
   local script_path="${REPO_ROOT}/${script}"
   if [[ ! -f "$script_path" ]]; then
@@ -547,7 +620,7 @@ run_shell_script() {
 
   local -a SRUN_PREFIX=()
   if [[ $use_slurm -eq 1 ]]; then
-    build_srun_prefix "$slurm_job_name"
+    build_srun_prefix "$slurm_job_name" "$slurm_params"
   fi
 
   local script_basename="${script##*/}"
@@ -737,6 +810,7 @@ run_manifest_command() {
   { # block for output redirection
     get_steps_array
     validate_slurm_steps
+    validate_step_dispatch_configuration
     validate_pecan_parallelism_mode
     local apptainer_arg=""
     if [[ $use_apptainer -eq 1 ]]; then
@@ -790,6 +864,7 @@ run_manifest_command() {
 run_get_demo_data() {
   get_steps_array
   validate_slurm_steps
+  validate_step_dispatch_configuration
   check_aws
   for i in "${!STEPS[@]}"; do
     check_r_libs_for_step "$i"
@@ -809,6 +884,7 @@ run_run_ensembles() {
   { # block for output redirection
     get_steps_array
     validate_slurm_steps
+    validate_step_dispatch_configuration
 
     if [[ $use_apptainer -eq 1 ]]; then
       ensure_apptainer_available