diff --git a/NEWS b/NEWS
index a0692ce8cae7923bb8493860b087d51efadc13c4..fe81829a2c16d84c6d9d8d392bbe885f94a8f2c5 100644
--- a/NEWS
+++ b/NEWS
@@ -3,9 +3,12 @@ documents those changes that are of interest to users and admins.
* Changes in SLURM 1.4.0-pre7
=============================
- -- Bug fix for preemtpion with select/cons_res when there are no idle nodes.
+ -- Bug fix for preemption with select/cons_res when there are no idle nodes.
-- Bug fix for use of srun options --exclusive and --cpus-per-task together
for job step resource allocation (tracking of cpus in use was bad).
+ -- Added the srun option --preserve-env to pass the current values of
+ environment variables SLURM_NNODES and SLURM_NPROCS through to the
+ executable, rather than computing them from commandline parameters.
* Changes in SLURM 1.4.0-pre6
=============================
diff --git a/doc/man/man1/srun.1 b/doc/man/man1/srun.1
index d2adb07a3323218b838e441500ef737a7f096154..bd99bd0cbeff82b4de99257ab4233431d341f058 100644
--- a/doc/man/man1/srun.1
+++ b/doc/man/man1/srun.1
@@ -263,6 +263,12 @@ redirected to different locations.
See \fBIO Redirection\fR below for more options.
If the specified file already exists, it will be overwritten.
+.TP
+\fB\-E\fR, \fB\-\-preserve-env\fR
+Pass the current values of environment variables SLURM_NNODES and
+SLURM_NPROCS through to the \fIexecutable\fR, rather than computing them
+from commandline parameters.
+
.TP
\fB\-\-epilog\fR=\fIexecutable\fR
\fBsrun\fR will run \fIexecutable\fR just after the job step completes.
diff --git a/src/api/step_launch.c b/src/api/step_launch.c
index aa4bcf62e6c4b2ad41ddf141ac009a75569339ee..e793c9a8421d3ea0e26452d8a650461eae1a44fa 100644
--- a/src/api/step_launch.c
+++ b/src/api/step_launch.c
@@ -195,7 +195,8 @@ int slurm_step_launch (slurm_step_ctx_t *ctx,
env_array_merge(&env, (const char **)params->env);
}
env_array_for_step(&env, ctx->step_resp,
- ctx->launch_state->resp_port[0]);
+ ctx->launch_state->resp_port[0],
+ params->preserve_env);
env_array_merge(&env, (const char **)mpi_env);
env_array_free(mpi_env);
diff --git a/src/common/env.c b/src/common/env.c
index ec95f12ed207073bb1986a1d2d8837f0ed512161..d111a1c21abf4d9d003f49dcc77ae597def4f3bd 100644
--- a/src/common/env.c
+++ b/src/common/env.c
@@ -956,7 +956,8 @@ env_array_for_batch_job(char ***dest, const batch_job_launch_msg_t *batch,
* overwriting any environment variables of the same name. If the address
* pointed to by "dest" is NULL, memory will automatically be xmalloc'ed.
* The array is terminated by a NULL pointer, and thus is suitable for
- * use by execle() and other env_array_* functions.
+ * use by execle() and other env_array_* functions. If preserve_env is
+ * true, the variables SLURM_NNODES and SLURM_NPROCS remain unchanged.
*
* Sets variables:
* SLURM_STEP_ID
@@ -979,7 +980,8 @@ env_array_for_batch_job(char ***dest, const batch_job_launch_msg_t *batch,
void
env_array_for_step(char ***dest,
const job_step_create_response_msg_t *step,
- uint16_t launcher_port)
+ uint16_t launcher_port,
+ bool preserve_env)
{
char *tmp;
@@ -998,10 +1000,12 @@ env_array_for_step(char ***dest,
/* OBSOLETE, but needed by MPI, do not remove */
env_array_overwrite_fmt(dest, "SLURM_STEPID", "%u", step->job_step_id);
- env_array_overwrite_fmt(dest, "SLURM_NNODES",
- "%hu", step->step_layout->node_cnt);
- env_array_overwrite_fmt(dest, "SLURM_NPROCS",
- "%u", step->step_layout->task_cnt);
+ if (!preserve_env) {
+ env_array_overwrite_fmt(dest, "SLURM_NNODES",
+ "%hu", step->step_layout->node_cnt);
+ env_array_overwrite_fmt(dest, "SLURM_NPROCS",
+ "%u", step->step_layout->task_cnt);
+ }
env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE", "%s", tmp);
env_array_overwrite_fmt(dest, "SLURM_SRUN_COMM_PORT",
"%hu", launcher_port);
diff --git a/src/common/env.h b/src/common/env.h
index 6c577dab44024a1d5e9ab9e247fc5c78404ba9a9..44581572bd249c53f66035d5fcc4a0914dc63a16 100644
--- a/src/common/env.h
+++ b/src/common/env.h
@@ -135,11 +135,12 @@ extern void env_array_for_batch_job(char ***dest,
const char* node_name);
/*
- * Set in "dest the environment variables relevant to a SLURM job step,
+ * Set in "dest" the environment variables relevant to a SLURM job step,
* overwriting any environment variables of the same name. If the address
* pointed to by "dest" is NULL, memory will automatically be xmalloc'ed.
* The array is terminated by a NULL pointer, and thus is suitable for
- * use by execle() and other env_array_* functions.
+ * use by execle() and other env_array_* functions. If preserve_env is
+ * true, the variables SLURM_NNODES and SLURM_NPROCS remain unchanged.
*
* Sets variables:
* SLURM_STEP_ID
@@ -163,7 +164,8 @@ extern void env_array_for_batch_job(char ***dest,
void
env_array_for_step(char ***dest,
const job_step_create_response_msg_t *step,
- uint16_t launcher_port);
+ uint16_t launcher_port,
+ bool preserve_env);
/*
* Return an empty environment variable array (contains a single
diff --git a/src/srun/opt.c b/src/srun/opt.c
index f580524362de9c5928836c9e6568e518ccc68032..a04056d6777d19368a945e6987c36fec87147f04 100644
--- a/src/srun/opt.c
+++ b/src/srun/opt.c
@@ -377,6 +377,7 @@ static void _opt_default()
opt.quit_on_intr = false;
opt.disable_status = false;
opt.test_only = false;
+ opt.preserve_env = false;
opt.quiet = 0;
_verbose = 0;
@@ -679,6 +680,7 @@ static void set_options(const int argc, char **argv)
{"slurmd-debug", required_argument, 0, 'd'},
{"chdir", required_argument, 0, 'D'},
{"error", required_argument, 0, 'e'},
+ {"preserve-env", no_argument, 0, 'E'},
{"geometry", required_argument, 0, 'g'},
{"hold", no_argument, 0, 'H'},
{"input", required_argument, 0, 'i'},
@@ -770,7 +772,7 @@ static void set_options(const int argc, char **argv)
{"wckey", required_argument, 0, LONG_OPT_WCKEY},
{NULL, 0, 0, 0}
};
- char *opt_string = "+aAbB:c:C:d:D:e:g:Hi:IjJ:kKlL:m:n:N:"
+ char *opt_string = "+aAbB:c:C:d:D:e:Eg:Hi:IjJ:kKlL:m:n:N:"
"o:Op:P:qQr:R:st:T:uU:vVw:W:x:XZ";
struct option *optz = spank_option_table_create (long_options);
@@ -851,6 +853,9 @@ static void set_options(const int argc, char **argv)
else
opt.efname = xstrdup(optarg);
break;
+ case (int)'E':
+ opt.preserve_env = true;
+ break;
case (int)'g':
if (verify_geometry(optarg, opt.geometry))
exit(1);
@@ -1878,6 +1883,7 @@ static void _opt_list()
xfree(str);
info("reboot : %s", opt.reboot ? "no" : "yes");
info("rotate : %s", opt.no_rotate ? "yes" : "no");
+ info("preserve_env : %s", tf_(opt.preserve_env));
if (opt.blrtsimage)
info("BlrtsImage : %s", opt.blrtsimage);
@@ -1925,6 +1931,7 @@ static bool _under_parallel_debugger (void)
return (MPIR_being_debugged != 0);
}
+
static void _usage(void)
{
printf(
@@ -1940,7 +1947,7 @@ static void _usage(void)
" [--kill-on-bad-exit] [--propagate[=rlimits] [--comment=name]\n"
" [--cpu_bind=...] [--mem_bind=...] [--network=type]\n"
" [--ntasks-per-node=n] [--ntasks-per-socket=n]\n"
-" [--ntasks-per-core=n] [--mem-per-cpu=MB]\n"
+" [--ntasks-per-core=n] [--mem-per-cpu=MB] [--preserve-env]\n"
#ifdef HAVE_BG /* Blue gene specific options */
" [--geometry=XxYxZ] [--conn-type=type] [--no-rotate] [--reboot]\n"
" [--blrts-image=path] [--linux-image=path]\n"
@@ -2015,6 +2022,7 @@ static void _help(void)
" -L, --licenses=names required license, comma separated\n"
" --checkpoint=time job step checkpoint interval\n"
" --checkpoint-path=dir path to store job step checkpoint image files\n"
+" -E, --preserve-env env vars for node and task counts override command-line flags\n"
#ifdef HAVE_PTY_H
" --pty run task zero in pseudo terminal\n"
#endif
diff --git a/src/srun/opt.h b/src/srun/opt.h
index 1d64ff587eeec87f70cc263a35e0f58841b7e958..c273790ce615c41b27eb3318ac29143eb83c4f9d 100644
--- a/src/srun/opt.h
+++ b/src/srun/opt.h
@@ -165,6 +165,7 @@ typedef struct srun_options {
char *task_epilog; /* --task-epilog= */
char *task_prolog; /* --task-prolog= */
char *licenses; /* --licenses, -L */
+ bool preserve_env; /* --preserve-env */
/* constraint options */
int32_t job_min_cpus; /* --mincpus=n */
diff --git a/src/srun/srun.c b/src/srun/srun.c
index faa28e8dfa7aa03c3545be17411e8f1fd420046b..7709ad8daa5fa06fab9826f1c645524855db5d8e 100644
--- a/src/srun/srun.c
+++ b/src/srun/srun.c
@@ -290,7 +290,6 @@ int srun(int ac, char **av)
/*
* Enhance environment for job
*/
- env->nprocs = opt.nprocs;
env->cpus_per_task = opt.cpus_per_task;
if (opt.ntasks_per_node != NO_VAL)
env->ntasks_per_node = opt.ntasks_per_node;
@@ -315,7 +314,6 @@ int srun(int ac, char **av)
&tasks);
env->select_jobinfo = job->select_jobinfo;
- env->nhosts = job->nhosts;
env->nodelist = job->nodelist;
env->task_count = _uint16_array_to_str(
job->nhosts, tasks);
@@ -377,6 +375,7 @@ int srun(int ac, char **av)
launch_params.ntasks_per_core = opt.ntasks_per_core;
launch_params.task_dist = opt.distribution;
launch_params.ckpt_path = xstrdup(opt.ckpt_path);
+ launch_params.preserve_env = opt.preserve_env;
/* job structure should now be filled in */
_setup_signals();