From 96fbe76930d78ece694b70dda488e7abbdad5384 Mon Sep 17 00:00:00 2001
From: "Christopher J. Morrone" <morrone2@llnl.gov>
Date: Fri, 28 Jul 2006 22:22:57 +0000
Subject: [PATCH] Beginning of the new environment variable management code.
---
doc/man/man1/salloc.1 | 2 +-
src/api/step_launch.c | 11 ++-
src/common/env.c | 183 ++++++++++++++++++++++++++++++++++++------
src/common/env.h | 96 ++++++++++++++++++++++
src/salloc/salloc.c | 7 +-
src/slaunch/opt.c | 6 +-
src/srun/srun.c | 72 ++++++++++-------
7 files changed, 314 insertions(+), 63 deletions(-)
diff --git a/doc/man/man1/salloc.1 b/doc/man/man1/salloc.1
index 53c9c11e51d..1849e31f4c2 100644
--- a/doc/man/man1/salloc.1
+++ b/doc/man/man1/salloc.1
@@ -8,7 +8,7 @@ salloc \- Obtain a SLURM job allocation (a set of nodes), execute a command, and
salloc [\fIoptions\fP] <\fIcommand\fP> [\fIcommand args\fR]
.SH "DESCRIPTION"
.LP
-salloc is used to allocate a SLURM job allocation, which is a set of resources (nodes), possibly with some set of constraints (e.g. number of processors per node). When salloc successfully obtains the requested allocation, it then runs the command specified by the user. Then, when the user specified command is complete, salloc relinquishes the job allocation.
+salloc is used to allocate a SLURM job allocation, which is a set of resources (nodes), possibly with some set of constraints (e.g. number of processors per node). When salloc successfully obtains the requested allocation, it then runs the command specified by the user. Finally, when the user specified command is complete, salloc relinquishes the job allocation.
The command may be any program the user wishes. Some typical commands are xterm, a shell script containing slaunch commands, and slaunch (see the EXAMPLES section).
.SH "OPTIONS"
diff --git a/src/api/step_launch.c b/src/api/step_launch.c
index ee7d9116c55..415388f4771 100644
--- a/src/api/step_launch.c
+++ b/src/api/step_launch.c
@@ -131,6 +131,7 @@ int slurm_step_launch (slurm_step_ctx ctx,
{
launch_tasks_request_msg_t launch;
int i;
+ char **env = NULL;
debug("Entering slurm_step_launch");
if (ctx == NULL || ctx->magic != STEP_CTX_MAGIC) {
@@ -166,8 +167,13 @@ int slurm_step_launch (slurm_step_ctx ctx,
launch.argv = params->argv;
launch.cred = ctx->step_resp->cred;
launch.job_step_id = ctx->step_resp->job_step_id;
- launch.envc = params->envc;
- launch.env = params->env;
+ env = env_array_create_for_step(ctx->step_resp,
+ "localhost",
+ 15500,
+ "127.0.0.1");
+ env_array_merge(&env, params->env);
+ launch.envc = envcount(env);
+ launch.env = env;
launch.cwd = params->cwd;
launch.nnodes = ctx->step_req->node_count;
launch.nprocs = ctx->step_req->num_tasks;
@@ -224,6 +230,7 @@ int slurm_step_launch (slurm_step_ctx ctx,
}
_launch_tasks(ctx, &launch);
+ env_array_free(env);
return SLURM_SUCCESS;
}
diff --git a/src/common/env.c b/src/common/env.c
index f4ac432b41c..0a3e422ff02 100644
--- a/src/common/env.c
+++ b/src/common/env.c
@@ -529,10 +529,89 @@ int setup_env(env_t *env)
return SLURM_SUCCESS;
}
-#if 0
/**********************************************************************
- * FIXME - Not yet fully implemented, still in the planning phase
+ * From here on are the new environment variable management functions,
+ * used by the "new" commands: salloc, sbatch, an slaunch.
**********************************************************************/
+
+/*
+ * Return a string representation of an array of uint32_t elements.
+ * Each value in the array is printed in decimal notation and elements
+ * are seperated 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().
+ */
+static char *_uint32_array_to_str(int array_len, const uint32_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 of uint32_t, 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().
+ */
+static char *_uint32_compressed_to_str(uint32_t array_len,
+ const uint32_t *array,
+ const uint32_t *array_reps)
+{
+ int i;
+ char *sep = ","; /* seperator */
+ char *str = xstrdup("");
+
+ 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;
+}
+
/*
* Create an array of pointers to environment variables strings relevant
* to a SLURM job allocation. The array is terminated by a NULL pointer,
@@ -551,12 +630,18 @@ char **
env_array_create_for_job(const resource_allocation_response_msg_t *alloc)
{
char **ptr;
+ char *tmp;
ptr = env_array_create();
env_array_append(&ptr, "SLURM_JOB_ID", "%u", alloc->job_id);
env_array_append(&ptr, "SLURM_JOB_NUM_NODES", "%u", alloc->node_cnt);
env_array_append(&ptr, "SLURM_JOB_NODELIST", "%s", alloc->node_list);
- env_array_append(&ptr, "SLURM_JOB_CPUS_PER_NODE", "%s", FIXME);
+
+ tmp = _uint32_compressed_to_str((uint32_t)alloc->num_cpu_groups,
+ alloc->cpus_per_node,
+ alloc->cpu_count_reps);
+ env_array_append(&ptr, "SLURM_JOB_CPUS_PER_NODE", "%s", tmp);
+ xfree(tmp);
return ptr;
}
@@ -593,19 +678,22 @@ env_array_create_for_step(const job_step_create_response_msg_t *step,
const char *ip_addr_str)
{
char **ptr;
+ char *tmp;
+ tmp = _uint32_array_to_str(step->step_layout->node_cnt,
+ step->step_layout->tasks);
ptr = env_array_create();
env_array_append(&ptr, "SLURM_STEP_ID", "%u", step->job_step_id);
env_array_append(&ptr, "SLURM_STEP_NUM_NODES",
"%hu", step->step_layout->node_cnt);
env_array_append(&ptr, "SLURM_STEP_NUM_TASKS",
- "%s", step->step_layout->task_cnt);
- env_array_append(&ptr, "SLURM_STEP_TASKS_PER_NODE", "%s", FIXME);
+ "%u", step->step_layout->task_cnt);
+ env_array_append(&ptr, "SLURM_STEP_TASKS_PER_NODE", "%s", tmp);
env_array_append(&ptr, "SLURM_STEP_LAUNCHER_HOSTNAME",
"%s", launcher_hostname);
env_array_append(&ptr, "SLURM_STEP_LAUNCHER_PORT",
"%hu", launcher_port);
- env_array_appent(&ptr, "SLURM_STEP_LAUNCHER_IPADDR",
+ env_array_append(&ptr, "SLURM_STEP_LAUNCHER_IPADDR",
"%s", ip_addr_str);
/* OBSOLETE */
@@ -613,15 +701,16 @@ env_array_create_for_step(const job_step_create_response_msg_t *step,
env_array_append(&ptr, "SLURM_NNODES",
"%hu", step->step_layout->node_cnt);
env_array_append(&ptr, "SLURM_NPROCS",
- "%s", step->step_layout->task_cnt);
- env_array_append(&ptr, "SLURM_TASKS_PER_NODE", "%s", FIXME);
+ "%u", step->step_layout->task_cnt);
+ env_array_append(&ptr, "SLURM_TASKS_PER_NODE", "%s", tmp);
env_array_append(&ptr, "SLURM_SRUN_COMM_HOST",
"%s", launcher_hostname);
env_array_append(&ptr, "SLURM_SRUN_COMM_PORT",
"%hu", launcher_port);
- env_array_appent(&ptr, "SLURM_LAUNCH_NODE_IPADDR",
+ env_array_append(&ptr, "SLURM_LAUNCH_NODE_IPADDR",
"%s", ip_addr_str);
+ xfree(tmp);
return ptr;
}
@@ -657,6 +746,7 @@ char **env_array_create(void)
return env_array;
}
+
/*
* 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
@@ -671,10 +761,9 @@ int env_array_append(char ***array_ptr, const char *name,
char **ep = NULL;
char *str = NULL;
va_list ap;
- int rc;
buf[0] = '\0';
- if (array_ptr == NULL || *array == NULL) {
+ if (array_ptr == NULL || *array_ptr == NULL) {
return 0;
}
@@ -709,10 +798,9 @@ int env_array_overwrite(char ***array_ptr, const char *name,
char **ep = NULL;
char *str = NULL;
va_list ap;
- int rc;
buf[0] = '\0';
- if (array_ptr == NULL || *array == NULL) {
+ if (array_ptr == NULL || *array_ptr == NULL) {
return 0;
}
@@ -735,7 +823,7 @@ int env_array_overwrite(char ***array_ptr, const char *name,
char **env_array_copy(const char **array)
{
-
+ return NULL;
}
/*
@@ -745,37 +833,82 @@ char **env_array_copy(const char **array)
*/
void env_array_merge(char ***dest_array, const char **src_array)
{
- char **ptr;
+/* char **ptr; */
if (src_array == NULL)
return;
- for (ptr = array; *ptr != NULL; ptr++) {
- env_array_overwrite(dest_array, *ptr);
- }
+/* FIXME - split the env entry on the =, then give to overwrite */
+/* for (ptr = (char **)src_array; *ptr != NULL; ptr++) { */
+/* env_array_overwrite(dest_array, *ptr); */
+/* } */
}
+/*
+ * Free the memory used by an environment variable array.
+ */
void env_array_free(char **env_array)
{
char **ptr;
- if (array == NULL)
+ if (env_array == NULL)
return;
- for (ptr = array; *ptr != NULL; ptr++) {
+ for (ptr = env_array; *ptr != NULL; ptr++) {
xfree(*ptr);
}
- xfree(array);
+ 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 = index(entry, '=');
+ len = ptr - entry;
+ if (len > name_len-1)
+ return 0;
+ strncpy(name, entry, len);
+ name[len] = '\0';
+
+ ptr = ptr + 1;
+ len = strlen(ptr);
+ if (len > value_len-1)
+ return 0;
+ strncpy(value, ptr, len);
+ value[len] = '\0';
+
+ return 1;
}
/*
* Work similarly to putenv() (from C stdlib), but uses setenv()
* under the covers. This avoids having pointers from the global
- * array "environ" to "string".
+ * array "environ" into "string".
+ *
+ * Return 1 on success, 0 on failure.
*/
static int _env_array_putenv(const char *string)
{
+ char name[1024];
+ char value[1024];
+ if (!_env_array_entry_splitter(string, name, 1024, value, 1024))
+ return 0;
+ if (setenv(name, value, 1) == -1)
+ return 0;
+
+ return 1;
}
/*
@@ -786,12 +919,10 @@ void env_array_set_environment(char **env_array)
{
char **ptr;
- if (array == NULL)
+ if (env_array == NULL)
return;
- for (ptr = array; *ptr != NULL; ptr++) {
+ for (ptr = env_array; *ptr != NULL; ptr++) {
_env_array_putenv(*ptr);
}
}
-
-#endif
diff --git a/src/common/env.h b/src/common/env.h
index e6710b4d1bc..7f275defc9d 100644
--- a/src/common/env.h
+++ b/src/common/env.h
@@ -75,4 +75,100 @@ void unsetenvp(char **env, const char *name);
char * getenvp(char **env, const char *name);
int setup_env(env_t *env);
+/**********************************************************************
+ * Newer environment variable handling scheme
+ **********************************************************************/
+/*
+ * Create an array of pointers to environment variables strings relevant
+ * to a SLURM job allocation. 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_NUM_NODES
+ * SLURM_JOB_NODELIST
+ * SLURM_JOB_CPUS_PER_NODE
+ *
+ * Sets OBSOLETE variables:
+ * ? probably only needed for users...
+ */
+char **
+env_array_create_for_job(const resource_allocation_response_msg_t *alloc);
+
+/*
+ * Create an array of pointers to environment variables strings relevant
+ * to a SLURM job step. The array is terminated by a NULL pointer,
+ * and thus is suitable for use by execle() and other env_array_* functions.
+ *
+ * Sets variables:
+ * SLURM_STEP_ID
+ * SLURM_STEP_NUM_NODES
+ * SLURM_STEP_NUM_TASKS
+ * SLURM_STEP_TASKS_PER_NODE
+ * SLURM_STEP_LAUNCHER_HOSTNAME
+ * SLURM_STEP_LAUNCHER_PORT
+ * SLURM_STEP_LAUNCHER_IPADDR
+ *
+ * Sets OBSOLETE variables:
+ * SLURM_STEPID
+ * SLURM_NNODES
+ * SLURM_NPROCS
+ * SLURM_NODELIST
+ * SLURM_TASKS_PER_NODE
+ * SLURM_SRUN_COMM_HOST
+ * SLURM_SRUN_COMM_PORT
+ * SLURM_LAUNCH_NODE_IPADDR
+ *
+ */
+char **
+env_array_create_for_step(const job_step_create_response_msg_t *step,
+ const char *launcher_hostname,
+ uint16_t launcher_port,
+ const char *ip_addr_str);
+
+/*
+ * Return an empty environment variable array (contains a single
+ * pointer to NULL).
+ */
+char **env_array_create(void);
+
+/*
+ * 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);
+
+/*
+ * Free the memory used by an environment variable array.
+ */
+void env_array_free(char **env_array);
+
+/*
+ * 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_fmt, ...);
+
+/*
+ * 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_fmt, ...);
+
+/*
+ * Set all of the environment variables in a supplied environment
+ * variable array.
+ */
+void env_array_set_environment(char **env_array);
+
#endif
diff --git a/src/salloc/salloc.c b/src/salloc/salloc.c
index 9763b69123b..bb4e4112b77 100644
--- a/src/salloc/salloc.c
+++ b/src/salloc/salloc.c
@@ -44,6 +44,7 @@
#include "src/common/xmalloc.h"
#include "src/common/xsignal.h"
#include "src/common/read_config.h"
+#include "src/common/env.h"
#include "src/salloc/opt.h"
#include "src/salloc/msg.h"
@@ -62,6 +63,7 @@ int main(int argc, char *argv[])
time_t before, after;
salloc_msg_thread_t *msg_thr;
int rc;
+ char **env = NULL;
log_init(xbasename(argv[0]), logopt, 0, NULL);
if (initialize_and_process_args(argc, argv) < 0) {
@@ -108,9 +110,10 @@ int main(int argc, char *argv[])
/*
* Run the user's command.
*/
- setenvfs("SLURM_JOBID=%d", alloc->job_id);
- setenvfs("SLURM_NNODES=%d", alloc->node_cnt);
+ env = env_array_create_for_job(alloc);
+ env_array_set_environment(env);
rc = run_command(command_argv);
+ env_array_free(env);
/*
* Relinquish the job allocation.
diff --git a/src/slaunch/opt.c b/src/slaunch/opt.c
index bf8e03871f7..b3880cfccde 100644
--- a/src/slaunch/opt.c
+++ b/src/slaunch/opt.c
@@ -732,7 +732,7 @@ struct env_vars {
};
env_vars_t env_vars[] = {
- {"SLURM_JOBID", OPT_INT, &opt.jobid, &opt.jobid_set },
+ {"SLURM_JOB_ID", OPT_INT, &opt.jobid, &opt.jobid_set },
{"SLAUNCH_JOBID", OPT_INT, &opt.jobid, &opt.jobid_set },
{"SLURMD_DEBUG", OPT_INT, &opt.slurmd_debug, NULL },
{"SLAUNCH_CPUS_PER_TASK",OPT_INT, &opt.cpus_per_task, &opt.cpus_set },
@@ -745,8 +745,8 @@ env_vars_t env_vars[] = {
{"SLAUNCH_GEOMETRY", OPT_GEOMETRY, NULL, NULL },
{"SLAUNCH_KILL_BAD_EXIT",OPT_INT, &opt.kill_bad_exit, NULL },
{"SLAUNCH_LABELIO", OPT_INT, &opt.labelio, NULL },
- {"SLURM_NNODES", OPT_INT, &opt.num_nodes, NULL },
- {"SLAUNCH_NNODES", OPT_INT, &opt.num_nodes, &opt.num_nodes_set},
+ {"SLURM_JOB_NUM_NODES", OPT_INT, &opt.num_nodes, NULL },
+ {"SLAUNCH_NUM_NODES", OPT_INT, &opt.num_nodes, &opt.num_nodes_set},
{"SLAUNCH_NO_ROTATE", OPT_NO_ROTATE, NULL, NULL },
{"SLAUNCH_NPROCS", OPT_INT, &opt.num_tasks, &opt.num_tasks_set},
{"SLAUNCH_OVERCOMMIT", OPT_OVERCOMMIT,NULL, NULL },
diff --git a/src/srun/srun.c b/src/srun/srun.c
index bd2c71d5e76..e89e3dd302a 100644
--- a/src/srun/srun.c
+++ b/src/srun/srun.c
@@ -107,7 +107,7 @@ static int _run_job_script(srun_job_t *job, env_t *env);
static void _set_prio_process_env(void);
static int _set_rlimit_env(void);
static int _set_umask_env(void);
-static char *_task_count_string(srun_job_t *job);
+static char *_uint32_array_to_str(int count, const uint32_t *array);
static void _switch_standalone(srun_job_t *job);
static int _become_user (void);
static int _print_script_exit_status(const char *argv0, int status);
@@ -290,7 +290,8 @@ int srun(int ac, char **av)
env->select_jobinfo = job->select_jobinfo;
env->nhosts = job->nhosts;
env->nodelist = job->nodelist;
- env->task_count = _task_count_string (job);
+ env->task_count = _uint32_array_to_str(
+ job->nhosts, job->step_layout->tasks);
env->jobid = job->jobid;
env->stepid = job->stepid;
}
@@ -403,35 +404,47 @@ int srun(int ac, char **av)
exit(exitcode);
}
-static char *
-_task_count_string (srun_job_t *job)
+/*
+ * Return a string representation of an array of uint32_t elements.
+ * Each value in the array is printed in decimal notation and elements
+ * are seperated 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().
+ */
+static char *_uint32_array_to_str(int array_len, const uint32_t *array)
{
- int i, last_val, last_cnt;
- char tmp[16];
- char *str = xstrdup ("");
- if(job->step_layout->tasks == NULL)
- return (str);
- last_val = job->step_layout->tasks[0];
- last_cnt = 1;
- for (i=1; i<job->nhosts; i++) {
- if (last_val == job->step_layout->tasks[i])
- last_cnt++;
- else {
- if (last_cnt > 1)
- sprintf(tmp, "%d(x%d),", last_val, last_cnt);
- else
- sprintf(tmp, "%d,", last_val);
- xstrcat(str, tmp);
- last_val = job->step_layout->tasks[i];
- last_cnt = 1;
+ 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;
}
- if (last_cnt > 1)
- sprintf(tmp, "%d(x%d)", last_val, last_cnt);
- else
- sprintf(tmp, "%d", last_val);
- xstrcat(str, tmp);
- return (str);
+
+ return str;
}
static void
@@ -904,7 +917,8 @@ static int _run_job_script (srun_job_t *job, env_t *env)
env->jobid = job->jobid;
env->nhosts = job->nhosts;
env->nodelist = job->nodelist;
- env->task_count = _task_count_string(job);
+ env->task_count = _uint32_array_to_str(
+ job->nhosts, job->step_layout->tasks);
}
if (setup_env(env) != SLURM_SUCCESS)
--
GitLab