_handle_core_select add handling for cyclic dist on sockets

The algorithm we had was correct for fcyclic distribution, which isn't
default. The default distribution is cyclic which means that we
cyclically place tasks on sockets, but if multiple CPUs are required for
socket we try to pick all CPUs for the task from the same socket.

Bug 14229

src/slurmctld/step_mgr.c

@@ -1947,7 +1947,8 @@ static bool _handle_core_select(step_record_t *step_ptr,
 				bitstr_t *avail_core_bitmap,
 				int job_node_inx, uint16_t sockets,
 				uint16_t cores, bool use_all_cores,
-				bool oversubscribing_cpus, int *cpu_cnt)
+				bool oversubscribing_cpus, int *cpu_cnt,
+				uint16_t cpus_per_task)
 {
 	int core_inx, i, sock_inx;
 	static int last_core_inx;
@@ -1987,7 +1988,9 @@ static bool _handle_core_select(step_record_t *step_ptr,
 					return true;
 			}
 		}
-	} else {
+	} else if (step_ptr->step_layout &&
+		   ((step_ptr->step_layout->task_dist & SLURM_DIST_SOCKMASK) ==
+		    SLURM_DIST_SOCKCFULL)) {
 		for (i = 0; i < cores; i++) {
 			if (oversubscribing_cpus)
 				core_inx = (last_core_inx + i) % cores;
@@ -1998,13 +2001,44 @@ static bool _handle_core_select(step_record_t *step_ptr,
 						     avail_core_bitmap,
 						     job_node_inx, sock_inx,
 						     core_inx, use_all_cores,
-						     oversubscribing_cpus))
-					continue;
-
+						     oversubscribing_cpus)) {
+						if (sock_inx == sockets)
+							sock_inx = 0;
+						continue;
+				}
 				if (--(*cpu_cnt) == 0)
 					return true;
 			}
 		}
+	} else { /* SLURM_DIST_SOCKCYCLIC */
+		int task_alloc_cpus = 0;
+		int start_core = 0;
+		int last_core = 0;
+		for (int j=0; j < bit_set_count(avail_core_bitmap); j++) {
+			for (sock_inx=0; sock_inx < sockets; sock_inx++) {
+				for (i = start_core; i < cores; i++) {
+					if (oversubscribing_cpus)
+						core_inx = (last_core_inx + i) % cores;
+					else
+						core_inx = i;
+					if (!_pick_step_core(step_ptr, job_resrcs_ptr,
+							     avail_core_bitmap,
+							     job_node_inx, sock_inx,
+							     core_inx, use_all_cores,
+							     oversubscribing_cpus))
+						continue;
+					if (--(*cpu_cnt) == 0)
+						return true;
+					if (++task_alloc_cpus ==
+					    cpus_per_task) {
+						task_alloc_cpus = 0;
+						last_core = i;
+						break;
+					}
+				}
+			}
+			start_core = last_core + 1;
+		}
 	}
 	return false;
 }
@@ -2017,7 +2051,7 @@ static int _pick_step_cores(step_record_t *step_ptr,
 			    uint16_t task_cnt, uint16_t cpus_per_core,
 			    int node_inx)
 {
-	uint16_t sockets, cores;
+	uint16_t sockets, cores, cpus_per_task;
 	int cpu_cnt = (int) task_cnt;
 	bool use_all_cores;
 	bitstr_t *all_gres_core_bitmap = NULL, *any_gres_core_bitmap = NULL;
@@ -2083,18 +2117,20 @@ static int _pick_step_cores(step_record_t *step_ptr,
 				job_resrcs_ptr->core_bitmap);
 		}
 	}
-
+	cpus_per_task = cpu_cnt / task_cnt;
 	/* select idle cores that fit all gres binding first */
 	if (_handle_core_select(step_ptr, job_resrcs_ptr,
 				all_gres_core_bitmap, job_node_inx,
-				sockets, cores, use_all_cores, false, &cpu_cnt))
+				sockets, cores, use_all_cores, false, &cpu_cnt,
+				cpus_per_task))
 		goto cleanup;
 
 	/* select idle cores that fit any gres binding second */
 	if (!bit_equal(all_gres_core_bitmap, any_gres_core_bitmap) &&
 	    _handle_core_select(step_ptr, job_resrcs_ptr,
 				any_gres_core_bitmap, job_node_inx,
-				sockets, cores, use_all_cores, false, &cpu_cnt))
+				sockets, cores, use_all_cores, false, &cpu_cnt,
+				cpus_per_task))
 		goto cleanup;
 
 	/* select any idle cores */
@@ -2105,7 +2141,8 @@ static int _pick_step_cores(step_record_t *step_ptr,
 		if (_handle_core_select(step_ptr, job_resrcs_ptr,
 					job_resrcs_ptr->core_bitmap,
 					job_node_inx, sockets, cores,
-					use_all_cores, false, &cpu_cnt))
+					use_all_cores, false, &cpu_cnt,
+					cpus_per_task))
 			goto cleanup;
 	}
 
@@ -2132,22 +2169,25 @@ static int _pick_step_cores(step_record_t *step_ptr,
 	/* oversubscribe cores that fit all gres binding first */
 	if (_handle_core_select(step_ptr, job_resrcs_ptr,
 				all_gres_core_bitmap, job_node_inx,
-				sockets, cores, use_all_cores, true, &cpu_cnt))
+				sockets, cores, use_all_cores, true, &cpu_cnt,
+				cpus_per_task))
 		goto cleanup;
 
 	/* oversubscribe cores that fit any gres binding second */
 	if (!bit_equal(all_gres_core_bitmap, any_gres_core_bitmap) &&
 	    _handle_core_select(step_ptr, job_resrcs_ptr,
 				any_gres_core_bitmap, job_node_inx,
-				sockets, cores, use_all_cores, true, &cpu_cnt))
+				sockets, cores, use_all_cores, true, &cpu_cnt,
+				cpus_per_task))
 		goto cleanup;
 
 	/* oversubscribe any cores */
 	if (!(step_ptr->job_ptr->bit_flags & GRES_ENFORCE_BIND) &&
 	    !bit_equal(any_gres_core_bitmap, job_resrcs_ptr->core_bitmap) &&
 	    _handle_core_select(step_ptr, job_resrcs_ptr,
 				job_resrcs_ptr->core_bitmap, job_node_inx,
-				sockets, cores, use_all_cores, true, &cpu_cnt))
+				sockets, cores, use_all_cores, true, &cpu_cnt,
+				cpus_per_task))
 		goto cleanup;