/* The job has specialized cores, synchronize user map with available cores */
static void _validate_map(launch_tasks_request_msg_t *req, char *avail_mask)
{
char *tmp_map, *save_ptr = NULL, *tok;
cpu_set_t avail_cpus;
bool superset = true;
CPU_ZERO(&avail_cpus);
(void) str_to_cpuset(&avail_cpus, avail_mask);
tmp_map = xstrdup(req->cpu_bind);
tok = strtok_r(tmp_map, ",", &save_ptr);
while (tok) {
int i = atoi(tok);
if (!CPU_ISSET(i, &avail_cpus)) {
/* The task's CPU map is completely invalid.
* Disable CPU map. */
superset = false;
break;
}
tok = strtok_r(NULL, ",", &save_ptr);
}
xfree(tmp_map);
if (!superset) {
info("task/affinity: Ignoring user CPU binding outside of job "
"step allocation");
req->cpu_bind_type &= (~CPU_BIND_MAP);
req->cpu_bind_type |= CPU_BIND_MASK;
xfree(req->cpu_bind);
req->cpu_bind = xstrdup(avail_mask);
}
}
/* The job has specialized cores, synchronize user mask with available cores */
static void _validate_mask(launch_tasks_request_msg_t *req, char *avail_mask)
{
char *new_mask = NULL, *save_ptr = NULL, *tok;
cpu_set_t avail_cpus, task_cpus;
bool superset = true;
CPU_ZERO(&avail_cpus);
(void) str_to_cpuset(&avail_cpus, avail_mask);
tok = strtok_r(req->cpu_bind, ",", &save_ptr);
while (tok) {
int i, overlaps = 0;
char mask_str[1 + CPU_SETSIZE / 4];
CPU_ZERO(&task_cpus);
(void) str_to_cpuset(&task_cpus, tok);
for (i = 0; i < CPU_SETSIZE; i++) {
if (!CPU_ISSET(i, &task_cpus))
continue;
if (CPU_ISSET(i, &avail_cpus)) {
overlaps++;
} else {
CPU_CLR(i, &task_cpus);
superset = false;
}
}
if (overlaps == 0) {
/* The task's CPU mask is completely invalid.
* Give it all allowed CPUs. */
for (i = 0; i < CPU_SETSIZE; i++) {
if (CPU_ISSET(i, &avail_cpus))
CPU_SET(i, &task_cpus);
}
}
cpuset_to_str(&task_cpus, mask_str);
if (new_mask)
xstrcat(new_mask, ",");
xstrcat(new_mask, mask_str);
tok = strtok_r(NULL, ",", &save_ptr);
}
if (!superset) {
info("task/affinity: Ignoring user CPU binding outside of job "
"step allocation");
}
xfree(req->cpu_bind);
req->cpu_bind = new_mask;
}
int slurm_get_cpuset(char *path, pid_t pid, size_t size, cpu_set_t *mask)
{
int fd, rc;
char file_path[PATH_MAX];
char mstr[1 + CPU_SETSIZE * 4];
snprintf(file_path, sizeof(file_path), "%s/cpus", path);
fd = open(file_path, O_RDONLY);
if (fd < 0) {
error("open(%s): %m", file_path);
return -1;
}
rc = read(fd, mstr, sizeof(mstr));
close(fd);
if (rc < 1) {
error("read(%s): %m", file_path);
return -1;
}
str_to_cpuset(mask, mstr);
snprintf(file_path, sizeof(file_path), "%s/tasks", path);
fd = open(file_path, O_CREAT | O_RDONLY, 0700);
if (fd < 0) {
error("open(%s): %m", file_path);
return -1;
}
rc = read(fd, mstr, sizeof(mstr));
close(fd);
if (rc < 1) {
error("read(%s): %m", file_path);
return -1;
}
/* FIXME: verify that pid is in mstr */
return 0;
}
/* mhm - we need something clever for all that CPU_SET() and CPU_ISSET() stuff */
int parse_affinity(cpu_set_t *mask, char *arg)
{
cpu_set_t tmp_aff;
char *tmp_arg;
size_t valid_len;
CPU_ZERO(&tmp_aff);
if(*arg == '0' && *(arg+1) == 'x') {
/* we're in standard hex mode */
str_to_cpuset(&tmp_aff, arg);
} else if( (valid_len=strspn(arg, "0123456789,.")) ) {
/* new list mode: schedtool -a 0,2 -> run on CPU0 and CPU2 */
/* split on ',' and '.', because '.' is near ',' :) */
while((tmp_arg=strsep(&arg, ",."))) {
int tmp_cpu;
if(isdigit((int)*tmp_arg)) {
tmp_cpu=atoi(tmp_arg);
CPU_SET(tmp_cpu, &tmp_aff);
#ifdef DEBUG
printf("tmp_arg: %s -> tmp_cpu: %d\n", tmp_arg, tmp_cpu);
#endif
}
}
} else {
decode_error("affinity %s is not parseable", arg);
exit(1);
}
*mask=tmp_aff;
return 0;
}
int get_cpuset(cpu_set_t *mask, slurmd_job_t *job)
{
int nummasks, maskid, i, threads;
char *curstr, *selstr;
char mstr[1 + CPU_SETSIZE / 4];
uint32_t local_id = job->envtp->localid;
char buftype[1024];
slurm_sprint_cpu_bind_type(buftype, job->cpu_bind_type);
debug3("get_cpuset (%s[%d]) %s", buftype, job->cpu_bind_type,
job->cpu_bind);
CPU_ZERO(mask);
if (job->cpu_bind_type & CPU_BIND_NONE) {
return true;
}
if (job->cpu_bind_type & CPU_BIND_RANK) {
threads = MAX(conf->threads, 1);
CPU_SET(job->envtp->localid % (job->cpus*threads), mask);
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDRANK) {
/* if HAVE_NUMA then bind this task ID to it's corresponding
* locality domain ID. Otherwise, bind this task ID to it's
* corresponding socket ID */
return _bind_ldom(local_id, mask);
}
if (!job->cpu_bind)
return false;
nummasks = 1;
maskid = 0;
selstr = NULL;
/* get number of strings present in cpu_bind */
curstr = job->cpu_bind;
while (*curstr) {
if (nummasks == local_id+1) {
selstr = curstr;
maskid = local_id;
break;
}
if (*curstr == ',')
nummasks++;
curstr++;
}
/* if we didn't already find the mask... */
if (!selstr) {
/* ...select mask string by wrapping task ID into list */
maskid = local_id % nummasks;
i = maskid;
curstr = job->cpu_bind;
while (*curstr && i) {
if (*curstr == ',')
i--;
curstr++;
}
if (!*curstr) {
return false;
}
selstr = curstr;
}
/* extract the selected mask from the list */
i = 0;
curstr = mstr;
while (*selstr && *selstr != ',' && i++ < (CPU_SETSIZE/4))
*curstr++ = *selstr++;
*curstr = '\0';
if (job->cpu_bind_type & CPU_BIND_MASK) {
/* convert mask string into cpu_set_t mask */
if (str_to_cpuset(mask, mstr) < 0) {
error("str_to_cpuset %s", mstr);
return false;
}
return true;
}
if (job->cpu_bind_type & CPU_BIND_MAP) {
unsigned int mycpu = 0;
if (strncmp(mstr, "0x", 2) == 0) {
mycpu = strtoul (&(mstr[2]), NULL, 16);
} else {
mycpu = strtoul (mstr, NULL, 10);
}
CPU_SET(mycpu, mask);
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDMASK) {
/* if HAVE_NUMA bind this task to the locality domains
* identified in mstr. Otherwise bind this task to the
* sockets identified in mstr */
int len = strlen(mstr);
char *ptr = mstr + len - 1;
uint32_t base = 0;
curstr = mstr;
/* skip 0x, it's all hex anyway */
if (len > 1 && !memcmp(mstr, "0x", 2L))
curstr += 2;
while (ptr >= curstr) {
char val = char_to_val(*ptr);
if (val == (char) -1)
return false;
if (val & 1)
_bind_ldom(base, mask);
if (val & 2)
_bind_ldom(base + 1, mask);
if (val & 4)
_bind_ldom(base + 2, mask);
if (val & 8)
_bind_ldom(base + 3, mask);
len--;
ptr--;
base += 4;
}
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDMAP) {
/* if HAVE_NUMA bind this task to the given locality
* domain. Otherwise bind this task to the given
* socket */
uint32_t myldom = 0;
if (strncmp(mstr, "0x", 2) == 0) {
myldom = strtoul (&(mstr[2]), NULL, 16);
} else {
myldom = strtoul (mstr, NULL, 10);
}
return _bind_ldom(myldom, mask);
}
return false;
}
int main(int argc, char *argv[])
{
cpu_set_t new_mask, cur_mask;
pid_t pid = 0;
int opt, err;
char mstr[1 + CPU_SETSIZE / 4];
char cstr[7 * CPU_SETSIZE];
int c_opt = 0;
struct option longopts[] = {
{ "pid", 0, NULL, 'p' },
{ "cpu-list", 0, NULL, 'c' },
{ "help", 0, NULL, 'h' },
{ "version", 0, NULL, 'V' },
{ NULL, 0, NULL, 0 }
};
while ((opt = getopt_long(argc, argv, "+pchV", longopts, NULL)) != -1) {
int ret = 1;
switch (opt) {
case 'p':
pid = atoi(argv[argc - 1]);
break;
case 'c':
c_opt = 1;
break;
case 'V':
printf("taskset version " VERSION "\n");
return 0;
case 'h':
ret = 0;
default:
show_usage(argv[0]);
return ret;
}
}
if ((!pid && argc - optind < 2)
|| (pid && (argc - optind < 1 || argc - optind > 2))) {
show_usage(argv[0]);
return 1;
}
if (pid) {
if (sched_getaffinity(pid, sizeof (cur_mask), &cur_mask) < 0) {
perror("sched_getaffinity");
fprintf(stderr, "failed to get pid %d's affinity\n",
pid);
return 1;
}
if (c_opt)
printf("pid %d's current affinity list: %s\n", pid,
cpuset_to_cstr(&cur_mask, cstr));
else
printf("pid %d's current affinity mask: %s\n", pid,
cpuset_to_str(&cur_mask, mstr));
if (argc - optind == 1)
return 0;
}
if (c_opt)
err = cstr_to_cpuset(&new_mask, argv[optind]);
else
err = str_to_cpuset(&new_mask, argv[optind]);
if (err) {
if (c_opt)
fprintf(stderr, "failed to parse CPU list %s\n",
argv[optind]);
else
fprintf(stderr, "failed to parse CPU mask %s\n",
argv[optind]);
return 1;
}
if (sched_setaffinity(pid, sizeof (new_mask), &new_mask)) {
perror("sched_setaffinity");
fprintf(stderr, "failed to set pid %d's affinity.\n", pid);
return 1;
}
if (sched_getaffinity(pid, sizeof (cur_mask), &cur_mask) < 0) {
perror("sched_getaffinity");
fprintf(stderr, "failed to get pid %d's affinity.\n", pid);
return 1;
}
if (pid) {
if (c_opt)
printf("pid %d's new affinity list: %s\n", pid,
cpuset_to_cstr(&cur_mask, cstr));
else
printf("pid %d's new affinity mask: %s\n", pid,
cpuset_to_str(&cur_mask, mstr));
} else {
argv += optind + 1;
execvp(argv[0], argv);
perror("execvp");
fprintf(stderr, "failed to execute %s\n", argv[0]);
return 1;
}
return 0;
}
int _get_sched_cpuset(hwloc_topology_t topology,
hwloc_obj_type_t hwtype, hwloc_obj_type_t req_hwtype,
cpu_set_t *mask, stepd_step_rec_t *job)
{
int nummasks, maskid, i, threads;
char *curstr, *selstr;
char mstr[1 + CPU_SETSIZE / 4];
uint32_t local_id = job->envtp->localid;
char buftype[1024];
/* For CPU_BIND_RANK, CPU_BIND_MASK and CPU_BIND_MAP, generate sched
* cpuset directly from cpu numbers.
* For CPU_BIND_LDRANK, CPU_BIND_LDMASK and CPU_BIND_LDMAP, generate
* sched cpuset from hwloc topology.
*/
slurm_sprint_cpu_bind_type(buftype, job->cpu_bind_type);
debug3("task/cgroup: (%s[%d]) %s", buftype,
job->cpu_bind_type, job->cpu_bind);
CPU_ZERO(mask);
if (job->cpu_bind_type & CPU_BIND_NONE) {
return true;
}
if (job->cpu_bind_type & CPU_BIND_RANK) {
threads = MAX(conf->threads, 1);
CPU_SET(job->envtp->localid % (job->cpus*threads), mask);
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDRANK) {
return _get_ldom_sched_cpuset(topology, hwtype, req_hwtype,
local_id, mask);
}
if (!job->cpu_bind)
return false;
nummasks = 1;
selstr = NULL;
/* get number of strings present in cpu_bind */
curstr = job->cpu_bind;
while (*curstr) {
if (nummasks == local_id+1) {
selstr = curstr;
break;
}
if (*curstr == ',')
nummasks++;
curstr++;
}
/* if we didn't already find the mask... */
if (!selstr) {
/* ...select mask string by wrapping task ID into list */
maskid = local_id % nummasks;
i = maskid;
curstr = job->cpu_bind;
while (*curstr && i) {
if (*curstr == ',')
i--;
curstr++;
}
if (!*curstr) {
return false;
}
selstr = curstr;
}
/* extract the selected mask from the list */
i = 0;
curstr = mstr;
while (*selstr && *selstr != ',' && i++ < (CPU_SETSIZE/4))
*curstr++ = *selstr++;
*curstr = '\0';
if (job->cpu_bind_type & CPU_BIND_MASK) {
/* convert mask string into cpu_set_t mask */
if (str_to_cpuset(mask, mstr) < 0) {
error("task/cgroup: str_to_cpuset %s", mstr);
return false;
}
return true;
}
if (job->cpu_bind_type & CPU_BIND_MAP) {
unsigned int mycpu = 0;
if (strncmp(mstr, "0x", 2) == 0) {
mycpu = strtoul (&(mstr[2]), NULL, 16);
} else {
mycpu = strtoul (mstr, NULL, 10);
}
CPU_SET(mycpu, mask);
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDMASK) {
int len = strlen(mstr);
char *ptr = mstr + len - 1;
uint32_t base = 0;
curstr = mstr;
/* skip 0x, it's all hex anyway */
if (len > 1 && !memcmp(mstr, "0x", 2L))
curstr += 2;
while (ptr >= curstr) {
char val = char_to_val(*ptr);
if (val == (char) -1)
return false;
if (val & 1)
_get_ldom_sched_cpuset(topology, hwtype,
req_hwtype, base, mask);
if (val & 2)
_get_ldom_sched_cpuset(topology, hwtype,
req_hwtype, base + 1, mask);
if (val & 4)
_get_ldom_sched_cpuset(topology, hwtype,
req_hwtype, base + 2, mask);
if (val & 8)
_get_ldom_sched_cpuset(topology, hwtype,
req_hwtype, base + 3, mask);
len--;
ptr--;
base += 4;
}
return true;
}
if (job->cpu_bind_type & CPU_BIND_LDMAP) {
uint32_t myldom = 0;
if (strncmp(mstr, "0x", 2) == 0) {
myldom = strtoul (&(mstr[2]), NULL, 16);
} else {
myldom = strtoul (mstr, NULL, 10);
}
return _get_ldom_sched_cpuset(topology, hwtype, req_hwtype,
myldom, mask);
}
return false;
}
