diff --git a/doc/man/man3/slurm_spawn.3 b/doc/man/man3/slurm_spawn.3
index 5ffe916337c1cecfa6750a621b3de5f1a6ddbece..cfd931c7bcd0fa1af6d140c443ca69237ad16078 100644
--- a/doc/man/man3/slurm_spawn.3
+++ b/doc/man/man3/slurm_spawn.3
@@ -1,4 +1,4 @@
-.TH "Slurm API" "3" "June 2004" "Morris Jette" "Slurm task spawn functions"
+.TH "Slurm API" "3" "July 2004" "Morris Jette" "Slurm task spawn functions"
.SH "NAME"
slurm_spawn \- Slurm task spawn functions
.SH "SYNTAX"
@@ -12,6 +12,16 @@ slurm_step_ctx \fBslurm_step_ctx_create\fR (
.br
);
.LP
+int \fBslurm_step_ctx_get\fR (
+.br
+ slurm_step_ctx \fIctx\fP,
+.br
+ int \fIctx_key\fP,
+.br
+ ...
+.br
+);
+.LP
int \fBslurm_step_ctx_set\fR (
.br
slurm_step_ctx \fIctx\fP,
@@ -55,7 +65,8 @@ Job step context. Created by \fBslurm_step_ctx_create\fR, used in subsequent
function calls, and destroyed by \fBslurm_step_ctx_destroy\fR.
.TP
\fIctx_key\fP
-Identifies the fields in \fIctx\fP to be set by \fBslurm_step_ctx_set\fR.
+Identifies the fields in \fIctx\fP to be collected by \fBslurm_step_ctx_get\fR
+or set by \fBslurm_step_ctx_set\fR.
.TP
\fIfd_array\fP
Array of socket file descriptors to be connected to the initiated tasks.
@@ -71,6 +82,11 @@ Signal to be sent to the spawned tasks.
leaks call \fBslurm_step_ctx_destroy\fR when the use of this context is
finished.
.LP
+\fBslurm_step_ctx_get\fR Get values from a job step context.
+\fIctx_key\fP identifies the fields to be gathered from the job step context.
+Subsequent arguments to this function are dependent upon the value
+of \fIctx_key\fP. See the \fBCONTEXT KEYS\fR section for details.
+.LP
\fBslurm_step_ctx_set\fR Set values in a job step context.
\fIctx_key\fP identifies the fields to be set in the job step context.
Subsequent arguments to this function are dependent upon the value
@@ -108,6 +124,25 @@ Sets the environment variable count and values for the executable.
Accepts two additional arguments, the first of type int and
the second of type char **. By default the current environment
variables are copied to started task's environment.
+.TP
+\fBSLURM_STEP_CTX_TASKS\fR
+Get the number of tasks per node for a given job.
+Accepts one additional argument of type uint32_t **.
+This argument will be set to point to an array with the
+task counts of each node in an element of the array.
+See \fBSLURM_STEP_CTX_TID\fR below to determine the
+task ID numbers associated with each of those tasks.
+.TP
+\fBSLURM_STEP_CTX_TID\fR
+Get the task ID numbers associated with the tasks allocated to
+a specific node.
+Accepts two additional arguments, the first of type int and
+the second of type uint32_t **. The first argument identifies
+the node number of interest (zero origin). The second argument
+will be set to point to an array with the task ID numbers of
+each task allocated to the node (also zero origin).
+See \fBSLURM_STEP_CTX_TASKS\fR above to determine how many
+tasks are associated with each node.
.SH "RETURN VALUE"
.LP
For \fB slurm_step_ctx_create\fR a context is return upon success. On error
@@ -159,7 +194,7 @@ int main (int argc, char *argv[])
.br
{
.br
- int i, nodes = 1, tasks;
+ int i, j, nodes = 1, tasks;
.br
job_desc_msg_t job_req;
.br
@@ -172,6 +207,8 @@ int main (int argc, char *argv[])
char *task_argv[2];
.br
int *fd_array;
+.br
+ uint32_t *task_cnt; *tid;
.LP
if (argc > 1) {
.br
@@ -222,11 +259,7 @@ int main (int argc, char *argv[])
exit(1);
.br
}
-.br
- printf("Starting %d tasks on %d nodes\n",
-.br
- step_req.num_tasks, step_req.node_count);
-.br
+.LP
task_argv[0] = "/bin/hostname";
.br
if (slurm_step_ctx_set(ctx, SLURM_STEP_CTX_ARGS,
@@ -248,6 +281,34 @@ int main (int argc, char *argv[])
exit(1);
.br
}
+.LP
+ printf("Started %d tasks on %d nodes\\n",
+.br
+ step_req.num_tasks, step_req.node_count);
+.br
+ if (slurm_step_ctx_get(ctx, SLURM_STEP_CTX_TASKS, &task_cnt)) {
+.br
+ slurm_perror("slurm_step_ctx_create");
+.br
+ exit(1);
+.br
+ }
+.br
+ for (i=0; i<step_req.node_count; i++) {
+.br
+ if (slurm_step_ctx_get(ctx, SLURM_STEP_CTX_TID, i, &tid)) {
+.br
+ slurm_perror("slurm_step_ctx_create");
+.br
+ exit(1);
+.br
+ }
+.br
+ for (j=0; j<task_cnt[i]; j++)
+.br
+ printf("tid[%d][%d] = %u\\n", i, j, tid[j]);
+.br
+
.LP
/* Interact with spawned tasks as desired */
.br
diff --git a/doc/man/man3/slurm_step_ctx_get.3 b/doc/man/man3/slurm_step_ctx_get.3
new file mode 100644
index 0000000000000000000000000000000000000000..02d58710e5d6dac5a874510de9300bd44a2ed25a
--- /dev/null
+++ b/doc/man/man3/slurm_step_ctx_get.3
@@ -0,0 +1 @@
+.so man3/slurm_spawn.3
diff --git a/slurm/slurm.h.in b/slurm/slurm.h.in
index 0f54629b1416cfe411b69f87c9b231251e23c788..ea4f0cd4031e2ee4ee8bd5f921bc7380f02ca2c7 100644
--- a/slurm/slurm.h.in
+++ b/slurm/slurm.h.in
@@ -162,11 +162,14 @@ enum node_states {
* Values can be can be ORed */
#define SHOW_ALL 1 /* Show info for "hidden" partitions */
-/* Define keys for ctx_key argument of slurm_step_ctx_set() */
+/* Define keys for ctx_key argument of slurm_step_ctx_get() and
+ * slurm_step_ctx_set() */
enum ctx_keys {
SLURM_STEP_CTX_ARGS, /* set argument count and values */
SLURM_STEP_CTX_CHDIR, /* set directory to run from */
- SLURM_STEP_CTX_ENV /* set environment variable count and values */
+ SLURM_STEP_CTX_ENV, /* set environment variable count and values */
+ SLURM_STEP_CTX_TASKS, /* get array of task count on each node */
+ SLURM_STEP_CTX_TID /* get array of task IDs for specified node */
};
/*****************************************************************************\
@@ -648,7 +651,13 @@ extern int slurm_complete_job_step PARAMS((
extern slurm_step_ctx slurm_step_ctx_create PARAMS((
job_step_create_request_msg_t *step_req));
-
+/*
+ * slurm_step_ctx_get - get parameters from a job step context.
+ * IN ctx - job step context generated by slurm_step_ctx_create
+ * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
+ */
+extern int slurm_step_ctx_get PARAMS((slurm_step_ctx ctx,
+ int ctx_key, ...));
/*
* slurm_step_ctx_set - set parameters in job step context.
* IN ctx - job step context generated by slurm_step_ctx_create
diff --git a/src/api/spawn.c b/src/api/spawn.c
index c0894aad2f5a8e16d9ead5da5e3d2ba99114b652..08fd5cd14af7aca7b6ae5baed2fa23c6a42eb8e4 100644
--- a/src/api/spawn.c
+++ b/src/api/spawn.c
@@ -145,6 +145,52 @@ slurm_step_ctx_create (job_step_create_request_msg_t *step_req)
return rc;
}
+/*
+ * slurm_step_ctx_get - get parameters from a job step context.
+ * IN ctx - job step context generated by slurm_step_ctx_create
+ * RET SLURM_SUCCESS or SLURM_ERROR (with slurm_errno set)
+ */
+extern int
+slurm_step_ctx_get (slurm_step_ctx ctx, int ctx_key, ...)
+{
+ va_list ap;
+ int rc = SLURM_SUCCESS;
+ uint32_t node_inx;
+ uint32_t **array_pptr = (uint32_t **) NULL;
+
+ if ((ctx == NULL) ||
+ (ctx->magic != STEP_CTX_MAGIC)) {
+ slurm_seterrno(EINVAL);
+ return SLURM_ERROR;
+ }
+
+ va_start(ap, ctx_key);
+ switch (ctx_key) {
+ case SLURM_STEP_CTX_TASKS:
+ array_pptr = (uint32_t **) va_arg(ap, void *);
+ *array_pptr = ctx->tasks;
+ break;
+
+ case SLURM_STEP_CTX_TID:
+ node_inx = va_arg(ap, uint32_t);
+ if ((node_inx < 0) || (node_inx > ctx->nhosts)) {
+ slurm_seterrno(EINVAL);
+ rc = SLURM_ERROR;
+ break;
+ }
+ array_pptr = (uint32_t **) va_arg(ap, void *);
+ *array_pptr = ctx->tids[node_inx];
+ break;
+
+ default:
+ slurm_seterrno(EINVAL);
+ rc = SLURM_ERROR;
+ }
+ va_end(ap);
+
+ return rc;
+}
+
/*
* slurm_step_ctx_set - set parameters in job step context.
* IN ctx - job step context generated by slurm_step_ctx_create
@@ -546,7 +592,7 @@ static int _envcount(char **env)
/* dump the contents of a job step context */
static void _dump_ctx(slurm_step_ctx ctx)
{
- int i;
+ int i, j;
if ((ctx == NULL) ||
(ctx->magic != STEP_CTX_MAGIC)) {
@@ -584,10 +630,13 @@ static void _dump_ctx(slurm_step_ctx ctx)
}
}
- for (i=0; i<ctx->nhosts; i++)
- printf("host=%s cpus=%u tasks=%u tid[0]=%u\n",
- ctx->host[i], ctx->cpus[i], ctx->tasks[i],
- ctx->tids[i][0]);
+ for (i=0; i<ctx->nhosts; i++) {
+ printf("host=%s cpus=%u tasks=%u",
+ ctx->host[i], ctx->cpus[i], ctx->tasks[i]);
+ for (j=0; j<ctx->tasks[i]; j++)
+ printf(" tid[%d]=%u", j, ctx->tids[i][j]);
+ printf("\n");
+ }
printf("\n");
}