diff --git a/doc/html/prolog_epilog.shtml b/doc/html/prolog_epilog.shtml index a75ea59479d..308d19efa3c 100644 --- a/doc/html/prolog_epilog.shtml +++ b/doc/html/prolog_epilog.shtml @@ -483,6 +483,10 @@ Available in SrunProlog, TaskProlog, SrunEpilog and TaskEpilog.
  • SLURM_JOB_ACCOUNT Account name used for the job.
  • +
  • SLURM_JOB_ALLOC_MEM_PER_NODE +Effective memory allocated to the job on the current slurmd node, in MiB. +Available in Prolog.
  • +
  • SLURM_JOB_COMMENT Comment added to the job. Available in Prolog, PrologSlurmctld, Epilog and EpilogSlurmctld.
  • diff --git a/src/plugins/prep/script/prep_script_slurmd.c b/src/plugins/prep/script/prep_script_slurmd.c index 45b7c2132f5..1d245a1a407 100644 --- a/src/plugins/prep/script/prep_script_slurmd.c +++ b/src/plugins/prep/script/prep_script_slurmd.c @@ -36,6 +36,7 @@ #include "config.h" #include +#include #include #include #include @@ -65,6 +66,8 @@ slurmd_conf_t *conf = NULL; static char **_build_env(job_env_t *job_env, slurm_cred_t *cred, bool is_epilog); +static bool _get_job_alloc_mem_per_node(slurm_cred_arg_t *cred_arg, + uint64_t *alloc_mem); static int _run_spank_job_script(const char *mode, char **env, uint32_t job_id, bool is_epilog); @@ -220,6 +223,43 @@ extern int slurmd_script(job_env_t *job_env, slurm_cred_t *cred, return rc; } +static bool _get_job_alloc_mem_per_node(slurm_cred_arg_t *cred_arg, + uint64_t *alloc_mem) +{ + int node_id, rep_idx = -1; + + if (!cred_arg->job_hostlist || !cred_arg->job_mem_alloc || + !cred_arg->job_mem_alloc_rep_count || !cred_arg->job_mem_alloc_size) + goto skip; + + if ((node_id = nodelist_find(cred_arg->job_hostlist, conf->node_name)) < + 0) + goto skip; + + for (uint32_t i = 0; i < cred_arg->job_mem_alloc_size; i++) + if (!cred_arg->job_mem_alloc_rep_count[i]) + goto skip; + + /* Batch credentials only carry the batch node's memory allocation. */ + if (cred_arg->step_id.step_id == SLURM_BATCH_SCRIPT) + rep_idx = 0; + else + rep_idx = slurm_get_rep_count_inx( + cred_arg->job_mem_alloc_rep_count, + cred_arg->job_mem_alloc_size, node_id); + if ((rep_idx < 0) || + ((uint32_t) rep_idx >= cred_arg->job_mem_alloc_size)) + goto skip; + + *alloc_mem = cred_arg->job_mem_alloc[rep_idx]; + return true; + +skip: + debug2("%s: unable to get SLURM_JOB_ALLOC_MEM_PER_NODE on node %s", + __func__, conf->node_name); + return false; +} + /* NOTE: call env_array_free() to free returned value */ static char **_build_env(job_env_t *job_env, slurm_cred_t *cred, bool is_epilog) @@ -292,10 +332,15 @@ static char **_build_env(job_env_t *job_env, slurm_cred_t *cred, if (cred) { slurm_cred_arg_t *cred_arg = slurm_cred_get_args(cred); + uint64_t alloc_mem; if (cred_arg->job_account) setenvf(&env, "SLURM_JOB_ACCOUNT", "%s", cred_arg->job_account); + if (!is_epilog && + _get_job_alloc_mem_per_node(cred_arg, &alloc_mem)) + setenvf(&env, "SLURM_JOB_ALLOC_MEM_PER_NODE", + "%" PRIu64, alloc_mem); if (cred_arg->job_comment) setenvf(&env, "SLURM_JOB_COMMENT", "%s", cred_arg->job_comment); diff --git a/testsuite/expect/test31.3 b/testsuite/expect/test31.3 index fd8023c0bdf..67c6a6756f9 100755 --- a/testsuite/expect/test31.3 +++ b/testsuite/expect/test31.3 @@ -79,7 +79,7 @@ reconfigure -fail # Test whether env vars exist in Prolog and Epilog log_info "Checking for regular env vars" -set job_id [submit_job -fail "-t1 -N1 -w$test_node -o/dev/null --comment=foo --extra=bar --wrap='hostname'"] +set job_id [submit_job -fail "-t1 -N1 -w$test_node --mem=1 -o/dev/null --comment=foo --extra=bar --wrap='hostname'"] wait_for_job $job_id "DONE" @@ -104,6 +104,8 @@ check_epilog "SLURM_JOB_UID=" check_prolog "SLURM_JOB_USER=" check_epilog "SLURM_JOB_USER=" +check_prolog "SLURM_JOB_ALLOC_MEM_PER_NODE=1" + check_prolog "SLURM_SCRIPT_CONTEXT=prolog_slurmd" check_epilog "SLURM_SCRIPT_CONTEXT=epilog_slurmd"