Add support for new API function slurm_step_ctx_get(), to get task id

layout for poe.

doc/man/man3/slurm_spawn.3

-.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

doc/man/man3/slurm_step_ctx_get.3

+.so man3/slurm_spawn.3

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

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");
 }