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"