From c0098be5d1de6f503eaa0d9ece1d35e9da6949ee Mon Sep 17 00:00:00 2001
From: Mark Grondona <mgrondona@llnl.gov>
Date: Mon, 13 Jan 2003 23:05:00 +0000
Subject: [PATCH] o move _compute_task_count() into job.c to allow -c option
to affect number of tasks in batch job steps.
---
src/srun/job.c | 35 +++++++++++++++++++++++++++++++++--
src/srun/srun.c | 25 ++++++++-----------------
2 files changed, 41 insertions(+), 19 deletions(-)
diff --git a/src/srun/job.c b/src/srun/job.c
index 9b7b684c6af..0b298857a9f 100644
--- a/src/srun/job.c
+++ b/src/srun/job.c
@@ -48,7 +48,9 @@ typedef struct allocation_info {
uint32_t jobid;
uint32_t stepid;
char *nodelist;
+ int nnodes;
slurm_addr *addrs;
+ int num_cpu_groups;
int *cpus_per_node;
int *cpu_count_reps;
} allocation_info_t;
@@ -62,6 +64,27 @@ _estimate_nports(int nclients, int cli_per_port)
return d.rem > 0 ? d.quot + 1 : d.quot;
}
+static int
+_compute_task_count(allocation_info_t *info)
+{
+ int i, cnt = 0;
+
+ if (opt.cpus_set) {
+ for (i = 0; i < info->num_cpu_groups; i++)
+ cnt += ( info->cpu_count_reps[i] *
+ (info->cpus_per_node[i]/opt.cpus_per_task));
+ }
+
+ return (cnt < info->nnodes) ? info->nnodes : cnt;
+}
+
+static void
+_set_nprocs(allocation_info_t *info)
+{
+ if (!opt.nprocs_set)
+ opt.nprocs = _compute_task_count(info);
+}
+
static job_t *
_job_create_internal(allocation_info_t *info)
@@ -73,6 +96,10 @@ _job_create_internal(allocation_info_t *info)
hostlist_t hl;
job_t *job;
+ /* Reset nprocs if necessary
+ */
+ _set_nprocs(info);
+
job = xmalloc(sizeof(*job));
slurm_mutex_init(&job->state_mutex);
@@ -120,7 +147,7 @@ _job_create_internal(allocation_info_t *info)
job->outbuf = (cbuf_t *) xmalloc(opt.nprocs * sizeof(cbuf_t));
job->errbuf = (cbuf_t *) xmalloc(opt.nprocs * sizeof(cbuf_t));
job->inbuf = (cbuf_t *) xmalloc(opt.nprocs * sizeof(cbuf_t));
- job->stdin_eof = (bool *) xmalloc(opt.nprocs * sizeof(bool));
+ job->stdin_eof = (bool *) xmalloc(opt.nprocs * sizeof(bool));
/* nhost host states */
@@ -177,8 +204,10 @@ job_create_allocation(resource_allocation_response_msg_t *resp)
allocation_info_t *info = xmalloc(sizeof(*info));
info->nodelist = resp->node_list;
+ info->nnodes = resp->node_cnt;
info->jobid = resp->job_id;
info->stepid = NO_VAL;
+ info->num_cpu_groups = resp->num_cpu_groups;
info->cpus_per_node = resp->cpus_per_node;
info->cpu_count_reps = resp->cpu_count_reps;
info->addrs = resp->node_addr;
@@ -225,9 +254,11 @@ job_create_noalloc(void)
info->jobid = (uint32_t) (lrand48() % 65550L + 1L);
info->stepid = 0;
info->nodelist = opt.nodelist;
+ info->nnodes = hostlist_count(hl);
- if (opt.nprocs < hostlist_count(hl))
+ /* if (opt.nprocs < info->nnodes)
opt.nprocs = hostlist_count(hl);
+ */
hostlist_destroy(hl);
info->cpus_per_node = &cpn;
diff --git a/src/srun/srun.c b/src/srun/srun.c
index e67b563d09b..d81708b9ea6 100644
--- a/src/srun/srun.c
+++ b/src/srun/srun.c
@@ -329,25 +329,10 @@ _allocate_nodes(void)
}
sigaction(SIGINT, &old_action, NULL);
}
- if (!opt.nprocs_set) /* can vary by min-max node count */
- opt.nprocs = _compute_task_count(resp);
-
+
return resp;
}
-static int _compute_task_count(allocation_resp *resp)
-{
- int i, cnt = 0;
-
- if (opt.cpus_set) {
- for (i=0; i<resp->num_cpu_groups; i++)
- cnt += (resp->cpu_count_reps[i] *
- (resp->cpus_per_node[i]/opt.cpus_per_task));
- }
- if (cnt < resp->node_cnt)
- cnt = resp->node_cnt;
- return cnt;
-}
static void
_sig_kill_alloc(int signum)
@@ -773,11 +758,17 @@ _set_batch_script_env(uint32_t jobid, uint32_t node_cnt)
opt.nprocs = node_cnt;
}
- if (setenvf("SLURM_NPROCS=%u", opt.nprocs)) {
+ if (opt.nprocs_set && setenvf("SLURM_NPROCS=%u", opt.nprocs)) {
error("Unable to set SLURM_NPROCS environment variable");
return -1;
}
+ if ( (opt.cpus_per_task > 0) &&
+ setenvf("SLURM_CPUS_PER_TASK=%u", opt.cpus_per_task)) {
+ error("Unable to set SLURM_CPUS_PER_TASK");
+ return -1;
+ }
+
if (opt.distribution != SRUN_DIST_UNKNOWN) {
dist = (opt.distribution == SRUN_DIST_BLOCK) ?
"block" : "cyclic";
--
GitLab