diff --git a/src/srun/allocate.c b/src/srun/allocate.c index f3ac7987f535cf2d6e07d0616e6bcab2e2a9678b..be88b964d5ed15c19579c6854d17d608cf484dba 100644 --- a/src/srun/allocate.c +++ b/src/srun/allocate.c @@ -404,11 +404,6 @@ job_desc_msg_create_from_opts (char *script) */ xassert (opt.batch); - if (opt.nprocs_set - && setenvf("SLURM_NPROCS=%u", opt.nprocs)) { - error("Unable to set SLURM_NPROCS environment " - "variable"); - } j->environment = environ; j->env_size = envcount (environ); j->script = script; diff --git a/src/srun/srun.c b/src/srun/srun.c index 530234bc22ded8e498e6300f7682d8d9deadce71..2a04f684f86ee88cf1251e923eaf68ab185efc1a 100644 --- a/src/srun/srun.c +++ b/src/srun/srun.c @@ -96,7 +96,7 @@ static int _run_batch_job (void); static void _run_job_script(job_t *job); static int _set_batch_script_env(job_t *job); static int _set_rlimit_env(void); -static char *_sprint_task_cnt(job_t *job); +static char *_task_count_string(job_t *job); static void _switch_standalone(job_t *job); static int _become_user (void); @@ -214,7 +214,7 @@ int srun(int ac, char **av) setenvf("SLURM_JOBID=%u", job->jobid); setenvf("SLURM_NPROCS=%d", opt.nprocs); setenvf("SLURM_NNODES=%d", job->nhosts); - setenvf("SLURM_TASKS_PER_NODE=%s", (task_cnt = _sprint_task_cnt(job))); + setenvf("SLURM_TASKS_PER_NODE=%s", task_cnt = _task_count_string (job)); setenvf("SLURM_DISTRIBUTION=%s", format_distribution_t (opt.distribution)); @@ -278,12 +278,12 @@ int srun(int ac, char **av) } static char * -_sprint_task_cnt(job_t *job) +_task_count_string (job_t *job) { int i, last_val, last_cnt; - char *task_str = xstrdup(""); char tmp[16]; - + char *str = xstrdup (""); + last_val = job->ntask[0]; last_cnt = 1; for (i=1; i<job->nhosts; i++) { @@ -294,7 +294,7 @@ _sprint_task_cnt(job_t *job) sprintf(tmp, "%d(x%d),", last_val, last_cnt); else sprintf(tmp, "%d,", last_val); - xstrcat(task_str, tmp); + xstrcat(str, tmp); last_val = job->ntask[i]; last_cnt = 1; } @@ -303,8 +303,8 @@ _sprint_task_cnt(job_t *job) sprintf(tmp, "%d(x%d)", last_val, last_cnt); else sprintf(tmp, "%d", last_val); - xstrcat(task_str, tmp); - return task_str; + xstrcat(str, tmp); + return (str); } static void @@ -371,6 +371,8 @@ _run_batch_job(void) return SLURM_ERROR; } + _set_batch_script_env (NULL); + if (!(req = job_desc_msg_create_from_opts (script))) fatal ("Unable to create job request"); @@ -564,44 +566,22 @@ static int _set_batch_script_env(job_t *job) { int rc = SLURM_SUCCESS; - char *dist = NULL, *task_cnt; + char *dist = NULL; + char *p; struct utsname name; - if (job->jobid > 0) { - if (setenvf("SLURM_JOBID=%u", job->jobid)) { - error("Unable to set SLURM_JOBID environment"); - rc = SLURM_FAILURE; - } - } - - if (job->nhosts > 0) { - if (setenvf("SLURM_NNODES=%u", job->nhosts)) { - error("Unable to set SLURM_NNODES environment var"); - rc = SLURM_FAILURE; - } - } - - if (job->nodelist) { - if (setenvf("SLURM_NODELIST=%s", job->nodelist)) { - error("Unable to set SLURM_NODELIST environment var."); - rc = SLURM_FAILURE; - } - } - if (opt.nprocs_set && setenvf("SLURM_NPROCS=%u", opt.nprocs)) { error("Unable to set SLURM_NPROCS environment variable"); rc = SLURM_FAILURE; } - if ( opt.cpus_set && setenvf("SLURM_CPUS_PER_TASK=%u", opt.cpus_per_task) ) { error("Unable to set SLURM_CPUS_PER_TASK"); rc = SLURM_FAILURE; } - - if (opt.distribution != SRUN_DIST_UNKNOWN) { + if (job && opt.distribution != SRUN_DIST_UNKNOWN) { dist = (opt.distribution == SRUN_DIST_BLOCK) ? "block" : "cyclic"; @@ -629,11 +609,41 @@ _set_batch_script_env(job_t *job) rc = SLURM_FAILURE; } - if (setenvf("SLURM_TASKS_PER_NODE=%s", (task_cnt = _sprint_task_cnt(job)))) { - error("Unable to set SLURM_TASKS_PER_NODE environment variable"); - rc = SLURM_FAILURE; + /* + * If no job has been allocated yet, just return. We are + * submitting a batch job. + */ + if (job == NULL) + return (rc); + + + if (job->jobid > 0) { + if (setenvf("SLURM_JOBID=%u", job->jobid)) { + error("Unable to set SLURM_JOBID environment"); + rc = SLURM_FAILURE; + } + } + + if (job->nhosts > 0) { + if (setenvf("SLURM_NNODES=%u", job->nhosts)) { + error("Unable to set SLURM_NNODES environment var"); + rc = SLURM_FAILURE; + } + } + + if (job->nodelist) { + if (setenvf("SLURM_NODELIST=%s", job->nodelist)) { + error("Unable to set SLURM_NODELIST environment var."); + rc = SLURM_FAILURE; + } + } + if ((p = _task_count_string (job))) { + if (setenvf ("SLURM_TASKS_PER_NODE=%s", p)) { + error ("Can't set SLURM_TASKS_PER_NODE env variable"); + rc = SLURM_FAILURE; + } + xfree (p); } - xfree(task_cnt); uname(&name); if (strcasecmp(name.sysname, "AIX") == 0) {