Google Git
Sign in
third-party-mirror / SchedMD / slurm / refs/heads/slurm-20.02 / . / src / common / env.c
blob: 394c0ffc15d60723f6e2285de49e825bfc7b59b8 [file] [log] [blame] [edit]
/*****************************************************************************\
* src/common/env.c - add an environment variable to environment vector
*****************************************************************************
* Copyright (C) 2002-2007 The Regents of the University of California.
* Copyright (C) 2008-2009 Lawrence Livermore National Security.
* Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
* Written by Mark Grondona <mgrondona@llnl.gov>, Danny Auble <da@llnl.gov>.
* CODE-OCEC-09-009. All rights reserved.
*
* This file is part of Slurm, a resource management program.
* For details, see <https://slurm.schedmd.com/>.
* Please also read the included file: DISCLAIMER.
*
* Slurm is free software; you can redistribute it and/or modify it under
* the terms of the GNU General Public License as published by the Free
* Software Foundation; either version 2 of the License, or (at your option)
* any later version.
*
* In addition, as a special exception, the copyright holders give permission
* to link the code of portions of this program with the OpenSSL library under
* certain conditions as described in each individual source file, and
* distribute linked combinations including the two. You must obey the GNU
* General Public License in all respects for all of the code used other than
* OpenSSL. If you modify file(s) with this exception, you may extend this
* exception to your version of the file(s), but you are not obligated to do
* so. If you do not wish to do so, delete this exception statement from your
* version. If you delete this exception statement from all source files in
* the program, then also delete it here.
*
* Slurm is distributed in the hope that it will be useful, but WITHOUT ANY
* WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
* FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
* details.
*
* You should have received a copy of the GNU General Public License along
* with Slurm; if not, write to the Free Software Foundation, Inc.,
* 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
\*****************************************************************************/
#include "config.h"
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/param.h> /* MAXPATHLEN */
#include <unistd.h>
#include "slurm/slurm.h"
#include "src/common/cpu_frequency.h"
#include "src/common/log.h"
#include "src/common/env.h"
#include "src/common/fd.h"
#include "src/common/node_select.h"
#include "src/common/macros.h"
#include "src/common/proc_args.h"
#include "src/common/read_config.h"
#include "src/common/slurm_opt.h"
#include "src/common/slurm_protocol_api.h"
#include "src/common/slurm_protocol_defs.h"
#include "src/common/slurm_step_layout.h"
#include "src/common/slurmdb_defs.h"
#include "src/common/strlcpy.h"
#include "src/common/xassert.h"
#include "src/common/xmalloc.h"
#include "src/common/xstring.h"
/*
* Define slurm-specific aliases for use by plugins, see slurm_xlator.h
* for details.
*/
strong_alias(setenvf, slurm_setenvpf);
strong_alias(unsetenvp, slurm_unsetenvp);
strong_alias(getenvp, slurm_getenvp);
strong_alias(env_array_create, slurm_env_array_create);
strong_alias(env_array_merge, slurm_env_array_merge);
strong_alias(env_array_copy, slurm_env_array_copy);
strong_alias(env_array_free, slurm_env_array_free);
strong_alias(env_array_append, slurm_env_array_append);
strong_alias(env_array_append_fmt, slurm_env_array_append_fmt);
strong_alias(env_array_overwrite, slurm_env_array_overwrite);
strong_alias(env_array_overwrite_fmt, slurm_env_array_overwrite_fmt);
strong_alias(env_array_overwrite_het_fmt, slurm_env_array_overwrite_het_fmt);
strong_alias(env_unset_environment, slurm_env_unset_environment);
#define ENV_BUFSIZE (256 * 1024)
#define MAX_ENV_STRLEN (32 * 4096) /* Needed for CPU_BIND and MEM_BIND on
* SGI systems with huge CPU counts */
/*
* Return pointer to `name' entry in environment if found, or
* pointer to the last entry (i.e. NULL) if `name' is not
* currently set in `env'
*
*/
static char **
_find_name_in_env(char **env, const char *name)
{
char **ep;
ep = env;
while (*ep != NULL) {
size_t cnt = 0;
while ( ((*ep)[cnt] == name[cnt])
&& ( name[cnt] != '\0')
&& ((*ep)[cnt] != '\0') )
++cnt;
if (name[cnt] == '\0' && (*ep)[cnt] == '=') {
break;
} else
++ep;
}
return (ep);
}
/*
* Extend memory allocation for env by 1 entry. Make last entry == NULL.
* return pointer to last env entry;
*/
static char **
_extend_env(char ***envp)
{
char **ep;
size_t newcnt = (xsize (*envp) / sizeof (char *)) + 1;
*envp = xrealloc (*envp, newcnt * sizeof (char *));
(*envp)[newcnt - 1] = NULL;
ep = &((*envp)[newcnt - 2]);
/*
* Find last non-NULL entry
*/
while (*ep == NULL)
--ep;
return (++ep);
}
/* return true if the environment variables should not be set for
* srun's --get-user-env option */
static bool _discard_env(char *name, char *value)
{
if ((xstrcmp(name, "DISPLAY") == 0) ||
(xstrcmp(name, "ENVIRONMENT") == 0) ||
(xstrcmp(name, "HOSTNAME") == 0))
return true;
return false;
}
/*
* Return the number of elements in the environment `env'
*/
int
envcount (char **env)
{
int envc = 0;
while (env && env[envc])
envc++;
return (envc);
}
/*
* _setenvfs() (stolen from pdsh)
*
* Set a variable in the callers environment. Args are printf style.
* XXX Space is allocated on the heap and will never be reclaimed.
* Example: setenvfs("RMS_RANK=%d", rank);
*/
int
setenvfs(const char *fmt, ...)
{
va_list ap;
char *buf, *bufcpy, *loc;
int rc, size;
buf = xmalloc(ENV_BUFSIZE);
va_start(ap, fmt);
vsnprintf(buf, ENV_BUFSIZE, fmt, ap);
va_end(ap);
size = strlen(buf);
bufcpy = xstrdup(buf);
xfree(buf);
if (size >= MAX_ENV_STRLEN) {
if ((loc = strchr(bufcpy, '=')))
loc[0] = '\0';
error("environment variable %s is too long", bufcpy);
xfree(bufcpy);
rc = ENOMEM;
} else {
rc = putenv(bufcpy);
}
return rc;
}
int setenvf(char ***envp, const char *name, const char *fmt, ...)
{
char *value;
va_list ap;
int size, rc;
if (!name || name[0] == '\0')
return EINVAL;
value = xmalloc(ENV_BUFSIZE);
va_start(ap, fmt);
vsnprintf(value, ENV_BUFSIZE, fmt, ap);
va_end(ap);
size = strlen(name) + strlen(value) + 2;
if (size >= MAX_ENV_STRLEN) {
error("environment variable %s is too long", name);
return ENOMEM;
}
if (envp && *envp) {
if (env_array_overwrite(envp, name, value) == 1)
rc = 0;
else
rc = 1;
} else {
rc = setenv(name, value, 1);
}
xfree(value);
return rc;
}
/*
* Remove environment variable `name' from "environment"
* contained in `env'
*
* [ This was taken almost verbatim from glibc's
* unsetenv() code. ]
*/
void unsetenvp(char **env, const char *name)
{
char **ep;
if (env == NULL)
return;
ep = env;
while ((ep = _find_name_in_env (ep, name)) && (*ep != NULL)) {
char **dp = ep;
xfree (*ep);
do
dp[0] = dp[1];
while (*dp++);
/* Continue loop in case `name' appears again. */
++ep;
}
return;
}
char *getenvp(char **env, const char *name)
{
size_t len;
char **ep;
if (!name || !env || !env[0])
return (NULL);
len = strlen(name);
ep = _find_name_in_env (env, name);
if (*ep != NULL)
return (&(*ep)[len+1]);
return NULL;
}
int setup_env(env_t *env, bool preserve_env)
{
int rc = SLURM_SUCCESS;
char *addr, *dist = NULL, *lllp_dist = NULL;
char addrbuf[INET_ADDRSTRLEN];
if (env == NULL)
return SLURM_ERROR;
/*
* Always force SLURM_CONF into the environment. This ensures the
* "configless" operation is working, and prevents the client commands
* from falling back to separate RPC requests in case the cache dir
* is unresponsive.
*/
if (setenvf(&env->env, "SLURM_CONF", "%s", getenv("SLURM_CONF"))) {
error("Unable to set SLURM_CONF environment variable");
rc = SLURM_ERROR;
}
/*
* Similarly, prevent this option from leaking in. SLURM_CONF would
* always take precedence, but tidy it up anyways.
*/
unsetenvp(env->env, "SLURM_CONF_SERVER");
if (!preserve_env && env->ntasks) {
if (setenvf(&env->env, "SLURM_NTASKS", "%d", env->ntasks)) {
error("Unable to set SLURM_NTASKS environment variable");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_NPROCS", "%d", env->ntasks)) {
error("Unable to set SLURM_NPROCS environment variable");
rc = SLURM_ERROR;
}
}
if (env->cpus_per_task &&
setenvf(&env->env, "SLURM_CPUS_PER_TASK", "%d",
env->cpus_per_task) ) {
error("Unable to set SLURM_CPUS_PER_TASK");
rc = SLURM_ERROR;
}
if (env->ntasks_per_node
&& setenvf(&env->env, "SLURM_NTASKS_PER_NODE", "%d",
env->ntasks_per_node) ) {
error("Unable to set SLURM_NTASKS_PER_NODE");
rc = SLURM_ERROR;
}
if (env->ntasks_per_socket
&& setenvf(&env->env, "SLURM_NTASKS_PER_SOCKET", "%d",
env->ntasks_per_socket) ) {
error("Unable to set SLURM_NTASKS_PER_SOCKET");
rc = SLURM_ERROR;
}
if (env->ntasks_per_core
&& setenvf(&env->env, "SLURM_NTASKS_PER_CORE", "%d",
env->ntasks_per_core) ) {
error("Unable to set SLURM_NTASKS_PER_CORE");
rc = SLURM_ERROR;
}
if (env->cpus_on_node
&& setenvf(&env->env, "SLURM_CPUS_ON_NODE", "%d",
env->cpus_on_node) ) {
error("Unable to set SLURM_CPUS_ON_NODE");
rc = SLURM_ERROR;
}
set_distribution(env->distribution, &dist, &lllp_dist);
if (dist)
if (setenvf(&env->env, "SLURM_DISTRIBUTION", "%s", dist)) {
error("Can't set SLURM_DISTRIBUTION env variable");
rc = SLURM_ERROR;
}
if ((env->distribution & SLURM_DIST_STATE_BASE) == SLURM_DIST_PLANE)
if (setenvf(&env->env, "SLURM_DIST_PLANESIZE", "%u",
env->plane_size)) {
error("Can't set SLURM_DIST_PLANESIZE env variable");
rc = SLURM_ERROR;
}
if (lllp_dist)
if (setenvf(&env->env, "SLURM_DIST_LLLP", "%s", lllp_dist)) {
error("Can't set SLURM_DIST_LLLP env variable");
rc = SLURM_ERROR;
}
if (env->cpu_bind_type) {
char *str_verbose, *str_bind1 = NULL, *str_bind2 = NULL;
char *str_bind_list, *str_bind_type = NULL, *str_bind = NULL;
if (!env->batch_flag) {
unsetenvp(env->env, "SLURM_CPU_BIND");
unsetenvp(env->env, "SLURM_CPU_BIND_LIST");
unsetenvp(env->env, "SLURM_CPU_BIND_TYPE");
unsetenvp(env->env, "SLURM_CPU_BIND_VERBOSE");
}
if (env->cpu_bind_type & CPU_BIND_VERBOSE)
str_verbose = "verbose";
else
str_verbose = "quiet";
if (env->cpu_bind_type & CPU_BIND_TO_THREADS) {
str_bind1 = "threads";
} else if (env->cpu_bind_type & CPU_BIND_TO_CORES) {
str_bind1 = "cores";
} else if (env->cpu_bind_type & CPU_BIND_TO_SOCKETS) {
str_bind1 = "sockets";
} else if (env->cpu_bind_type & CPU_BIND_TO_LDOMS) {
str_bind1 = "ldoms";
} else if (env->cpu_bind_type & CPU_BIND_TO_BOARDS) {
str_bind1 = "boards";
}
if (env->cpu_bind_type & CPU_BIND_NONE) {
str_bind2 = "none";
} else if (env->cpu_bind_type & CPU_BIND_RANK) {
str_bind2 = "rank";
} else if (env->cpu_bind_type & CPU_BIND_MAP) {
str_bind2 = "map_cpu:";
} else if (env->cpu_bind_type & CPU_BIND_MASK) {
str_bind2 = "mask_cpu:";
} else if (env->cpu_bind_type & CPU_BIND_LDRANK) {
str_bind2 = "rank_ldom";
} else if (env->cpu_bind_type & CPU_BIND_LDMAP) {
str_bind2 = "map_ldom:";
} else if (env->cpu_bind_type & CPU_BIND_LDMASK) {
str_bind2 = "mask_ldom:";
}
if (env->cpu_bind)
str_bind_list = env->cpu_bind;
else
str_bind_list = "";
/* combine first and second part with a comma if needed */
if (str_bind1)
xstrcat(str_bind_type, str_bind1);
if (str_bind1 && str_bind2)
xstrcatchar(str_bind_type, ',');
if (str_bind2)
xstrcat(str_bind_type, str_bind2);
xstrcat(str_bind, str_verbose);
if (str_bind_type) {
xstrcatchar(str_bind, ',');
xstrcat(str_bind, str_bind_type);
xstrcat(str_bind, str_bind_list);
} else
str_bind_type = xstrdup("");
if (!env->batch_flag) {
if (setenvf(&env->env, "SLURM_CPU_BIND", "%s", str_bind)) {
error("Unable to set SLURM_CPU_BIND");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_CPU_BIND_LIST", "%s",
str_bind_list)) {
error("Unable to set SLURM_CPU_BIND_LIST");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_CPU_BIND_TYPE", "%s",
str_bind_type)) {
error("Unable to set SLURM_CPU_BIND_TYPE");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_CPU_BIND_VERBOSE", "%s",
str_verbose)) {
error("Unable to set SLURM_CPU_BIND_VERBOSE");
rc = SLURM_ERROR;
}
}
xfree(str_bind);
xfree(str_bind_type);
}
if (env->mem_bind_type) {
char *str_verbose, *str_bind_type = NULL, *str_bind_list;
char *str_prefer = NULL, *str_bind = NULL;
char *str_bind_sort = NULL;
if (env->batch_flag) {
unsetenvp(env->env, "SBATCH_MEM_BIND");
unsetenvp(env->env, "SBATCH_MEM_BIND_LIST");
unsetenvp(env->env, "SBATCH_MEM_BIND_PREFER");
unsetenvp(env->env, "SBATCH_MEM_BIND_TYPE");
unsetenvp(env->env, "SBATCH_MEM_BIND_VERBOSE");
} else {
unsetenvp(env->env, "SLURM_MEM_BIND");
unsetenvp(env->env, "SLURM_MEM_BIND_LIST");
unsetenvp(env->env, "SLURM_MEM_BIND_PREFER");
unsetenvp(env->env, "SLURM_MEM_BIND_SORT");
unsetenvp(env->env, "SLURM_MEM_BIND_TYPE");
unsetenvp(env->env, "SLURM_MEM_BIND_VERBOSE");
}
if (env->mem_bind_type & MEM_BIND_VERBOSE)
str_verbose = "verbose";
else
str_verbose = "quiet";
if (env->mem_bind_type & MEM_BIND_PREFER)
str_prefer = "prefer";
if (env->mem_bind_type & MEM_BIND_NONE) {
str_bind_type = "none";
} else if (env->mem_bind_type & MEM_BIND_RANK) {
str_bind_type = "rank";
} else if (env->mem_bind_type & MEM_BIND_MAP) {
str_bind_type = "map_mem:";
} else if (env->mem_bind_type & MEM_BIND_MASK) {
str_bind_type = "mask_mem:";
} else if (env->mem_bind_type & MEM_BIND_LOCAL) {
str_bind_type = "local";
}
if (env->mem_bind_type & MEM_BIND_SORT)
str_bind_sort = "sort";
if (env->mem_bind)
str_bind_list = env->mem_bind;
else
str_bind_list = "";
xstrcat(str_bind, str_verbose);
if (str_prefer) {
xstrcatchar(str_bind, ',');
xstrcat(str_bind, str_prefer);
}
if (str_bind_type) {
xstrcatchar(str_bind, ',');
xstrcat(str_bind, str_bind_type);
xstrcat(str_bind, str_bind_list);
} else
str_bind_type = "";
if (env->batch_flag) {
if (setenvf(&env->env, "SBATCH_MEM_BIND", "%s", str_bind)) {
error("Unable to set SBATCH_MEM_BIND");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SBATCH_MEM_BIND_LIST", "%s",
str_bind_list)) {
error("Unable to set SBATCH_MEM_BIND_LIST");
rc = SLURM_ERROR;
}
if (str_prefer &&
setenvf(&env->env, "SBATCH_MEM_BIND_PREFER", "%s",
str_prefer)) {
error("Unable to set SBATCH_MEM_BIND_PREFER");
rc = SLURM_ERROR;
}
if (str_bind_sort &&
setenvf(&env->env, "SBATCH_MEM_BIND_SORT", "%s",
str_bind_sort)) {
error("Unable to set SBATCH_MEM_BIND_SORT");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SBATCH_MEM_BIND_TYPE", "%s",
str_bind_type)) {
error("Unable to set SBATCH_MEM_BIND_TYPE");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SBATCH_MEM_BIND_VERBOSE", "%s",
str_verbose)) {
error("Unable to set SBATCH_MEM_BIND_VERBOSE");
rc = SLURM_ERROR;
}
} else {
if (setenvf(&env->env, "SLURM_MEM_BIND", "%s", str_bind)) {
error("Unable to set SLURM_MEM_BIND");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_MEM_BIND_LIST", "%s",
str_bind_list)) {
error("Unable to set SLURM_MEM_BIND_LIST");
rc = SLURM_ERROR;
}
if (str_prefer &&
setenvf(&env->env, "SLURM_MEM_BIND_PREFER", "%s",
str_prefer)) {
error("Unable to set SLURM_MEM_BIND_PREFER");
rc = SLURM_ERROR;
}
if (str_bind_sort &&
setenvf(&env->env, "SLURM_MEM_BIND_SORT", "%s",
str_bind_sort)) {
error("Unable to set SLURM_MEM_BIND_SORT");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_MEM_BIND_TYPE", "%s",
str_bind_type)) {
error("Unable to set SLURM_MEM_BIND_TYPE");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_MEM_BIND_VERBOSE", "%s",
str_verbose)) {
error("Unable to set SLURM_MEM_BIND_VERBOSE");
rc = SLURM_ERROR;
}
}
xfree(str_bind);
}
if (cpu_freq_set_env("SLURM_CPU_FREQ_REQ", env->cpu_freq_min,
env->cpu_freq_max, env->cpu_freq_gov) != SLURM_SUCCESS)
rc = SLURM_ERROR;
if (env->overcommit
&& (setenvf(&env->env, "SLURM_OVERCOMMIT", "%s", "1"))) {
error("Unable to set SLURM_OVERCOMMIT environment variable");
rc = SLURM_ERROR;
}
if (env->slurmd_debug
&& setenvf(&env->env, "SLURMD_DEBUG", "%d", env->slurmd_debug)) {
error("Can't set SLURMD_DEBUG environment variable");
rc = SLURM_ERROR;
}
if (env->labelio
&& setenvf(&env->env, "SLURM_LABELIO", "1")) {
error("Unable to set SLURM_LABELIO environment variable");
rc = SLURM_ERROR;
}
if (env->jobid >= 0) {
if (setenvf(&env->env, "SLURM_JOB_ID", "%d", env->jobid)) {
error("Unable to set SLURM_JOB_ID environment");
rc = SLURM_ERROR;
}
/* and for backwards compatibility... */
if (setenvf(&env->env, "SLURM_JOBID", "%d", env->jobid)) {
error("Unable to set SLURM_JOBID environment");
rc = SLURM_ERROR;
}
}
if (env->job_name) {
if (setenvf(&env->env, "SLURM_JOB_NAME", "%s", env->job_name)) {
error("Unable to set SLURM_JOB_NAME environment");
rc = SLURM_ERROR;
}
}
/*
* These aren't relevant to a system not using Slurm as the
* launcher. Since there isn't a flag for that we check for
* the flags we do have.
*/
if (env->task_pid &&
setenvf(&env->env, "SLURM_TASK_PID", "%d",
(int)env->task_pid)) {
error("Unable to set SLURM_TASK_PID environment "
"variable");
rc = SLURM_ERROR;
}
if ((env->nodeid >= 0) &&
setenvf(&env->env, "SLURM_NODEID", "%d", env->nodeid)) {
error("Unable to set SLURM_NODEID environment");
rc = SLURM_ERROR;
}
if ((env->procid >= 0) &&
setenvf(&env->env, "SLURM_PROCID", "%d", env->procid)) {
error("Unable to set SLURM_PROCID environment");
rc = SLURM_ERROR;
}
if ((env->localid >= 0) &&
setenvf(&env->env, "SLURM_LOCALID", "%d", env->localid)) {
error("Unable to set SLURM_LOCALID environment");
rc = SLURM_ERROR;
}
if (env->stepid >= 0) {
if (setenvf(&env->env, "SLURM_STEP_ID", "%d", env->stepid)) {
error("Unable to set SLURM_STEP_ID environment");
rc = SLURM_ERROR;
}
/* and for backwards compatibility... */
if (setenvf(&env->env, "SLURM_STEPID", "%d", env->stepid)) {
error("Unable to set SLURM_STEPID environment");
rc = SLURM_ERROR;
}
}
if (!preserve_env && env->nhosts
&& setenvf(&env->env, "SLURM_NNODES", "%d", env->nhosts)) {
error("Unable to set SLURM_NNODES environment var");
rc = SLURM_ERROR;
}
if (env->nhosts
&& setenvf(&env->env, "SLURM_JOB_NUM_NODES", "%d", env->nhosts)) {
error("Unable to set SLURM_JOB_NUM_NODES environment var");
rc = SLURM_ERROR;
}
if (env->nodelist &&
setenvf(&env->env, "SLURM_NODELIST", "%s", env->nodelist)) {
error("Unable to set SLURM_NODELIST environment var.");
rc = SLURM_ERROR;
}
if (env->partition
&& setenvf(&env->env, "SLURM_JOB_PARTITION", "%s", env->partition)) {
error("Unable to set SLURM_JOB_PARTITION environment var.");
rc = SLURM_ERROR;
}
if (!preserve_env && env->task_count
&& setenvf (&env->env,
"SLURM_TASKS_PER_NODE", "%s", env->task_count)) {
error ("Can't set SLURM_TASKS_PER_NODE env variable");
rc = SLURM_ERROR;
}
if (env->comm_port
&& setenvf (&env->env, "SLURM_SRUN_COMM_PORT", "%u",
env->comm_port)) {
error ("Can't set SLURM_SRUN_COMM_PORT env variable");
rc = SLURM_ERROR;
}
if (env->cli) {
slurm_print_slurm_addr (env->cli, addrbuf, INET_ADDRSTRLEN);
/*
* XXX: Eventually, need a function for slurm_addrs that
* returns just the IP address (not addr:port)
*/
if ((dist = strchr (addrbuf, ':')) != NULL)
*dist = '\0';
setenvf (&env->env, "SLURM_LAUNCH_NODE_IPADDR", "%s", addrbuf);
}
if (env->sgtids &&
setenvf(&env->env, "SLURM_GTIDS", "%s", env->sgtids)) {
error("Unable to set SLURM_GTIDS environment variable");
rc = SLURM_ERROR;
}
if (env->pty_port
&& setenvf(&env->env, "SLURM_PTY_PORT", "%hu", env->pty_port)) {
error("Can't set SLURM_PTY_PORT env variable");
rc = SLURM_ERROR;
}
if (env->ws_col
&& setenvf(&env->env, "SLURM_PTY_WIN_COL", "%hu", env->ws_col)) {
error("Can't set SLURM_PTY_WIN_COL env variable");
rc = SLURM_ERROR;
}
if (env->ws_row
&& setenvf(&env->env, "SLURM_PTY_WIN_ROW", "%hu", env->ws_row)) {
error("Can't set SLURM_PTY_WIN_ROW env variable");
rc = SLURM_ERROR;
}
if (env->restart_cnt &&
setenvf(&env->env, "SLURM_RESTART_COUNT", "%u", env->restart_cnt)) {
error("Can't set SLURM_RESTART_COUNT env variable");
rc = SLURM_ERROR;
}
if (env->user_name) {
if (setenvf(&env->env, "SLURM_JOB_UID", "%u",
(unsigned int) env->uid)) {
error("Can't set SLURM_JOB_UID env variable");
rc = SLURM_ERROR;
}
if (setenvf(&env->env, "SLURM_JOB_USER", "%s", env->user_name)){
error("Can't set SLURM_JOB_USER env variable");
rc = SLURM_ERROR;
}
}
if (env->account) {
if (setenvf(&env->env,
"SLURM_JOB_ACCOUNT",
"%s",
env->account)) {
error("%s: can't set SLURM_JOB_ACCOUNT env variable",
__func__);
rc = SLURM_ERROR;
}
}
if (env->qos) {
if (setenvf(&env->env,
"SLURM_JOB_QOS",
"%s",
env->qos)) {
error("%s: can't set SLURM_JOB_QOS env variable",
__func__);
rc = SLURM_ERROR;
}
}
if (env->resv_name) {
if (setenvf(&env->env,
"SLURM_JOB_RESERVATION",
"%s",
env->resv_name)) {
error("%s: can't set SLURM_JOB_RESERVATION env variable",
__func__);
rc = SLURM_ERROR;
}
}
if (slurmctld_conf.slurmctld_addr)
addr = slurmctld_conf.slurmctld_addr;
else
addr = slurmctld_conf.control_addr[0];
setenvf(&env->env, "SLURM_WORKING_CLUSTER", "%s:%s:%d:%d:%d",
slurmctld_conf.cluster_name, addr,
slurmctld_conf.slurmctld_port, SLURM_PROTOCOL_VERSION,
select_get_plugin_id());
return rc;
}
/**********************************************************************
* From here on are the new environment variable management functions,
* used by the "new" commands: salloc, sbatch, and the step launch APIs.
**********************************************************************/
/*
* Return a string representation of an array of uint16_t elements.
* Each value in the array is printed in decimal notation and elements
* are separated by a comma. If sequential elements in the array
* contain the same value, the value is written out just once followed
* by "(xN)", where "N" is the number of times the value is repeated.
*
* Example:
* The array "1, 2, 1, 1, 1, 3, 2" becomes the string "1,2,1(x3),3,2"
*
* Returns an xmalloc'ed string. Free with xfree().
*/
extern char *uint16_array_to_str(int array_len, const uint16_t *array)
{
int i;
int previous = 0;
char *sep = ","; /* seperator */
char *str = xstrdup("");
if (array == NULL)
return str;
for (i = 0; i < array_len; i++) {
if ((i+1 < array_len) && (array[i] == array[i+1])) {
previous++;
continue;
}
if (i == array_len-1) /* last time through loop */
sep = "";
if (previous > 0) {
xstrfmtcat(str, "%u(x%u)%s",
array[i], previous+1, sep);
} else {
xstrfmtcat(str, "%u%s", array[i], sep);
}
previous = 0;
}
return str;
}
/*
* The cpus-per-node representation in Slurm (and perhaps tasks-per-node
* in the future) is stored in a compressed format comprised of two
* equal-length arrays, and an integer holding the array length. In one
* array an element represents a count (number of cpus, number of tasks,
* etc.), and the corresponding element in the other array contains the
* number of times the count is repeated sequentially in the uncompressed
* something-per-node array.
*
* This function returns the string representation of the compressed
* array. Free with xfree().
*/
extern char *uint32_compressed_to_str(uint32_t array_len,
const uint16_t *array,
const uint32_t *array_reps)
{
int i;
char *sep = ","; /* seperator */
char *str = xstrdup("");
if (!array || !array_reps)
return str;
for (i = 0; i < array_len; i++) {
if (i == array_len-1) /* last time through loop */
sep = "";
if (array_reps[i] > 1) {
xstrfmtcat(str, "%u(x%u)%s",
array[i], array_reps[i], sep);
} else {
xstrfmtcat(str, "%u%s", array[i], sep);
}
}
return str;
}
/*
* Set in "dest" the environment variables relevant to a Slurm job
* allocation, 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.
*
* Sets the variables:
* SLURM_JOB_ID
* SLURM_JOB_NAME
* SLURM_JOB_NUM_NODES
* SLURM_JOB_NODELIST
* SLURM_JOB_CPUS_PER_NODE
* SLURM_NODE_ALIASES
* SLURM_NTASKS_PER_NODE
*
* dest OUT - array in which to the set environment variables
* alloc IN - resource allocation response
* desc IN - job allocation request
* het_job_offset IN - component offset into hetjob, -1 if not hetjob
*
* Sets OBSOLETE variables (needed for MPI, do not remove):
* SLURM_JOBID
* SLURM_NNODES
* SLURM_NODELIST
* SLURM_TASKS_PER_NODE
*/
extern int env_array_for_job(char ***dest,
const resource_allocation_response_msg_t *alloc,
const job_desc_msg_t *desc, int het_job_offset)
{
char *tmp = NULL;
char *dist = NULL, *lllp_dist = NULL;
char *key, *value;
slurm_step_layout_t *step_layout = NULL;
int i, rc = SLURM_SUCCESS;
slurm_step_layout_req_t step_layout_req;
uint16_t cpus_per_task_array[1];
uint32_t cpus_task_reps[1];
if (!alloc || !desc)
return SLURM_ERROR;
memset(&step_layout_req, 0, sizeof(slurm_step_layout_req_t));
step_layout_req.num_tasks = desc->num_tasks;
step_layout_req.num_hosts = alloc->node_cnt;
cpus_per_task_array[0] = desc->cpus_per_task;
cpus_task_reps[0] = alloc->node_cnt;
if (het_job_offset < 1) {
env_array_overwrite_fmt(dest, "SLURM_JOB_ID", "%u",
alloc->job_id);
}
env_array_overwrite_het_fmt(dest, "SLURM_JOB_ID", het_job_offset,
"%u", alloc->job_id);
env_array_overwrite_het_fmt(dest, "SLURM_JOB_NAME", het_job_offset,
"%s", desc->name);
env_array_overwrite_het_fmt(dest, "SLURM_JOB_NUM_NODES", het_job_offset,
"%u", step_layout_req.num_hosts);
env_array_overwrite_het_fmt(dest, "SLURM_JOB_NODELIST", het_job_offset,
"%s", alloc->node_list);
env_array_overwrite_het_fmt(dest, "SLURM_NODE_ALIASES", het_job_offset,
"%s", alloc->alias_list);
env_array_overwrite_het_fmt(dest, "SLURM_JOB_PARTITION", het_job_offset,
"%s", alloc->partition);
set_distribution(desc->task_dist, &dist, &lllp_dist);
if (dist) {
env_array_overwrite_het_fmt(dest, "SLURM_DISTRIBUTION",
het_job_offset, "%s", dist);
}
if ((desc->task_dist & SLURM_DIST_STATE_BASE) == SLURM_DIST_PLANE) {
env_array_overwrite_het_fmt(dest, "SLURM_DIST_PLANESIZE",
het_job_offset, "%u",
desc->plane_size);
}
if (lllp_dist) {
env_array_overwrite_het_fmt(dest, "SLURM_DIST_LLLP",
het_job_offset, "%s", lllp_dist);
}
tmp = uint32_compressed_to_str(alloc->num_cpu_groups,
alloc->cpus_per_node,
alloc->cpu_count_reps);
env_array_overwrite_het_fmt(dest, "SLURM_JOB_CPUS_PER_NODE",
het_job_offset, "%s", tmp);
xfree(tmp);
if (alloc->pn_min_memory & MEM_PER_CPU) {
uint64_t tmp_mem = alloc->pn_min_memory & (~MEM_PER_CPU);
env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_CPU",
het_job_offset, "%"PRIu64"",
tmp_mem);
} else if (alloc->pn_min_memory) {
uint64_t tmp_mem = alloc->pn_min_memory;
env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_NODE",
het_job_offset, "%"PRIu64"",
tmp_mem);
}
/* OBSOLETE, but needed by MPI, do not remove */
env_array_overwrite_het_fmt(dest, "SLURM_JOBID", het_job_offset, "%u",
alloc->job_id);
env_array_overwrite_het_fmt(dest, "SLURM_NNODES", het_job_offset, "%u",
step_layout_req.num_hosts);
env_array_overwrite_het_fmt(dest, "SLURM_NODELIST", het_job_offset, "%s",
alloc->node_list);
if (step_layout_req.num_tasks == NO_VAL) {
/* If we know how many tasks we are going to do then
we set SLURM_TASKS_PER_NODE */
int i = 0;
/* If no tasks were given we can figure it out here
* by totalling up the cpus and then dividing by the
* number of cpus per task */
step_layout_req.num_tasks = 0;
for (i = 0; i < alloc->num_cpu_groups; i++) {
step_layout_req.num_tasks += alloc->cpu_count_reps[i]
* alloc->cpus_per_node[i];
}
if ((int)desc->cpus_per_task > 1
&& desc->cpus_per_task != NO_VAL16)
step_layout_req.num_tasks /= desc->cpus_per_task;
//num_tasks = desc->min_cpus;
}
if ((desc->task_dist & SLURM_DIST_STATE_BASE) == SLURM_DIST_ARBITRARY) {
step_layout_req.node_list = desc->req_nodes;
env_array_overwrite_het_fmt(dest, "SLURM_ARBITRARY_NODELIST",
het_job_offset, "%s",
step_layout_req.node_list);
} else
step_layout_req.node_list = alloc->node_list;
step_layout_req.cpus_per_node = alloc->cpus_per_node;
step_layout_req.cpu_count_reps = alloc->cpu_count_reps;
step_layout_req.cpus_per_task = cpus_per_task_array;
step_layout_req.cpus_task_reps = cpus_task_reps;
step_layout_req.task_dist = desc->task_dist;
step_layout_req.plane_size = desc->plane_size;
if (!(step_layout = slurm_step_layout_create(&step_layout_req)))
return SLURM_ERROR;
tmp = uint16_array_to_str(step_layout->node_cnt, step_layout->tasks);
slurm_step_layout_destroy(step_layout);
env_array_overwrite_het_fmt(dest, "SLURM_TASKS_PER_NODE",
het_job_offset,
"%s", tmp);
xfree(tmp);
if (alloc->account) {
env_array_overwrite_het_fmt(dest, "SLURM_JOB_ACCOUNT",
het_job_offset, "%s",
alloc->account);
}
if (alloc->qos) {
env_array_overwrite_het_fmt(dest, "SLURM_JOB_QOS",
het_job_offset,
"%s", alloc->qos);
}
if (alloc->resv_name) {
env_array_overwrite_het_fmt(dest, "SLURM_JOB_RESERVATION",
het_job_offset, "%s",
alloc->resv_name);
}
if (alloc->env_size) { /* Used to set Burst Buffer environment */
for (i = 0; i < alloc->env_size; i++) {
tmp = xstrdup(alloc->environment[i]);
key = tmp;
value = strchr(tmp, '=');
if (value) {
value[0] = '\0';
value++;
env_array_overwrite_het_fmt(dest, key,
het_job_offset,
"%s",
value);
}
xfree(tmp);
}
}
if (desc->acctg_freq) {
env_array_overwrite_het_fmt(dest, "SLURM_ACCTG_FREQ",
het_job_offset, "%s",
desc->acctg_freq);
};
if (desc->network) {
env_array_overwrite_het_fmt(dest, "SLURM_NETWORK",
het_job_offset, "%s",
desc->network);
}
if (desc->overcommit != NO_VAL8) {
env_array_overwrite_het_fmt(dest, "SLURM_OVERCOMMIT",
het_job_offset, "%u",
desc->overcommit);
}
/* Add default task counts for srun, if not already set */
if (desc->bitflags & JOB_NTASKS_SET) {
env_array_overwrite_het_fmt(dest, "SLURM_NTASKS",
het_job_offset,
"%d", desc->num_tasks);
/* maintain for old scripts */
env_array_overwrite_het_fmt(dest, "SLURM_NPROCS",
het_job_offset,
"%d", desc->num_tasks);
}
if (desc->bitflags & JOB_CPUS_SET) {
env_array_overwrite_het_fmt(dest, "SLURM_CPUS_PER_TASK",
het_job_offset, "%d",
desc->cpus_per_task);
}
if (desc->ntasks_per_node && (desc->ntasks_per_node != NO_VAL16)) {
env_array_overwrite_het_fmt(dest, "SLURM_NTASKS_PER_NODE",
het_job_offset, "%d",
desc->ntasks_per_node);
}
return rc;
}
/*
* Set in "dest" the environment variables strings relevant to a Slurm batch
* job allocation, 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.
*
* Sets the variables:
* SLURM_CLUSTER_NAME
* SLURM_JOB_ID
* SLURM_JOB_NUM_NODES
* SLURM_JOB_NODELIST
* SLURM_JOB_CPUS_PER_NODE
* SLURM_NODE_ALIASES
* ENVIRONMENT=BATCH
* HOSTNAME
*
* Sets OBSOLETE variables (needed for MPI, do not remove):
* SLURM_JOBID
* SLURM_NNODES
* SLURM_NODELIST
* SLURM_NTASKS
* SLURM_TASKS_PER_NODE
*/
extern int
env_array_for_batch_job(char ***dest, const batch_job_launch_msg_t *batch,
const char *node_name)
{
char *tmp = NULL, *cluster_name;
uint32_t num_cpus = 0;
int i;
slurm_step_layout_t *step_layout = NULL;
uint16_t cpus_per_task;
uint32_t task_dist;
slurm_step_layout_req_t step_layout_req;
uint16_t cpus_per_task_array[1];
uint32_t cpus_task_reps[1];
if (!batch)
return SLURM_ERROR;
memset(&step_layout_req, 0, sizeof(slurm_step_layout_req_t));
step_layout_req.num_tasks = batch->ntasks;
/*
* There is no explicit node count in the batch structure,
* so we need to calculate the node count.
*/
for (i = 0; i < batch->num_cpu_groups; i++) {
step_layout_req.num_hosts += batch->cpu_count_reps[i];
num_cpus += batch->cpu_count_reps[i] * batch->cpus_per_node[i];
}
cluster_name = slurm_get_cluster_name();
if (cluster_name) {
env_array_overwrite_fmt(dest, "SLURM_CLUSTER_NAME", "%s",
cluster_name);
xfree(cluster_name);
}
env_array_overwrite_fmt(dest, "SLURM_JOB_ID", "%u", batch->job_id);
env_array_overwrite_fmt(dest, "SLURM_JOB_NUM_NODES", "%u",
step_layout_req.num_hosts);
if (batch->array_task_id != NO_VAL) {
env_array_overwrite_fmt(dest, "SLURM_ARRAY_JOB_ID", "%u",
batch->array_job_id);
env_array_overwrite_fmt(dest, "SLURM_ARRAY_TASK_ID", "%u",
batch->array_task_id);
}
env_array_overwrite_fmt(dest, "SLURM_JOB_NODELIST", "%s", batch->nodes);
env_array_overwrite_fmt(dest, "SLURM_JOB_PARTITION", "%s",
batch->partition);
env_array_overwrite_fmt(dest, "SLURM_NODE_ALIASES", "%s",
batch->alias_list);
tmp = uint32_compressed_to_str(batch->num_cpu_groups,
batch->cpus_per_node,
batch->cpu_count_reps);
env_array_overwrite_fmt(dest, "SLURM_JOB_CPUS_PER_NODE", "%s", tmp);
xfree(tmp);
env_array_overwrite_fmt(dest, "ENVIRONMENT", "BATCH");
if (node_name)
env_array_overwrite_fmt(dest, "HOSTNAME", "%s", node_name);
/* OBSOLETE, but needed by MPI, do not remove */
env_array_overwrite_fmt(dest, "SLURM_JOBID", "%u", batch->job_id);
env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u",
step_layout_req.num_hosts);
env_array_overwrite_fmt(dest, "SLURM_NODELIST", "%s", batch->nodes);
if ((batch->cpus_per_task != 0) &&
(batch->cpus_per_task != NO_VAL16))
cpus_per_task = batch->cpus_per_task;
else
cpus_per_task = 1; /* default value */
cpus_per_task_array[0] = cpus_per_task;
cpus_task_reps[0] = step_layout_req.num_hosts;
/* Only overwrite this if it is set. They are set in
* sbatch directly and could have changed. */
if (getenvp(*dest, "SLURM_CPUS_PER_TASK"))
env_array_overwrite_fmt(dest, "SLURM_CPUS_PER_TASK", "%u",
cpus_per_task);
if (step_layout_req.num_tasks) {
env_array_append_fmt(dest, "SLURM_NTASKS", "%u",
step_layout_req.num_tasks);
/* keep around for old scripts */
env_array_append_fmt(dest, "SLURM_NPROCS", "%u",
step_layout_req.num_tasks);
} else {
step_layout_req.num_tasks = num_cpus / cpus_per_task;
}
if ((step_layout_req.node_list =
getenvp(*dest, "SLURM_ARBITRARY_NODELIST"))) {
task_dist = SLURM_DIST_ARBITRARY;
} else {
step_layout_req.node_list = batch->nodes;
task_dist = SLURM_DIST_BLOCK;
}
step_layout_req.cpus_per_node = batch->cpus_per_node;
step_layout_req.cpu_count_reps = batch->cpu_count_reps;
step_layout_req.cpus_per_task = cpus_per_task_array;
step_layout_req.cpus_task_reps = cpus_task_reps;
step_layout_req.task_dist = task_dist;
step_layout_req.plane_size = NO_VAL16;
if (!(step_layout = slurm_step_layout_create(&step_layout_req)))
return SLURM_ERROR;
tmp = uint16_array_to_str(step_layout->node_cnt, step_layout->tasks);
slurm_step_layout_destroy(step_layout);
env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE", "%s", tmp);
xfree(tmp);
if (batch->pn_min_memory & MEM_PER_CPU) {
uint64_t tmp_mem = batch->pn_min_memory & (~MEM_PER_CPU);
env_array_overwrite_fmt(dest, "SLURM_MEM_PER_CPU", "%"PRIu64"",
tmp_mem);
} else if (batch->pn_min_memory) {
uint64_t tmp_mem = batch->pn_min_memory;
env_array_overwrite_fmt(dest, "SLURM_MEM_PER_NODE", "%"PRIu64"",
tmp_mem);
}
/* Set the SLURM_JOB_ACCOUNT, SLURM_JOB_QOS
* and SLURM_JOB_RESERVATION if set by
* the controller.
*/
if (batch->account) {
env_array_overwrite_fmt(dest,
"SLURM_JOB_ACCOUNT",
"%s",
batch->account);
}
if (batch->qos) {
env_array_overwrite_fmt(dest,
"SLURM_JOB_QOS",
"%s",
batch->qos);
}
if (batch->resv_name) {
env_array_overwrite_fmt(dest,
"SLURM_JOB_RESERVATION",
"%s",
batch->resv_name);
}
return SLURM_SUCCESS;
}
/*
* 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. If preserve_env is
* true, the variables SLURM_NNODES, SLURM_NTASKS and SLURM_TASKS_PER_NODE
* remain unchanged.
*
* Sets variables:
* SLURM_STEP_ID
* SLURM_STEP_NUM_NODES
* SLURM_STEP_NUM_TASKS
* SLURM_STEP_TASKS_PER_NODE
* SLURM_STEP_LAUNCHER_PORT
* SLURM_STEP_LAUNCHER_IPADDR
* SLURM_STEP_RESV_PORTS
* SLURM_STEP_SUB_MP
*
* Sets OBSOLETE variables:
* SLURM_STEPID
* SLURM_NNODES
* SLURM_NTASKS
* SLURM_NODELIST
* SLURM_TASKS_PER_NODE
* SLURM_SRUN_COMM_PORT
* SLURM_LAUNCH_NODE_IPADDR
*
*/
extern void
env_array_for_step(char ***dest,
const job_step_create_response_msg_t *step,
launch_tasks_request_msg_t *launch,
uint16_t launcher_port,
bool preserve_env)
{
char *tmp, *tpn;
uint32_t node_cnt, task_cnt;
if (!step || !launch)
return;
node_cnt = step->step_layout->node_cnt;
env_array_overwrite_fmt(dest, "SLURM_STEP_ID", "%u", step->job_step_id);
if (launch->het_job_node_list) {
tmp = launch->het_job_node_list;
env_array_overwrite_fmt(dest, "SLURM_NODELIST", "%s", tmp);
env_array_overwrite_fmt(dest, "SLURM_JOB_NODELIST", "%s", tmp);
} else {
tmp = step->step_layout->node_list;
env_array_append_fmt(dest, "SLURM_JOB_NODELIST", "%s", tmp);
}
env_array_overwrite_fmt(dest, "SLURM_STEP_NODELIST", "%s", tmp);
if (launch->het_job_nnodes && (launch->het_job_nnodes != NO_VAL))
node_cnt = launch->het_job_nnodes;
env_array_overwrite_fmt(dest, "SLURM_STEP_NUM_NODES", "%u", node_cnt);
if (launch->het_job_ntasks && (launch->het_job_ntasks != NO_VAL))
task_cnt = launch->het_job_ntasks;
else
task_cnt = step->step_layout->task_cnt;
env_array_overwrite_fmt(dest, "SLURM_STEP_NUM_TASKS", "%u", task_cnt);
if (launch->het_job_task_cnts) {
tpn = uint16_array_to_str(launch->het_job_nnodes,
launch->het_job_task_cnts);
env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE", "%s",
tpn);
env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u",
launch->het_job_nnodes);
} else {
tpn = uint16_array_to_str(step->step_layout->node_cnt,
step->step_layout->tasks);
if (!preserve_env) {
env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE",
"%s", tpn);
}
}
env_array_overwrite_fmt(dest, "SLURM_STEP_TASKS_PER_NODE", "%s", tpn);
env_array_overwrite_fmt(dest, "SLURM_STEP_LAUNCHER_PORT",
"%hu", launcher_port);
if (step->resv_ports) {
env_array_overwrite_fmt(dest, "SLURM_STEP_RESV_PORTS",
"%s", step->resv_ports);
}
/* OBSOLETE, but needed by some MPI implementations, do not remove */
env_array_overwrite_fmt(dest, "SLURM_STEPID", "%u", step->job_step_id);
if (!preserve_env) {
env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u", node_cnt);
env_array_overwrite_fmt(dest, "SLURM_NTASKS", "%u", task_cnt);
/* keep around for old scripts */
env_array_overwrite_fmt(dest, "SLURM_NPROCS",
"%u", step->step_layout->task_cnt);
}
env_array_overwrite_fmt(dest, "SLURM_SRUN_COMM_PORT",
"%hu", launcher_port);
xfree(tpn);
}
/*
* Enviroment variables set elsewhere
* ----------------------------------
*
* Set by slurmstepd:
* SLURM_STEP_NODEID
* SLURM_STEP_PROCID
* SLURM_STEP_LOCALID
*
* OBSOLETE set by slurmstepd:
* SLURM_NODEID
* SLURM_PROCID
* SLURM_LOCALID
*/
/***********************************************************************
* Environment variable array support functions
***********************************************************************/
/*
* Return an empty environment variable array (contains a single
* pointer to NULL).
*/
char **env_array_create(void)
{
char **env_array;
env_array = xmalloc(sizeof(char *));
env_array[0] = NULL;
return env_array;
}
static int _env_array_update(char ***array_ptr, const char *name,
const char *value, bool over_write)
{
char **ep = NULL;
char *str = NULL;
if (array_ptr == NULL)
return 0;
if (*array_ptr == NULL)
*array_ptr = env_array_create();
ep = _find_name_in_env(*array_ptr, name);
if (*ep != NULL) {
if (!over_write)
return 0;
xfree (*ep);
} else {
ep = _extend_env(array_ptr);
}
xstrfmtcat(str, "%s=%s", name, value);
*ep = str;
return 1;
}
/*
* Append a single environment variable to an environment variable array,
* if and only if a variable by that name does not already exist in the
* array.
*
* "value_fmt" supports printf-style formatting.
*
* Return 1 on success, and 0 on error.
*/
int env_array_append_fmt(char ***array_ptr, const char *name,
const char *value_fmt, ...)
{
int rc;
char *value;
va_list ap;
value = xmalloc(ENV_BUFSIZE);
va_start(ap, value_fmt);
vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
va_end(ap);
rc = env_array_append(array_ptr, name, value);
xfree(value);
return rc;
}
/*
* Append a single environment variable to an environment variable array,
* if and only if a variable by that name does not already exist in the
* array.
*
* Return 1 on success, and 0 on error.
*/
int env_array_append(char ***array_ptr, const char *name,
const char *value)
{
return _env_array_update(array_ptr, name, value, false);
}
/*
* Append a single environment variable to an environment variable array
* if a variable by that name does not already exist. If a variable
* by the same name is found in the array, it is overwritten with the
* new value.
*
* "value_fmt" supports printf-style formatting.
*
* Return 1 on success, and 0 on error.
*/
int env_array_overwrite_fmt(char ***array_ptr, const char *name,
const char *value_fmt, ...)
{
int rc;
char *value;
va_list ap;
value = xmalloc(ENV_BUFSIZE);
va_start(ap, value_fmt);
vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
va_end(ap);
rc = env_array_overwrite(array_ptr, name, value);
xfree(value);
return rc;
}
/*
* Append a single environment variable to an environment variable array
* if a variable by that name does not already exist. If a variable
* by the same name is found in the array, it is overwritten with the
* new value.
*
* "value_fmt" supports printf-style formatting.
*
* Return 1 on success, and 0 on error.
*/
int env_array_overwrite_het_fmt(char ***array_ptr, const char *name,
int het_job_offset,
const char *value_fmt, ...)
{
int rc;
char *value;
va_list ap;
value = xmalloc(ENV_BUFSIZE);
va_start(ap, value_fmt);
vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
va_end(ap);
if (het_job_offset != -1) {
char *het_comp_name = NULL;
/* Continue support for old hetjob terminology. */
xstrfmtcat(het_comp_name, "%s_PACK_GROUP_%d", name,
het_job_offset);
rc = env_array_overwrite(array_ptr, het_comp_name, value);
xfree(het_comp_name);
xstrfmtcat(het_comp_name, "%s_HET_GROUP_%d", name,
het_job_offset);
rc = env_array_overwrite(array_ptr, het_comp_name, value);
xfree(het_comp_name);
} else
rc = env_array_overwrite(array_ptr, name, value);
xfree(value);
return rc;
}
/*
* Append a single environment variable to an environment variable array
* if a variable by that name does not already exist. If a variable
* by the same name is found in the array, it is overwritten with the
* new value.
*
* Return 1 on success, and 0 on error.
*/
int env_array_overwrite(char ***array_ptr, const char *name,
const char *value)
{
return _env_array_update(array_ptr, name, value, true);
}
/*
* Copy env_array must be freed by env_array_free
*/
char **env_array_copy(const char **array)
{
char **ptr = NULL;
env_array_merge(&ptr, array);
return ptr;
}
/*
* Free the memory used by an environment variable array.
*/
void env_array_free(char **env_array)
{
char **ptr;
if (env_array == NULL)
return;
for (ptr = env_array; *ptr != NULL; ptr++) {
xfree(*ptr);
}
xfree(env_array);
}
/*
* Given an environment variable "name=value" string,
* copy the name portion into the "name" buffer, and the
* value portion into the "value" buffer.
*
* Return 1 on success, 0 on failure.
*/
static int _env_array_entry_splitter(const char *entry,
char *name, int name_len,
char *value, int value_len)
{
char *ptr;
int len;
ptr = xstrchr(entry, '=');
if (ptr == NULL) /* Bad parsing, no '=' found */
return 0;
/*
* need to consider the byte pointed by ptr.
* example: entry = 0x0 = "a=b"
* ptr = 0x1
* len = ptr - entry + 1 = 2 because we need
* 2 characters to store 'a\0'
*/
len = ptr - entry + 1;
if (len > name_len)
return 0;
strlcpy(name, entry, len);
ptr++;
/* account for '\0' here */
len = strlen(ptr) + 1;
if (len > value_len)
return 0;
strlcpy(value, ptr, len);
return 1;
}
/*
* Work similarly to putenv() (from C stdlib), but uses setenv()
* under the covers. This avoids having pointers from the global
* array "environ" into "string".
*
* Return 1 on success, 0 on failure.
*/
static int _env_array_putenv(const char *string)
{
int rc = 0;
char name[256], *value;
value = xmalloc(ENV_BUFSIZE);
if ((_env_array_entry_splitter(string, name, sizeof(name),
value, ENV_BUFSIZE)) &&
(setenv(name, value, 1) != -1))
rc = 1;
xfree(value);
return rc;
}
/*
* Set all of the environment variables in a supplied environment
* variable array.
*/
void env_array_set_environment(char **env_array)
{
char **ptr;
if (env_array == NULL)
return;
for (ptr = env_array; *ptr != NULL; ptr++) {
_env_array_putenv(*ptr);
}
}
/*
* Unset all of the environment variables in a user's current
* environment.
*
* (Note: becuae the environ array is decrementing with each
* unsetenv, only increment the ptr on a failure to unset.)
*/
void env_unset_environment(void)
{
extern char **environ;
char **ptr;
char name[256], *value;
value = xmalloc(ENV_BUFSIZE);
for (ptr = (char **)environ; *ptr != NULL; ) {
if ((_env_array_entry_splitter(*ptr, name, sizeof(name),
value, ENV_BUFSIZE)) &&
(unsetenv(name) != -1))
;
else
ptr++;
}
xfree(value);
}
/*
* Merge all of the environment variables in src_array into the
* array dest_array. Any variables already found in dest_array
* will be overwritten with the value from src_array.
*/
void env_array_merge(char ***dest_array, const char **src_array)
{
char **ptr;
char name[256], *value;
if (src_array == NULL)
return;
value = xmalloc(ENV_BUFSIZE);
for (ptr = (char **)src_array; *ptr != NULL; ptr++) {
if (_env_array_entry_splitter(*ptr, name, sizeof(name),
value, ENV_BUFSIZE))
env_array_overwrite(dest_array, name, value);
}
xfree(value);
}
/*
* Merge the environment variables in src_array beginning with "SLURM" into the
* array dest_array. Any variables already found in dest_array will be
* overwritten with the value from src_array.
*/
void env_array_merge_slurm(char ***dest_array, const char **src_array)
{
char **ptr;
char name[256], *value;
if (src_array == NULL)
return;
value = xmalloc(ENV_BUFSIZE);
for (ptr = (char **)src_array; *ptr != NULL; ptr++) {
if (_env_array_entry_splitter(*ptr, name, sizeof(name),
value, ENV_BUFSIZE) &&
(xstrncmp(name, "SLURM", 5) == 0))
env_array_overwrite(dest_array, name, value);
}
xfree(value);
}
/*
* Strip out trailing carriage returns and newlines
*/
static void _strip_cr_nl(char *line)
{
int len = strlen(line);
char *ptr;
for (ptr = line+len-1; ptr >= line; ptr--) {
if (*ptr=='\r' || *ptr=='\n') {
*ptr = '\0';
} else {
return;
}
}
}
/* Return the net count of curly brackets in a string
* '{' adds one and '}' subtracts one (zero means it is balanced).
* Special case: return -1 if no open brackets are found */
static int _bracket_cnt(char *value)
{
int count = 0, i;
for (i=0; value[i]; i++) {
if (value[i] == '{')
count++;
else if (value[i] == '}')
count--;
}
return count;
}
/*
* Load user environment from a specified file or file descriptor.
*
* This will read in a user specified file or fd, that is invoked
* via the --export-file option in sbatch. The NAME=value entries must
* be NULL separated to support special characters in the environment
* definitions.
*
* (Note: This is being added to a minor release. For the
* next major release, it might be a consideration to merge
* this functionality with that of load_env_cache and update
* env_cache_builder to use the NULL character.)
*/
char **env_array_from_file(const char *fname)
{
char *buf = NULL, *ptr = NULL, *eptr = NULL;
char *value, *p;
char **env = NULL;
char name[256];
int buf_size = BUFSIZ, buf_left;
int file_size = 0, tmp_size;
int separator = '\0';
int fd;
if (!fname)
return NULL;
/*
* If file name is a numeric value, then it is assumed to be a
* file descriptor.
*/
fd = (int)strtol(fname, &p, 10);
if ((*p != '\0') || (fd < 3) || (fd > sysconf(_SC_OPEN_MAX)) ||
(fcntl(fd, F_GETFL) < 0)) {
fd = open(fname, O_RDONLY);
if (fd == -1) {
error("Could not open user environment file %s", fname);
return NULL;
}
verbose("Getting environment variables from %s", fname);
} else
verbose("Getting environment variables from fd %d", fd);
/*
* Read in the user's environment data.
*/
buf = ptr = xmalloc(buf_size);
buf_left = buf_size;
while ((tmp_size = read(fd, ptr, buf_left))) {
if (tmp_size < 0) {
if (errno == EINTR)
continue;
error("read(environment_file): %m");
break;
}
buf_left -= tmp_size;
file_size += tmp_size;
if (buf_left == 0) {
buf_size += BUFSIZ;
xrealloc(buf, buf_size);
}
ptr = buf + file_size;
buf_left = buf_size - file_size;
}
close(fd);
/*
* Parse the buffer into individual environment variable names
* and build the environment.
*/
env = env_array_create();
value = xmalloc(ENV_BUFSIZE);
for (ptr = buf; ; ptr = eptr+1) {
eptr = strchr(ptr, separator);
if ((ptr == eptr) || (eptr == NULL))
break;
if (_env_array_entry_splitter(ptr, name, sizeof(name),
value, ENV_BUFSIZE) &&
(!_discard_env(name, value))) {
/*
* Unset the SLURM_SUBMIT_DIR if it is defined so
* that this new value does not get overwritten
* in the subsequent call to env_array_merge().
*/
if (xstrcmp(name, "SLURM_SUBMIT_DIR") == 0)
unsetenv(name);
env_array_overwrite(&env, name, value);
}
}
xfree(buf);
xfree(value);
return env;
}
/*
* Load user environment from a cache file located in
* <state_save_location>/env_username
*/
static char **_load_env_cache(const char *username)
{
char *state_save_loc, fname[MAXPATHLEN];
char *line, name[256], *value;
char **env = NULL;
FILE *fp;
int i;
state_save_loc = slurm_get_state_save_location();
i = snprintf(fname, sizeof(fname), "%s/env_cache/%s", state_save_loc,
username);
xfree(state_save_loc);
if (i < 0) {
error("Environment cache filename overflow");
return NULL;
}
if (!(fp = fopen(fname, "r"))) {
error("Could not open user environment cache at %s: %m",
fname);
return NULL;
}
verbose("Getting cached environment variables at %s", fname);
env = env_array_create();
line = xmalloc(ENV_BUFSIZE);
value = xmalloc(ENV_BUFSIZE);
while (1) {
if (!fgets(line, ENV_BUFSIZE, fp))
break;
_strip_cr_nl(line);
if (_env_array_entry_splitter(line, name, sizeof(name),
value, ENV_BUFSIZE) &&
(!_discard_env(name, value))) {
if (value[0] == '(') {
/* This is a bash function.
* It may span multiple lines */
while (_bracket_cnt(value) > 0) {
if (!fgets(line, ENV_BUFSIZE, fp))
break;
_strip_cr_nl(line);
if ((strlen(value) + strlen(line)) >
(ENV_BUFSIZE - 2))
break;
strcat(value, "\n");
strcat(value, line);
}
}
env_array_overwrite(&env, name, value);
}
}
xfree(line);
xfree(value);
fclose(fp);
return env;
}
/*
* Return an array of strings representing the specified user's default
* environment variables following a two-prongged approach.
* 1. Execute (more or less): "/bin/su - <username> -c /usr/bin/env"
* Depending upon the user's login scripts, this may take a very
* long time to complete or possibly never return
* 2. Load the user environment from a cache file. This is used
* in the event that option 1 times out. This only happens if no_cache isn't
* set. If it is set then NULL will be returned if the normal load fails.
*
* timeout value is in seconds or zero for default (2 secs)
* mode is 1 for short ("su <user>"), 2 for long ("su - <user>")
* On error, returns NULL.
*
* NOTE: The calling process must have an effective uid of root for
* this function to succeed.
*/
char **env_array_user_default(const char *username, int timeout, int mode,
bool no_cache)
{
char *line = NULL, *last = NULL, name[MAXPATHLEN], *value, *buffer;
char **env = NULL;
char *starttoken = "XXXXSLURMSTARTPARSINGHEREXXXX";
char *stoptoken = "XXXXSLURMSTOPPARSINGHEREXXXXX";
char cmdstr[256], *env_loc = NULL;
char *stepd_path = NULL;
int fd1, fd2, fildes[2], found, fval, len, rc, timeleft;
int buf_read, buf_rem, config_timeout;
pid_t child;
struct timeval begin, now;
struct pollfd ufds;
struct stat buf;
if (geteuid() != (uid_t)0) {
error("SlurmdUser must be root to use --get-user-env");
return NULL;
}
config_timeout = slurm_get_env_timeout();
if (config_timeout == 0) /* just read directly from cache */
return _load_env_cache(username);
if (stat(SUCMD, &buf))
fatal("Could not locate command: "SUCMD);
if (stat("/bin/echo", &buf))
fatal("Could not locate command: /bin/echo");
stepd_path = slurm_get_stepd_loc();
if (stat(stepd_path, &buf) == 0) {
xstrcat(stepd_path, " getenv");
env_loc = stepd_path;
} else if (stat("/bin/env", &buf) == 0)
env_loc = "/bin/env";
else if (stat("/usr/bin/env", &buf) == 0)
env_loc = "/usr/bin/env";
else
fatal("Could not location command: env");
snprintf(cmdstr, sizeof(cmdstr),
"/bin/echo; /bin/echo; /bin/echo; "
"/bin/echo %s; %s; /bin/echo %s",
starttoken, env_loc, stoptoken);
xfree(stepd_path);
if (pipe(fildes) < 0) {
fatal("pipe: %m");
return NULL;
}
child = fork();
if (child == -1) {
fatal("fork: %m");
return NULL;
}
if (child == 0) {
setenv("ENVIRONMENT", "BATCH", 1);
setpgid(0, 0);
close(0);
if ((fd1 = open("/dev/null", O_RDONLY)) == -1)
error("%s: open(/dev/null): %m", __func__);
dup2(fildes[1], 1);
close(2);
if ((fd2 = open("/dev/null", O_WRONLY)) == -1)
error("%s: open(/dev/null): %m", __func__);
if (mode == 1)
execl(SUCMD, "su", username, "-c", cmdstr, NULL);
else if (mode == 2)
execl(SUCMD, "su", "-", username, "-c", cmdstr, NULL);
else { /* Default system configuration */
#ifdef LOAD_ENV_NO_LOGIN
execl(SUCMD, "su", username, "-c", cmdstr, NULL);
#else
execl(SUCMD, "su", "-", username, "-c", cmdstr, NULL);
#endif
}
if (fd1 >= 0) /* Avoid Coverity resource leak notification */
(void) close(fd1);
if (fd2 >= 0) /* Avoid Coverity resource leak notification */
(void) close(fd2);
_exit(1);
}
close(fildes[1]);
if ((fval = fcntl(fildes[0], F_GETFL, 0)) < 0)
error("fcntl(F_GETFL) failed: %m");
else if (fcntl(fildes[0], F_SETFL, fval | O_NONBLOCK) < 0)
error("fcntl(F_SETFL) failed: %m");
gettimeofday(&begin, NULL);
ufds.fd = fildes[0];
ufds.events = POLLIN;
/* Read all of the output from /bin/su into buffer */
if (timeout == 0)
timeout = config_timeout; /* != 0 test above */
found = 0;
buf_read = 0;
buffer = xmalloc(ENV_BUFSIZE);
while (1) {
gettimeofday(&now, NULL);
timeleft = timeout * 1000;
timeleft -= (now.tv_sec - begin.tv_sec) * 1000;
timeleft -= (now.tv_usec - begin.tv_usec) / 1000;
if (timeleft <= 0) {
verbose("timeout waiting for "SUCMD" to complete");
kill(-child, 9);
break;
}
if ((rc = poll(&ufds, 1, timeleft)) <= 0) {
if (rc == 0) {
verbose("timeout waiting for "SUCMD" to complete");
break;
}
if ((errno == EINTR) || (errno == EAGAIN))
continue;
error("poll(): %m");
break;
}
if (!(ufds.revents & POLLIN)) {
if (ufds.revents & POLLHUP) { /* EOF */
found = 1; /* success */
} else if (ufds.revents & POLLERR) {
error("POLLERR");
} else {
error("poll() revents=%d", ufds.revents);
}
break;
}
buf_rem = ENV_BUFSIZE - buf_read;
if (buf_rem == 0) {
error("buffer overflow loading env vars");
break;
}
rc = read(fildes[0], &buffer[buf_read], buf_rem);
if (rc > 0)
buf_read += rc;
else if (rc == 0) { /* EOF */
found = 1; /* success */
break;
} else { /* error */
error("read(env pipe): %m");
break;
}
}
close(fildes[0]);
for (config_timeout=0; ; config_timeout++) {
kill(-child, SIGKILL); /* Typically a no-op */
if (config_timeout)
sleep(1);
if (waitpid(child, &rc, WNOHANG) > 0)
break;
if (config_timeout >= 2) {
/*
* Non-killable processes are indicative of file system
* problems. The process will remain as a zombie, but
* slurmd/salloc will not otherwise be effected.
*/
error("Failed to kill program loading user environment");
break;
}
}
if (!found) {
error("Failed to load current user environment variables");
xfree(buffer);
return no_cache ? _load_env_cache(username) : NULL;
}
/* First look for the start token in the output */
len = strlen(starttoken);
found = 0;
line = strtok_r(buffer, "\n", &last);
while (!found && line) {
if (!xstrncmp(line, starttoken, len)) {
found = 1;
break;
}
line = strtok_r(NULL, "\n", &last);
}
if (!found) {
error("Failed to get current user environment variables");
xfree(buffer);
return no_cache ? _load_env_cache(username) : NULL;
}
/* Process environment variables until we find the stop token */
len = strlen(stoptoken);
found = 0;
env = env_array_create();
line = strtok_r(NULL, "\n", &last);
value = xmalloc(ENV_BUFSIZE);
while (!found && line) {
if (!xstrncmp(line, stoptoken, len)) {
found = 1;
break;
}
if (_env_array_entry_splitter(line, name, sizeof(name),
value, ENV_BUFSIZE) &&
(!_discard_env(name, value))) {
if (value[0] == '(') {
/* This is a bash function.
* It may span multiple lines */
while (_bracket_cnt(value) > 0) {
line = strtok_r(NULL, "\n", &last);
if (!line)
break;
if ((strlen(value) + strlen(line)) >
(ENV_BUFSIZE - 2))
break;
strcat(value, "\n");
strcat(value, line);
}
}
env_array_overwrite(&env, name, value);
}
line = strtok_r(NULL, "\n", &last);
}
xfree(value);
xfree(buffer);
if (!found) {
error("Failed to get all user environment variables");
env_array_free(env);
return no_cache ? _load_env_cache(username) : NULL;
}
return env;
}
/*
* Set TRES related env vars. Set here rather than env_array_for_job() since
* we don't have array of opt values and the raw values are not stored in the
* job_desc_msg_t structure (only the strings with possibly combined TRES)
*
* opt IN - options set by command parsing
* dest IN/OUT - location to write environment variables
* het_job_offset IN - component offset into hetjob, -1 if not hetjob
*/
extern void set_env_from_opts(slurm_opt_t *opt, char ***dest,
int het_job_offset)
{
if (opt->cpus_per_gpu) {
env_array_overwrite_het_fmt(dest, "SLURM_CPUS_PER_GPU",
het_job_offset, "%d",
opt->cpus_per_gpu);
}
if (opt->gpus) {
env_array_overwrite_het_fmt(dest, "SLURM_GPUS",
het_job_offset, "%s",
opt->gpus);
}
if (opt->gpu_bind) {
env_array_overwrite_het_fmt(dest, "SLURM_GPU_BIND",
het_job_offset, "%s",
opt->gpu_bind);
}
if (opt->gpu_freq) {
env_array_overwrite_het_fmt(dest, "SLURM_GPU_FREQ",
het_job_offset, "%s",
opt->gpu_freq);
}
if (opt->gpus_per_node) {
env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_NODE",
het_job_offset, "%s",
opt->gpus_per_node);
}
if (opt->gpus_per_socket) {
env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_SOCKET",
het_job_offset, "%s",
opt->gpus_per_socket);
}
if (opt->gpus_per_task) {
env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_TASK",
het_job_offset, "%s",
opt->gpus_per_task);
}
if (opt->mem_per_gpu != NO_VAL64) {
env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_GPU",
het_job_offset, "%"PRIu64,
opt->mem_per_gpu);
}
}
extern char *find_quote_token(char *tmp, char *sep, char **last)
{
char *start;
int i, quote_single = 0, quote_double = 0;
xassert(last);
if (*last)
start = *last;
else
start = tmp;
if (start[0] == '\0')
return NULL;
for (i = 0; ; i++) {
if (start[i] == '\'') {
if (quote_single)
quote_single--;
else
quote_single++;
} else if (start[i] == '\"') {
if (quote_double)
quote_double--;
else
quote_double++;
} else if (((start[i] == sep[0]) || (start[i] == '\0')) &&
(quote_single == 0) && (quote_double == 0)) {
if (((start[0] == '\'') && (start[i-1] == '\'')) ||
((start[0] == '\"') && (start[i-1] == '\"'))) {
start++;
i -= 2;
}
if (start[i] == '\0')
*last = &start[i];
else
*last = &start[i] + 1;
start[i] = '\0';
return start;
} else if (start[i] == '\0') {
error("Improperly formed environment variable (%s)",
start);
*last = &start[i];
return start;
}
}
}