From 89177fce34edb5ad0059a41548888d05588cc1c5 Mon Sep 17 00:00:00 2001 From: Joshua Bakita Date: Mon, 16 Jun 2025 19:29:07 -0400 Subject: Rewrite nvtaskset and implementation of partitioning for unmodified tasks Rather than requiring libsmctrl.so to be preloaded, we now wrap libcuda.so.1. All CUDA-using applications will load libcuda.so.1, ensuring that our wrapper will always be dynamically loaded, no matter if LD_PRELOAD is enabled, or if a program has been staticly linked. All that needs to be done is that the location of our "fake" libcuda.so.1 need to be put within the loader search path. This can be done by setting LD_LIBRARY_PATH, or by installing our wrapper into /lib/x86_64-linux-gnu. The mask can still be set via the LIBSMCTRL_MASK environment variable, but the easier-to-use nvtaskset tool is now the recommended way to view or change the supreme TPC mask for any CUDA-using application. This allows launching a program on the first two GPCs via a command as simple as: ./nvtaskset -g 0-1 ./a_program a_program_args (Note that use of the -g option requires the nvdebug kernel module to first be loaded.) These changes support the final version of the ECRTS'25 paper. Note that nvtaskset does not yet fully support multi-GPU systems. Bugfixes: - Fix crash that would occur if both libsmctrl.so and libsmctrl.a were built into an application. - Correctly use GPU ID when initializing a context in `libsmctrl_test_gpc_info`. - Include `nvtaskset` as a prerequisite for `libsmctrl_test_supreme_mask`. - Fix malfunction of `libsmctrl_test_gpc_info` if CUDA_VISIBLE_DEVICES is set. Other minor changes: - Adds make target to run all the tests. - Fixes typos in comments. - Enables -Wall build option. - Upgrades supreme mask from 64 to 128 bits. - Removes `detect_parker_soc()` from the global namespace. - Adjusts test messages to be more succinct. - Updates README with overview of how to partition unmodified applications, more details on the tests, and information on the new ECRTS'25 paper. --- nvtaskset.c | 520 ++++++++++++++++++++++++++++++++++++++++-------------------- 1 file changed, 347 insertions(+), 173 deletions(-) (limited to 'nvtaskset.c') diff --git a/nvtaskset.c b/nvtaskset.c index 74f88d7..4901cbe 100644 --- a/nvtaskset.c +++ b/nvtaskset.c @@ -1,210 +1,384 @@ // Copyright 2025 Joshua Bakita -// taskset-like utility for the GPU +// Show or change the GPU core affinity for a CUDA process +// taskset-like utility for NVIDIA GPUs #define _GNU_SOURCE // For program_invocation_name #include +#include #include #include +#include +#include #include #include #include #include -#include -#include -#include +#include #include #include // To help with getting GPC info #include "libsmctrl.h" -const char* maintainer = ""; -const char* version = "nvtaskset 2025.03"; -const char* desc = "taskset-like utility for NVIDIA GPUs."; +#define LINK_NAME "/memfd:libsmctrl" + +// TODO: Write automated tests: +// - Change region of non-existent PID +// - Change region of permission denied PID +// - Change region of non-GPU PID +// - Change GPC list +// - Change TPC list +// - Change TPC mask +// - Start TPC mask +// - Start TPC list +// - Start GPC list +// - Start with subargument containing - +// - Query GPC list +// - Query TPC list +// - Query TPC mask +// - Set GPC list w/ non-existant GPC +// - Set TPC list w/ non-existant TPC + +// Private symbols from libsmctrl +extern bool libsmctrl_is_mps_running(); +extern uint128_t strtou128(const char *nptr, char **endptr, int base); + +const char *argp_program_bug_address = ""; +const char *argp_program_version = "nvtaskset 2025.06"; +const char desc[] = "Show or change the GPU core affinity for a CUDA process\v" + "Warning: When using GPC lists, this tool currently " + "derives TPC to GPC mappings from the first NVIDIA GPU in " + "the system (by PCI bus ID) device. To use the mappings " + "for a different device, use the `libsmctrl_test_get_info` " + "tool to get the bitmask of TPCs associated with each GPC, " + "OR them, and then set that bitmask via this tool. Better " + "multi-GPU support is intended for a future release.\n\n" + "Inspired by the Linux taskset utility."; +const char args_doc[] = "[mask | list] [pid | cmd [args...]]"; const struct argp_option opts[] = { + {"gpc-list", 'g', NULL, 0, "Specify partition as a list of GPCs"}, + {"tpc-list", 't', NULL, 0, "Specify partition as a list of TPCs"}, + {"pid", 'p', NULL, 0, "Operate on an existing PID"}, {0} }; -unsigned __int128 strtou128(const char *nptr, char **endptr, int base) { - unsigned __int128 result = 0; - if (base != 16) - error(1, EINVAL, "Internal error"); - // Skip a "0x" prefix. Safe due to early evaluation - if (*nptr == '0' && (*(nptr + 1) == 'x' || *(nptr + 1) == 'X')) - nptr += 2; - // Until hitting an invalid character - while (1) { - if (*nptr >= 'a' && *nptr <= 'f') - result = result << 4 | (*nptr - 'a' + 10); - else if (*nptr >= 'A' && *nptr <= 'F') - result = result << 4 | (*nptr - 'A' + 10); - else if (*nptr >= '0' && *nptr <= '9') - result = result << 4 | (*nptr - '0'); - else - break; - nptr++; - } - if (endptr) - *endptr = (char*)nptr; - return result; -} - +// Create a CUDA context and query the associated GPC to TPC mappings +// Based off logic in libsmctrl_test_gpc_info void libsmctrl_get_gpc_info_ext_easy(uint32_t* num_gpcs, uint128_t** masks, int gpu_id) { int res; CUcontext ctx; - // XXX: Copied from libsmctrl_test_gpc_info - // Tell CUDA to use PCI device id ordering (to match nvdebug) - putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID"); - // A CUDA context is required before reading the topology information - if ((res = cuInit(0))) { - const char* name; - cuGetErrorName(res, &name); - fprintf(stderr, "%s: Unable to initialize CUDA, error %s\n", program_invocation_name, name); - exit(1); - } - if ((res = cuCtxCreate(&ctx, 0, 0))) { - const char* name; - cuGetErrorName(res, &name); - fprintf(stderr, "%s: Unable to create a CUDA context, error %s\n", program_invocation_name, name); - exit(1); - } - // Pull topology information from libsmctrl - if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) { - error(0, res, "libsmctrl_get_gpc_info() failed"); - if (res == ENOENT) - fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name); - if (res == EIO) - fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name); - exit(1); - } - // Not copied + char *old_order = NULL; + // Tell CUDA to use PCI device id ordering (to match nvdebug) + putenv((char*)"CUDA_DEVICE_ORDER=PCI_BUS_ID"); + // Allow CUDA to see all devices (to better match nvdebug) + if (getenv("CUDA_VISIBLE_DEVICES")) { + if (!(old_order = strdup(getenv("CUDA_VISIBLE_DEVICES")))) + error(1, errno, "Unable to allocate environment string"); + unsetenv("CUDA_VISIBLE_DEVICES"); + } + // A CUDA context is required before reading the topology information + if ((res = cuInit(0))) { + const char* name; + cuGetErrorName(res, &name); + error(1, 0, "Unable to create a initialize CUDA, error %s\n", name); + } + if ((res = cuCtxCreate(&ctx, 0, gpu_id))) { + const char* name; + cuGetErrorName(res, &name); + error(1, 0, "Unable to create a CUDA context, error %s\n", name); + } + // Pull topology information from libsmctrl + if ((res = libsmctrl_get_gpc_info_ext(num_gpcs, masks, gpu_id)) != 0) { + error(0, res, "libsmctrl_get_gpc_info() failed"); + if (res == ENOENT) + fprintf(stderr, "%s: Is the nvdebug kernel module loaded?\n", program_invocation_name); + if (res == EIO) + fprintf(stderr, "%s: Is the GPU powered on, i.e., is there an active context?\n", program_invocation_name); + exit(1); + } + // Restore the environment (in case we exec() later) unsetenv("CUDA_DEVICE_ORDER"); + if (old_order) { + setenv("CUDA_VISIBLE_DEVICES", old_order, 1); + free(old_order); + } } -int main(int argc, char **argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s -p \n", argv[0]); - fprintf(stderr, " %s \n", argv[0]); - fprintf(stderr, " %s --gpc-list \n", argv[0]); - fprintf(stderr, " has a bit set for each TPC to be enabled\n"); - return 1; - } - // TODO: Use a proper argument parser - if (strcmp("-p", argv[1]) == 0) { // Setting mask on running task - char *end; - pid_t target_pid = strtoul(argv[2], &end, 10); - // strtoul stores a pointer to the first invalid character in `end` - if (*end != '\0') { - fprintf(stderr, "Invalid character \"%c\" in PID argument.\n", *end); - return 1; +int parse_list(bool use_gpcs, char* list, uint128_t *mask_out) { + // We support the same ranges as taskset, e.g., X,Y,Z and X,Y-Z + uint32_t num_xpcs = 0; // Either TPC or GPC count, i.e., "X"PC + uint128_t* masks = NULL; + // TODO: Allow specifying GPU ID, rather than assuming 0! + if (use_gpcs) + libsmctrl_get_gpc_info_ext_easy(&num_xpcs, &masks, 0); + else + libsmctrl_get_tpc_info_cuda(&num_xpcs, 0); + uint128_t mask = 0; + int range_start_xpc = -1; + char* start = list; + int len = strlen(list); + // Convert comma-seperated GPC/TPC list into a mask + for (int i = 0; i < len + 1; i++) { + if (list[i] == ',' || list[i] == '\0') { + list[i] = '\0'; + int xpc = atoi(start); + if (xpc > num_xpcs - 1) + error(1, EINVAL, "%s is not a valid %s ID", start, use_gpcs ? "GPC" : "TPC"); + // Handle ranges + if (range_start_xpc != -1) { + if (range_start_xpc >= xpc) + error(1, EINVAL, "Malformed %s range", use_gpcs ? "GPC" : "TPC"); + while (range_start_xpc <= xpc) { + if (use_gpcs) + mask |= masks[range_start_xpc]; + else + mask |= (uint128_t)1 << range_start_xpc; + range_start_xpc++; + } + range_start_xpc = -1; + } else { + if (use_gpcs) + mask |= masks[xpc]; + else + mask |= (uint128_t)1 << xpc; + } + start = list + i + 1; } - unsigned __int128 mask = strtou128(argv[3], &end, 16); - if (*end != '\0') { - fprintf(stderr, "Invalid character \"%c\" in mask argument.\n", *end); - return 1; + // Range start + if (list[i] == '-') { + list[i] = '\0'; + range_start_xpc = atoi(start); + start = list + i + 1; } - // The shared memory lookup key is the lower 16-bits of the PID | "sm" - key_t shm_key = target_pid << 16 | (int)'s' << 8 | (int) 'm'; - // Get a handle to the 128-bit shared memory region - int shmid = shmget(shm_key, 16, 0); - if (shmid == -1) - error(1, errno, "Unable to find control region for PID %d", target_pid); - // Open the shared memory region - unsigned __int128 *supreme_mask = shmat(shmid, NULL, 0); - if (supreme_mask == (void*)-1) - error(1, errno, "Unable to open control region for PID %d", target_pid); - // Write the requested mask into the shared memory region - *supreme_mask = mask; - } else { // Starting a new task with a mask - // TODO: Check other locations for nvidia-cuda-mps-control if its not on the path - // TODO: Use dup2() to redirect MPS startup messages - int ret = system("echo -n | nvidia-cuda-mps-control"); - if (ret == -1) - error(1, errno, "Unable to run subshell to check MPS status"); - if (ret != 0) { // Control deamon not yet started - fprintf(stderr, "nvtaskset: MPS control deamon does not appear to be running. Automatically starting...\n"); - ret = system("nvidia-cuda-mps-control -d"); - if (ret == -1) - error(1, errno, "Unable to run subshell to start MPS"); - if (ret == 1) { - fprintf(stderr, "nvtaskset: Error starting MPS control deamon. Terminating...\n"); - return 1; + } + *mask_out = mask; + return 0; +} + +// Always returns a valid string +char* compose_list(uint128_t mask) { + // List will always be shorter than every TPC, comma-seperated + // 128 TPCs, with 10 1-char, 90 2-char, 28 3-char, 127 commas, and 1 null + static char list[10 + 90*2 + 28*3 + 128]; + char* tail = list; + int last_enabled = -2; + bool in_range; + for (int i = 0; i < 128; i++) { + bool enabled = (mask >> i) & 1; + if (in_range) { + if (enabled) { + last_enabled = i; + } else { + tail += sprintf(tail, "%d,", last_enabled); + in_range = false; } - fprintf(stderr, "nvtaskset: Done. Use \"echo quit | nvidia-cuda-mps-control\" to terminate it later as desired.\n"); + continue; + } + if (enabled) { + if (last_enabled == i - 1) { + in_range = true; + tail += sprintf(tail, "-"); + } else { + tail += sprintf(tail, "%d", i); + } + last_enabled = i; + } else { + if (last_enabled == i - 1) { + tail += sprintf(tail, ","); + } + } + } + // Strip trailing comma + if (*(tail - 1) == ',') + *(tail - 1) = '\0'; + return list; +} + +// Always returns a valid string +// (Terminates the program on error) +char* compose_gpc_list(uint128_t mask) { + uint32_t num_gpcs = 0; + uint128_t* masks = NULL; + libsmctrl_get_gpc_info_ext_easy(&num_gpcs, &masks, 0); + uint128_t gpc_mask = 0; + // Try to find correspondence between a list of TPCs and GPCs + for (int gpc = 0; gpc < num_gpcs; gpc++) { + if ((masks[gpc] & mask) == masks[gpc]) { + gpc_mask |= 1 << gpc; + mask &= ~masks[gpc]; } - // Tell loader to initialize libsmctrl.so first - // TODO: Append, rather than overwrite LD_PRELOAD - setenv("LD_PRELOAD", "libsmctrl.so", 1); - // Explictly set the number of channels, otherwise CUDA will only use two - // (see paper for why that causes problems) - setenv("CUDA_DEVICE_MAX_CONNECTIONS", "8", 1); - // Check if a mask, or a list of GPCs is being provided - if (strcmp(argv[1], "--gpc-list") == 0) { - // TODO: Support the full syntax that taskset supports - // We just support X,Y,Z for now - uint32_t num_gpcs = 0; - uint128_t* masks = NULL; - // TODO: Allow specifying GPU ID, rather than assuming 0! - libsmctrl_get_gpc_info_ext_easy(&num_gpcs, &masks, 0); - uint128_t mask = 0; - int range_start_gpc = -1; - char* start = argv[2]; - int len = strlen(argv[2]); - // TODO: Handle invalid input cleanly. - // Convert comma-seperated GPC list into a mask - for (int i = 0; i < len + 1; i++) { - if (argv[2][i] == ',' || argv[2][i] == '\0') { - argv[2][i] = '\0'; - int gpc = atoi(start); - if (gpc > num_gpcs - 1) { - fprintf(stderr, "Invalid GPC ID '%s'!\n", start); + } + if (mask) + error(1, EINVAL, "Unable to interpret affinity as GPC list; try -t instead of -g"); + return compose_list(gpc_mask); +} + + +uint128_t* get_mask_hndl(pid_t target_pid) { + char fd_path[277]; + int fd; + uint128_t *mask_hndl; + DIR *dp; + struct dirent *entry; + // Search for the file descriptor which represents the libsmctrl control + // region. + snprintf(fd_path, 277, "/proc/%d/fd/", target_pid); + if (!(dp = opendir(fd_path))) { + if (errno == ENOENT) + error(1, 0, "Unable to find PID %d.", target_pid); + else + error(1, errno, "Unable to access PID %d", target_pid); + } + while (entry = readdir(dp)) { + char link[sizeof(LINK_NAME)]; + snprintf(fd_path, 277, "/proc/%d/fd/%s", target_pid, entry->d_name); + readlink(fd_path, link, sizeof(LINK_NAME)); + if (strncmp(LINK_NAME, link, sizeof(LINK_NAME) - 1) == 0) + break; + } + closedir(dp); + if (!entry) + error(1, 0, "Unable to find libsmctrl-wrapper control region for PID %d.", target_pid); + // Access the shared memory region for libsmctrl control. + if ((fd = open(fd_path, O_RDWR)) == -1) + error(1, errno, "Unable to open libsmctrl-wrapper control file %s", fd_path); + mask_hndl = mmap(NULL, 16, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + if (mask_hndl == MAP_FAILED) + error(1, errno, "Unable to memory-map libsmctrl-wrapper control file %s", fd_path); + close(fd); + return mask_hndl; +} + +static error_t arg_parser(int key, char* arg, struct argp_state *state){ + static bool is_cmd = true; + static bool is_query = false; + static bool is_list = false; + static bool use_gpcs = false; + static uint128_t mask = 0; + static pid_t target_pid = 0; + static char **sub_argv = NULL; + char *end; + // Handle what to do in case of each option + switch (key) { + case 'g': + if (is_list) + argp_error(state, "Only one of -g and -t may be specified.\n"); + use_gpcs = true; + is_list = true; + break; + case 't': + if (is_list) + argp_error(state, "Only one of -g and -t may be specified.\n"); + is_list = true; + break; + case 'p': + is_cmd = false; + break; + case ARGP_KEY_ARG: + // Options: + // 1. -p and one argument -> Query mask for PID + // 2. -p and two arguments -> Set mask for PID + // 3. No -p and at least one argument -> Set mask and launch command + // (otherwise: invalid) + if (state->arg_num == 0 && !is_cmd && state->argc - state->next == 0) + is_query = true; + // Handle invalid and valid query cases + if (is_query) { + if (state->arg_num == 0) { + target_pid = strtoul(arg, &end, 10); + if (*end != '\0') + argp_error(state, "Invalid character \"%c\" in PID argument.\n", *end); + break; + } else + return ARGP_ERR_UNKNOWN; + } + // Handle non-query cases + if (state->arg_num == 0 && state->argc - state->next != 0) { + if (is_list) { + parse_list(use_gpcs, arg, &mask); + } else { + // strtoul stores a pointer to the first invalid character in `end` + mask = strtou128(arg, &end, 16); + if (*end != '\0') + argp_error(state, "Invalid character \"%c\" in mask argument.\n", *end); + } + } else if (state->arg_num == 1 && !is_cmd) { + target_pid = strtoul(arg, &end, 10); + if (*end != '\0') + argp_error(state, "Invalid character \"%c\" in PID argument.\n", *end); + } else + return ARGP_ERR_UNKNOWN; + break; + case ARGP_KEY_ARGS: + if (!is_cmd) + return ARGP_ERR_UNKNOWN; + sub_argv = state->argv + state->next; + break; + case ARGP_KEY_END: + if (is_query && state->arg_num < 1) + argp_usage(state); + else if (!is_query && state->arg_num < 2) + argp_usage(state); + break; + case ARGP_KEY_FINI: + if (is_query) { + // query PID + uint128_t* mask_hndl = get_mask_hndl(target_pid); + uint128_t enable_mask = ~*mask_hndl; + if (use_gpcs & is_list) + printf("PID %d's current GPC affinity list: %s\n", target_pid, compose_gpc_list(enable_mask)); + else if (use_gpcs & !is_list) + argp_error(state, "Unsupported to print query as a GPC mask.\n"); + else if (is_list) + printf("PID %d's current TPC affinity list: %s\n", target_pid, compose_list(enable_mask)); + else + printf("PID %d's current TPC affinity mask: 0x%.0lx%016lx\n", target_pid, (uint64_t)(enable_mask >> 64), (uint64_t)enable_mask); + } else if (is_cmd) { + // start MPS (as needed) + if (!libsmctrl_is_mps_running()) { + fprintf(stderr, "nvtaskset: MPS control deamon does not appear to be running. Automatically starting...\n"); + int ret = system("nvidia-cuda-mps-control -d"); + if (ret == -1) + error(1, errno, "Unable to run subshell to start MPS"); + if (ret == 1) { + fprintf(stderr, "nvtaskset: Error starting MPS control deamon. Terminating...\n"); return 1; } - // Handle ranges - if (range_start_gpc != -1) { - if (range_start_gpc >= gpc) { - fprintf(stderr, "Invalid GPC range!\n"); - return 1; - } - while (range_start_gpc <= gpc) { - //printf("gpc %i\n", range_start_gpc); - mask |= masks[range_start_gpc]; - range_start_gpc++; - } - range_start_gpc = -1; - } else { - //printf("gpc %i\n", gpc); - mask |= masks[gpc]; - } - start = argv[2] + i + 1; + fprintf(stderr, "nvtaskset: Done. Use \"echo quit | nvidia-cuda-mps-control\" to terminate it later as desired.\n"); } - // Range start - if (argv[2][i] == '-') { - argv[2][i] = '\0'; - range_start_gpc = atoi(start); - start = argv[2] + i + 1; + // launch subprocess + // Convert to string, prefix with ~, and set env var + char mask_str[32+3+1]; // 32 hexits, "~0x", and '\0' + snprintf(mask_str, 36, "~0x%.0lx%016lx", (uint64_t)(mask >> 64), (uint64_t)mask); + setenv("LIBSMCTRL_MASK", mask_str, 1); + // Start task + execvp(sub_argv[0], sub_argv); + error(1, errno, "Unable to launch task '%s'", sub_argv[0]); + } else { + if (!libsmctrl_is_mps_running()) + printf("Warning: NVIDIA MPS is not running. CUDA programs will not co-run! Run nvidia-cuda-mps-control -d before launching any CUDA-using programs that should co-run.\n"); + // change mask on PID + uint128_t* mask_hndl = get_mask_hndl(target_pid); + if (!is_list) { + printf("PID %d's current TPC affinity mask: 0x%.0lx%016lx\n", target_pid, ~(uint64_t)(*mask_hndl >> 64), ~(uint64_t)*mask_hndl); + printf("PID %d's new TPC affinity mask: 0x%.0lx%016lx\n", target_pid, (uint64_t)(mask >> 64), (uint64_t)mask); + } else { + printf("PID %d's current TPC affinity list: %s\n", target_pid, compose_list(~*mask_hndl)); + printf("PID %d's new TPC affinity list: %s\n", target_pid, compose_list(mask)); } + // Write the requested mask into the shared memory region + *mask_hndl = ~mask; } - // Convert to string, prefix with ~, and set env var - char mask_str[32+3+1]; // 32 hexits, "~0x", and '\0' - snprintf(mask_str, 36, "~0x%lx%016lx", (uint64_t)(mask >> 64), (uint64_t)mask); - //printf("nvtaskset: Using mask string %s\n", mask_str); - setenv("LIBSMCTRL_MASK", mask_str, 1); - // Start task - execvp(argv[3], argv+3); - error(1, errno, "Unable to launch task '%s'", argv[3]); - } else { - // Tell libsmctrl what mask to use - char* mask = malloc(strlen(argv[1]) + 2); - mask[0] = '~'; // Make an enable mask - strcpy(mask+1, argv[1]); - setenv("LIBSMCTRL_MASK", mask, 1); - free(mask); // setenv() made a copy - // Start task - execvp(argv[2], argv+2); - error(1, errno, "Unable to launch task '%s'", argv[2]); - } + break; + default: + return ARGP_ERR_UNKNOWN; } - fprintf(stderr, "Invalid arguments\n"); - return 1; + return 0; +} + +struct argp argp = {opts, arg_parser, args_doc, desc}; + +int main(int argc, char **argv) { + argp_parse(&argp, argc, argv, ARGP_IN_ORDER, 0, NULL); + return 0; } -- cgit v1.2.2