diff --git a/src/slurmd/slurmstepd/multi_prog.c b/src/slurmd/slurmstepd/multi_prog.c
index 92b7936aec986ed3f7cb23ae8caf9eee720b96e4..42f37eeab2fe6e374e6c6d01b02f6dc01090995f 100644
--- a/src/slurmd/slurmstepd/multi_prog.c
+++ b/src/slurmd/slurmstepd/multi_prog.c
@@ -296,9 +296,14 @@ fail:
return -1;
}
-/* Parse an MPMD file and determine count and layout of each task for use
- * with Cray systems. Builds the mpmd_set structure in the job record. */
-extern void multi_prog_parse(stepd_step_rec_t *job)
+/*
+ * Parse an MPMD file and determine count and layout of each task for use
+ * with Cray systems. Builds the mpmd_set structure in the job record.
+ *
+ * IN/OUT job - job step details, builds mpmd_set structure
+ * IN gtid - Array of global task IDs, indexed by node_id and task
+ */
+extern void multi_prog_parse(stepd_step_rec_t *job, uint32_t **gtid)
{
int i, j, line_num = 0, rank_id, total_ranks = 0;
char *line = NULL, *local_data = NULL;
@@ -372,17 +377,27 @@ extern void multi_prog_parse(stepd_step_rec_t *job)
#if _DEBUG
info("MPMD num_pe:%u", job->ntasks); /* Total rank count */
info("MPMD num_pe_here:%u", job->node_tasks); /* Node's rank count */
-//FIXME: DO WE HAVE OR NEED OTHER NODE'S RANK INFO?
- for (i = 0; i < job->node_tasks; i++) {
- if (!job->task) {
- error("MPMD task is NULL");
+ info("MPMD node_index:%u", job->nodeid); /* This node's index */
+ for (i = 0; i < job->nnodes; i++) {
+ if (!job->task_cnts) {
+ error("MPMD job->task_cnts is NULL");
+ break;
+ }
+ if (!job->task_cnts[i]) {
+ error("MPMD job->task_cnts[%d] is NULL", i);
+ break;
+ }
+ if (!gtid) {
+ error("MPMD gtid is NULL");
break;
}
- if (!job->task[i]) {
- error("MPMD task[%d] is NULL", i);
+ if (!gtid[i]) {
+ error("MPMD gtid[%d] is NULL", i);
break;
}
- info("MPMD placement[%d] rank:%u", i, job->task[i]->gtid);
+ for (j = 0; j < job->task_cnts[i]; j++) {
+ info("MPMD placement node[%d] rank:%u", i, gtid[i][j]);
+ }
}
#endif
job->mpmd_set = xmalloc(sizeof(mpmd_set_t));
diff --git a/src/slurmd/slurmstepd/multi_prog.h b/src/slurmd/slurmstepd/multi_prog.h
index de8e024dbc1e93ae1ba800cbaa83ae08e9e131f8..422e1b2c58906aa66b60e99ffa02d2cfb187a4a5 100644
--- a/src/slurmd/slurmstepd/multi_prog.h
+++ b/src/slurmd/slurmstepd/multi_prog.h
@@ -41,9 +41,14 @@
#include "slurmstepd_job.h"
-/* Parse an MPMD file and determine count and layout of each task for use
- * with Cray systems. Builds the mpmd_set structure in the job record. */
-extern void multi_prog_parse(stepd_step_rec_t *job);
+/*
+ * Parse an MPMD file and determine count and layout of each task for use
+ * with Cray systems. Builds the mpmd_set structure in the job record.
+ *
+ * IN/OUT job - job step details, builds mpmd_set structure
+ * IN gtid - Array of global task IDs, indexed by node_id and task
+ */
+extern void multi_prog_parse(stepd_step_rec_t *job, uint32_t **gtid);
/* Free memory associated with a job's MPMD data structure built by
* multi_prog_parse() and used for Cray system. */
diff --git a/src/slurmd/slurmstepd/slurmstepd_job.c b/src/slurmd/slurmstepd/slurmstepd_job.c
index d60d24a7c286c28c478ec72ae7e21654ea406b7c..4e6635e676758262633626ad205e3202b16fcfe2 100644
--- a/src/slurmd/slurmstepd/slurmstepd_job.c
+++ b/src/slurmd/slurmstepd/slurmstepd_job.c
@@ -73,7 +73,7 @@
static char ** _array_copy(int n, char **src);
static void _array_free(char ***array);
static void _srun_info_destructor(void *arg);
-static void _job_init_task_info(stepd_step_rec_t *job, uint32_t *gtid,
+static void _job_init_task_info(stepd_step_rec_t *job, uint32_t **gtid,
char *ifname, char *ofname, char *efname);
static void _task_info_destroy(stepd_step_task_info_t *t, uint16_t multi_prog);
@@ -215,10 +215,10 @@ _expand_stdio_filename(char *filename, int gtaskid, stepd_step_rec_t *job)
}
static void
-_job_init_task_info(stepd_step_rec_t *job, uint32_t *gtid,
+_job_init_task_info(stepd_step_rec_t *job, uint32_t **gtid,
char *ifname, char *ofname, char *efname)
{
- int i;
+ int i, node_id = job->nodeid;
char *in, *out, *err;
if (job->node_tasks == 0) {
@@ -231,10 +231,11 @@ _job_init_task_info(stepd_step_rec_t *job, uint32_t *gtid,
xmalloc(job->node_tasks * sizeof(stepd_step_task_info_t *));
for (i = 0; i < job->node_tasks; i++){
- in = _expand_stdio_filename(ifname, gtid[i], job);
- out = _expand_stdio_filename(ofname, gtid[i], job);
- err = _expand_stdio_filename(efname, gtid[i], job);
- job->task[i] = task_info_create(i, gtid[i], in, out, err);
+ in = _expand_stdio_filename(ifname, gtid[node_id][i], job);
+ out = _expand_stdio_filename(ofname, gtid[node_id][i], job);
+ err = _expand_stdio_filename(efname, gtid[node_id][i], job);
+ job->task[i] = task_info_create(i, gtid[node_id][i], in, out,
+ err);
if (!job->multi_prog) {
job->task[i]->argc = job->argc;
job->task[i]->argv = job->argv;
@@ -243,9 +244,10 @@ _job_init_task_info(stepd_step_rec_t *job, uint32_t *gtid,
if (job->multi_prog) {
// if (switch/cray)
-// multi_prog_parse(job);
+// multi_prog_parse(job, gtid);
for (i = 0; i < job->node_tasks; i++){
- multi_prog_get_argv(job->argv[1], job->env, gtid[i],
+ multi_prog_get_argv(job->argv[1], job->env,
+ gtid[node_id][i],
&job->task[i]->argc,
&job->task[i]->argv,
job->argc, job->argv);
@@ -335,6 +337,9 @@ stepd_step_rec_create(launch_tasks_request_msg_t *msg)
job->state = SLURMSTEPD_STEP_STARTING;
job->node_tasks = msg->tasks_to_launch[nodeid];
+ i = sizeof(uint16_t) * msg->nnodes;
+ job->task_cnts = xmalloc(i);
+ memcpy(job->task_cnts, msg->tasks_to_launch, i);
job->ntasks = msg->ntasks;
job->jobid = msg->job_id;
job->stepid = msg->job_step_id;
@@ -463,7 +468,7 @@ stepd_step_rec_create(launch_tasks_request_msg_t *msg)
list_append(job->sruns, (void *) srun);
- _job_init_task_info(job, msg->global_task_ids[nodeid],
+ _job_init_task_info(job, msg->global_task_ids,
msg->ifname, msg->ofname, msg->efname);
return job;
@@ -608,6 +613,7 @@ stepd_step_rec_destroy(stepd_step_rec_t *job)
xfree(job->task_epilog);
xfree(job->job_alloc_cores);
xfree(job->step_alloc_cores);
+ xfree(job->task_cnts);
xfree(job->user_name);
xfree(job);
}
diff --git a/src/slurmd/slurmstepd/slurmstepd_job.h b/src/slurmd/slurmstepd/slurmstepd_job.h
index bb88d313595d2180a56c361664e7c6ba1649a7a5..cc6f87da7a39818ed4ed2afeadedb88cd5eb21d2 100644
--- a/src/slurmd/slurmstepd/slurmstepd_job.h
+++ b/src/slurmd/slurmstepd/slurmstepd_job.h
@@ -128,6 +128,7 @@ typedef struct {
uint32_t ntasks; /* total number of tasks in current job */
uint32_t nodeid; /* relative position of this node in job */
uint32_t node_tasks; /* number of tasks on *this* node */
+ uint16_t *task_cnts; /* Number of tasks on each node in job */
uint32_t cpus_per_task; /* number of cpus desired per task */
uint32_t debug; /* debug level for job slurmd */
uint32_t job_mem; /* MB of memory reserved for the job */