diff --git a/agent/cmd/cuinterpose/README.md b/agent/cmd/cuinterpose/README.md index db56db9f..e6236295 100644 --- a/agent/cmd/cuinterpose/README.md +++ b/agent/cmd/cuinterpose/README.md @@ -4,10 +4,10 @@ All rights reserved. SPDX-License-Identifier: Apache-2.0 --> -# Same-node POSIX CUDA VMM checkpoint and restore +# Same-node POSIX CUDA VMM and multicast checkpoint and restore This interposer implements checkpoint and restore for CUDA VMM allocations -shared between processes on one node with +and complete CUDA multicast groups shared between processes on one node with `CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR`. ## Contract @@ -51,17 +51,16 @@ Legacy CUDA IPC remains driver-owned and is unsupported by this shim. hex participant ID. Otherwise, the shim creates one when the process starts. The orchestrator must externally quiesce the application at the checkpoint -boundary. Allocation sharing, import, mapping, access updates, kernel work, -communication-library setup, and teardown must not be in flight. +boundary. Allocation sharing, mapping, access updates, multicast setup or +teardown, kernel work, and communication-library setup must not be in flight. During ordinary execution: -- CUDA generic allocation handles for POSIX-capable allocations tracked by the - shim are tagged logical tokens; the corresponding real driver handles remain - internal. -- Untracked CUDA generic allocation handles remain real driver handles. -- A POSIX-capable allocation becomes checkpoint-managed when its ticket fd - is exported. +- CUDA generic handles for tracked POSIX resources are tagged logical tokens; + the corresponding real driver handles remain internal. +- Untracked CUDA generic handles remain real driver handles. +- A POSIX-capable unicast allocation becomes checkpoint-managed when its ticket + FD is exported. - POSIX exports return sealed ticket FDs containing the creator, allocation, endpoint, and authorization identities. - A ticket import resolves through the creator's local Unix endpoint. A @@ -69,18 +68,22 @@ During ordinary execution: closed immediately, and never returned to the application. - A raw external POSIX import is passed directly to CUDA and is not tracked by the shim. +- Multicast create/import/export, device membership, bindings, mappings, and + access are passed to CUDA and recorded. Handle ownership is explicit: | Source | Application receives | Managed table | Driver-handle owner | | --- | --- | --- | --- | -| POSIX-capable create, tracked retain, or ticket import | Tagged logical token | Logical token to driver handle | Shim | +| POSIX-capable create, tracked retain, or ticket import | Tagged logical token | Logical token to unicast or multicast resource | Shim | | Raw external import or other pass-through | Real driver handle | None | Application | | Checkpoint carrier | Nothing | No separate entry | Shim | `posix.c` owns the sealed ticket format and remote creator exchange. `symbols.c` owns `dlsym`, `cuGetProcAddress`, and the replacement table. -`interpose.c` owns CUDA interception, allocation state, and local raw exports. +`interpose.c` owns generic CUDA interception and unicast allocation state. +`multicast.c` owns multicast devices, bindings, mappings, teardown, +reconstruction, and validation. Before native CUDA checkpoint, the snapshot agent asks the native coordinator to validate the complete participant topology. The original `cuMemCreate` @@ -99,6 +102,12 @@ handle translations, and performs final topology validation. Any failed validation fails closed. There is no rollback after checkpoint preparation mutates CUDA state. +Multicast is a topology overlay on the unicast allocations whose bytes native +CUDA checkpoints. Prepare removes multicast mappings, bindings, and objects +before unicast prepare. Restore creates fresh multicast objects, restores +device membership and bindings, remaps the original virtual addresses, and +validates the complete group before the workload resumes. + The native coordinator atomically writes `cuinterposer.state`, its opaque durable topology sidecar, in the checkpoint directory. Go only orders native VMM prepare and restore; it does not serialize or inspect the topology. The @@ -119,15 +128,24 @@ replacement table, using the requested CUDA version to select the ABI. Chaining with another `dlsym()` interposer and preserving original-caller `RTLD_NEXT` lookup scope are unsupported. +The multicast surface includes create, add-device, bind-memory, bind-address, +unbind, granularity, and the generic handle and mapping consumers. CUDA 13.1+ +resolver requests select the device-explicit bind ABI when build headers expose +it. + ## Testing -Run the native interposer integration test from the repository root: +Run the native interposer integration tests from the repository root: ```bash -uv run --project agent/cmd/cuinterpose/tests \ - pytest agent/cmd/cuinterpose/tests/test_cucheckpoint.py -v +cuda-checkpoint --launch-job \ + uv run --project agent/cmd/cuinterpose/tests \ + pytest agent/cmd/cuinterpose/tests/test_cucheckpoint.py -v ``` +The suite includes a POSIX baseline and a two-GPU multicast case with a +multimem collective and CUDA graph replay. + ## Qualification Disjoint GPU qualification is limited to source GPUs 0/1 restored onto @@ -169,11 +187,9 @@ Applications must finish setup and externally quiesce the complete process group first. Children created with `fork()` and no subsequent `exec()` lazily receive a fresh -random process-local participant identity, a child-PID control socket and thread, -and empty allocation, handle, and mapping bookkeeping on their first intercepted -VMM operation. This supports the fork-before-CUDA-initialization lifecycle used -by vLLM. An explicitly configured parent participant ID is never reused by a -forked child. +process-local participant identity, control socket, and empty unicast and +multicast bookkeeping on their first intercepted VMM operation. This supports +the fork-before-CUDA-initialization lifecycle used by vLLM. Forking after the shim has tracked VMM allocations, handles, or mappings is unsupported. The child deliberately discards its inherited copies without @@ -182,12 +198,16 @@ preserved. ### Fail-closed limitations -The coordinator fails closed for a missing creator anchor, incomplete +The coordinator fails closed for a missing unicast creator anchor, incomplete participants or topology, more than eight access descriptors, access ranges that partially overlap a tracked mapping, and reconstruction or final validation failures. One access call may cover multiple complete mappings. There is no rollback after checkpoint preparation mutates CUDA state. +Only complete same-node POSIX multicast groups are supported. Non-POSIX handle +types, partial groups, duplicate or missing devices, missing member bindings, +and multiple metadata creators fail closed. + ### Future compatibility FABRIC/IMEX handles, non-POSIX allocation sharing, raw-FD compatibility, and diff --git a/agent/cmd/cuinterpose/coordinator.c b/agent/cmd/cuinterpose/coordinator.c index 9c1d89e1..cb569110 100644 --- a/agent/cmd/cuinterpose/coordinator.c +++ b/agent/cmd/cuinterpose/coordinator.c @@ -31,6 +31,26 @@ struct participant { uint32_t count; }; +struct multicast { + uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]; + char creator[CUINTERPOSER_ID_SIZE]; + uint64_t size; + uint64_t handle_types; + uint64_t flags; + uint32_t num_devices; + uint32_t creators; + uint32_t devices; + uint32_t bindings; + struct multicast_device* device_list; + struct multicast* next; +}; + +struct multicast_device { + int32_t device; + bool bound; + struct multicast_device* next; +}; + struct allocation { uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]; char creator[CUINTERPOSER_ID_SIZE]; @@ -59,6 +79,45 @@ connect_endpoint(const char* endpoint) return fd; } +static struct multicast* +find_multicast(struct multicast* multicasts, const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]) +{ + struct multicast* multicast; + + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + if (memcmp(multicast->id, id, CUINTERPOSER_ALLOCATION_ID_SIZE) == 0) + return multicast; + } + return NULL; +} + +static struct multicast_device* +find_multicast_device(struct multicast* multicast, int32_t device) +{ + struct multicast_device* current; + + for (current = multicast->device_list; current != NULL; current = current->next) { + if (current->device == device) + return current; + } + return NULL; +} + +static void +free_multicasts(struct multicast* multicasts) +{ + while (multicasts != NULL) { + struct multicast* next = multicasts->next; + while (multicasts->device_list != NULL) { + struct multicast_device* device_next = multicasts->device_list->next; + free(multicasts->device_list); + multicasts->device_list = device_next; + } + free(multicasts); + multicasts = next; + } +} + static int exchange(struct participant* participant, uint16_t operation, struct cuinterposer_record** records, uint32_t* count) { @@ -127,6 +186,7 @@ static int validate_topology(struct participant* participants, size_t participant_count, struct allocation** output) { struct allocation* allocations = NULL; + struct multicast* multicasts = NULL; const char* reason = NULL; uint64_t value = UINT64_MAX; size_t participant_index; @@ -154,6 +214,7 @@ validate_topology(struct participant* participants, size_t participant_count, st for (record_index = 0; record_index < participant->count; record_index++) { const struct cuinterposer_record* record = &participant->records[record_index]; struct allocation* allocation = find_allocation(allocations, record->allocation_id); + struct multicast* multicast = find_multicast(multicasts, record->allocation_id); if (record->kind == CUINTERPOSER_ALLOCATION) { if (allocation == NULL) { @@ -210,6 +271,88 @@ validate_topology(struct participant* participants, size_t participant_count, st } if ((record->flags & CUINTERPOSER_CREATOR) != 0) allocation->creator_mapping = true; + } else if (record->kind == CUINTERPOSER_MULTICAST) { + if (record->handle_types != CUINTERPOSER_POSIX_HANDLE_TYPE || record->num_devices == 0 || + record->allocation_size == 0 || !is_lower_hex_id(record->creator_participant)) { + reason = "invalid multicast properties"; + goto failed; + } + if (multicast == NULL) { + multicast = calloc(1, sizeof(*multicast)); + if (multicast == NULL) { + reason = "multicast metadata allocation failed"; + goto failed; + } + memcpy(multicast->id, record->allocation_id, sizeof(multicast->id)); + multicast->size = record->allocation_size; + multicast->handle_types = record->handle_types; + multicast->flags = record->object_flags; + multicast->num_devices = record->num_devices; + snprintf(multicast->creator, sizeof(multicast->creator), "%s", record->creator_participant); + multicast->next = multicasts; + multicasts = multicast; + } else if ( + multicast->size != record->allocation_size || multicast->handle_types != record->handle_types || + multicast->flags != record->object_flags || multicast->num_devices != record->num_devices || + strcmp(multicast->creator, record->creator_participant) != 0) { + reason = "inconsistent multicast properties"; + goto failed; + } + if ((record->flags & CUINTERPOSER_CREATOR) != 0) { + if (strcmp(participant->id, multicast->creator) != 0) { + reason = "invalid multicast creator"; + goto failed; + } + multicast->creators++; + } + } else if (record->kind == CUINTERPOSER_MULTICAST_DEVICE) { + struct multicast_device* device; + if (multicast == NULL) { + reason = "multicast device precedes object"; + goto failed; + } + if (find_multicast_device(multicast, record->device) != NULL) { + reason = "duplicate multicast device"; + goto failed; + } + device = calloc(1, sizeof(*device)); + if (device == NULL) { + reason = "multicast device metadata allocation failed"; + goto failed; + } + device->device = record->device; + device->next = multicast->device_list; + multicast->device_list = device; + multicast->devices++; + } else if (record->kind == CUINTERPOSER_MULTICAST_BINDING) { + struct allocation* member = find_allocation(allocations, record->member_id); + struct multicast_device* device; + if (multicast == NULL || record->size == 0 || record->offset > multicast->size || + record->size > multicast->size - record->offset || + (record->binding_kind != CUINTERPOSER_MULTICAST_BIND_MEM && + record->binding_kind != CUINTERPOSER_MULTICAST_BIND_ADDR) || + (record->api_version != 1 && record->api_version != 2)) { + reason = "invalid multicast binding"; + goto failed; + } + if ((record->binding_kind == CUINTERPOSER_MULTICAST_BIND_MEM && (member == NULL || record->address != 0)) || + (record->binding_kind == CUINTERPOSER_MULTICAST_BIND_ADDR && record->address == 0)) { + reason = "invalid multicast member"; + goto failed; + } + device = find_multicast_device(multicast, record->device); + if (device == NULL) { + reason = "multicast binding device is absent"; + goto failed; + } + device->bound = true; + multicast->bindings++; + } else if (record->kind == CUINTERPOSER_MULTICAST_MAPPING) { + if (multicast == NULL || record->address == 0 || record->size == 0 || record->offset > multicast->size || + record->size > multicast->size - record->offset || record->access_count > CUINTERPOSER_MAX_ACCESS) { + reason = "invalid multicast mapping"; + goto failed; + } } else { reason = "unknown record kind"; value = record->kind; @@ -248,7 +391,38 @@ validate_topology(struct participant* participants, size_t participant_count, st } } } + { + struct multicast* multicast; + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + if (multicast->creators != 1) { + reason = "multicast group must have exactly one creator"; + value = multicast->creators; + goto failed; + } + if (multicast->devices != multicast->num_devices) { + reason = "incomplete multicast device group"; + value = multicast->devices; + goto failed; + } + if (multicast->bindings < multicast->num_devices) { + reason = "incomplete multicast binding group"; + value = multicast->bindings; + goto failed; + } + { + struct multicast_device* device; + for (device = multicast->device_list; device != NULL; device = device->next) { + if (!device->bound) { + reason = "multicast device has no binding"; + value = (uint64_t)(uint32_t)device->device; + goto failed; + } + } + } + } + } *output = allocations; + free_multicasts(multicasts); return 0; failed: if (value != UINT64_MAX) @@ -268,6 +442,7 @@ validate_topology(struct participant* participants, size_t participant_count, st free(allocations); allocations = next; } + free_multicasts(multicasts); return -1; } @@ -332,7 +507,7 @@ write_state(struct participant* participants, size_t count, FILE* output) size_t index; qsort(participants, count, sizeof(*participants), participant_compare); - if (fprintf(output, "snapshot-cuda-posix-v1\n") < 0) + if (fprintf(output, "snapshot-cuda-posix-v2\n") < 0) return -1; for (index = 0; index < count; index++) { struct participant* participant = &participants[index]; @@ -388,7 +563,7 @@ read_state(FILE* input, struct participant** output, size_t* output_count) struct participant* participants = NULL; size_t count = 0; - if (fgets(line, sizeof(line), input) == NULL || strcmp(line, "snapshot-cuda-posix-v1\n") != 0) + if (fgets(line, sizeof(line), input) == NULL || strcmp(line, "snapshot-cuda-posix-v2\n") != 0) return -1; while (fgets(line, sizeof(line), input) != NULL) { struct participant* participant; @@ -506,6 +681,21 @@ restore_unicast(struct participant* participants, size_t count) return command_all(participants, count, CUINTERPOSER_RESTORE_IMPORTERS); } +static int +restore_multicast(struct participant* participants, size_t count) +{ + /* Create, then import (same accept/export constraint as unicast). AddDevice on + * every rank, then bind: BindMem waits for the complete team, and the shim + * handles one request at a time. */ + if (command_all(participants, count, CUINTERPOSER_RESTORE_MULTICAST_CREATORS) != 0) + return -1; + if (command_all(participants, count, CUINTERPOSER_RESTORE_MULTICAST_IMPORTERS) != 0) + return -1; + if (command_all(participants, count, CUINTERPOSER_RESTORE_MULTICAST_DEVICES) != 0) + return -1; + return command_all(participants, count, CUINTERPOSER_RESTORE_MULTICAST); +} + static int same_participants(struct participant* expected, size_t expected_count, struct participant* actual, size_t actual_count) { @@ -583,7 +773,9 @@ main(int argc, char** argv) return errno == ENOENT ? EXIT_SUCCESS : EXIT_FAILURE; } if (proc_root[0] == '\0') { - const char* control = getenv("DYN_SNAPSHOT_CONTROL_DIR"); + const char* control = getenv("SNAPSHOT_CONTROL_DIR"); + if (control == NULL || control[0] == '\0') + control = getenv("DYN_SNAPSHOT_CONTROL_DIR"); if (control != NULL && control[0] != '\0') control_dir = control; } @@ -630,6 +822,12 @@ main(int argc, char** argv) fprintf(stderr, "prepare failed: topology validate\n"); goto done; } + /* Carriers are local to this request. Every rank must finish multicast + * teardown before PREPARE unmaps unicast. */ + if (command_all(participants, participant_count, CUINTERPOSER_PREPARE_MULTICAST) != 0) { + fprintf(stderr, "prepare failed: multicast teardown\n"); + goto done; + } if (command_all(participants, participant_count, CUINTERPOSER_PREPARE) != 0) { fprintf(stderr, "prepare failed: participant prepare\n"); goto done; @@ -649,7 +847,8 @@ main(int argc, char** argv) if (identify(participants, participant_count) != 0 || same_participants(expected, expected_count, participants, participant_count) != 0) goto done; - if (restore_unicast(participants, participant_count) != 0) { + if (restore_unicast(participants, participant_count) != 0 || + restore_multicast(participants, participant_count) != 0) { goto done; } for (index = 0; index < participant_count; index++) { diff --git a/agent/cmd/cuinterpose/interpose.c b/agent/cmd/cuinterpose/interpose.c index 10223a57..1d0aeedb 100644 --- a/agent/cmd/cuinterpose/interpose.c +++ b/agent/cmd/cuinterpose/interpose.c @@ -18,13 +18,18 @@ #include #include +#include "multicast.h" #include "posix.h" #include "protocol.h" #include "symbols.h" #include "util.h" +#undef cuMulticastBindAddr +#undef cuMulticastBindMem + #define CONTROL_DIR "/snapshot-control" -#define CONTROL_TIMEOUT_SECONDS 30 +/* See export_timeout_seconds() in posix.c for why 30 s was too tight. */ +#define CONTROL_TIMEOUT_SECONDS_DEFAULT 300 #define LOGICAL_HANDLE_TAG UINT64_C(0xd94d000000000000) #define LOGICAL_HANDLE_TAG_MASK UINT64_C(0xffff000000000000) #define LOGICAL_HANDLE_VALUE_MASK UINT64_C(0x0000ffffffffffff) @@ -85,8 +90,13 @@ struct allocation { enum phase { PHASE_ACTIVE, PHASE_CARRIERS, + PHASE_MULTICAST_DETACHED, PHASE_PREPARED, PHASE_CREATORS_RESTORED, + PHASE_UNICAST_RESTORED, + PHASE_MULTICAST_CREATED, + PHASE_MULTICAST_IMPORTED, + PHASE_MULTICAST_JOINED, PHASE_FAILED, }; @@ -107,6 +117,29 @@ static char socket_path[sizeof(((struct sockaddr_un*)0)->sun_path)]; static int listener = -1; static bool endpoint_needs_initialization; static uint64_t next_logical_handle = 1; +static struct handle* resolve_managed_handle(CUmemGenericAllocationHandle logical); +static struct mapping* find_mapping_at(CUdeviceptr address); +static struct mapping* first_mapping(const struct allocation* allocation); +static struct handle* first_live_handle(const struct allocation* allocation); + +static void +release_state_lock(void) +{ + pthread_mutex_unlock(&state_lock); +} + +static void +acquire_state_lock(void) +{ + pthread_mutex_lock(&state_lock); +} + +/* Caller must hold state_lock. Used to revalidate after a lock drop. */ +static bool +state_is_active(void) +{ + return current_phase == PHASE_ACTIVE; +} static void set_failure(const char* message) @@ -137,6 +170,66 @@ find_allocation(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]) return NULL; } +static int +multicast_member_from_handle(CUmemGenericAllocationHandle logical, struct cuinterposer_multicast_member* member) +{ + struct handle* handle = resolve_managed_handle(logical); + + if (handle == NULL || handle->driver == 0) + return -1; + memcpy(member->id, handle->allocation->id, sizeof(member->id)); + member->handle = handle->driver; + member->device = handle->allocation->properties.location.id; + return 0; +} + +static int +multicast_member_from_address(CUdeviceptr address, size_t size, struct cuinterposer_multicast_member* member) +{ + struct mapping* mapping = find_mapping_at(address); + + if (mapping == NULL || size > mapping->size || address - mapping->address > mapping->size - size) + return -1; + memcpy(member->id, mapping->allocation->id, sizeof(member->id)); + member->address = address; + member->allocation_offset = mapping->offset + (size_t)(address - mapping->address); + member->device = mapping->allocation->properties.location.id; + return 0; +} + +static void +multicast_mark_member_shared(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]) +{ + struct allocation* allocation = find_allocation(id); + + if (allocation != NULL) + allocation->shared = true; +} + +static int +multicast_member_from_id(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE], struct cuinterposer_multicast_member* member) +{ + retain_fn retain = (retain_fn)cuinterposer_lookup_real_symbol("cuMemRetainAllocationHandle"); + struct allocation* allocation = find_allocation(id); + struct handle* handle; + struct mapping* mapping; + + if (allocation == NULL) + return -1; + memcpy(member->id, allocation->id, sizeof(member->id)); + member->device = allocation->properties.location.id; + handle = first_live_handle(allocation); + if (handle != NULL && handle->driver != 0) { + member->handle = handle->driver; + return 0; + } + mapping = first_mapping(allocation); + if (mapping == NULL || retain == NULL || retain(&member->handle, (void*)(uintptr_t)mapping->address) != CUDA_SUCCESS) + return -1; + member->temporary_handle = true; + return 0; +} + static bool is_logical_handle(CUmemGenericAllocationHandle handle) { @@ -347,16 +440,21 @@ request_export(const struct cuinterposer_posix_ticket* ticket, int* output, char if (error != NULL && error_size != 0) error[0] = '\0'; if (strcmp(ticket->creator_participant, participant_id) == 0) { - struct allocation* allocation; - CUresult export_result; + CUresult export_result = CUDA_ERROR_INVALID_HANDLE; pthread_mutex_lock(&state_lock); - allocation = find_allocation(ticket->allocation_id); - if (allocation == NULL || - memcmp(allocation->authorization, ticket->authorization, sizeof(allocation->authorization)) != 0) { + if (ticket->resource_kind == CUINTERPOSER_RESOURCE_MULTICAST) { + export_result = cuinterposer_multicast_export_raw(ticket->allocation_id, ticket->authorization, output); + } else { + struct allocation* allocation = find_allocation(ticket->allocation_id); + if (allocation != NULL && + memcmp(allocation->authorization, ticket->authorization, sizeof(allocation->authorization)) == 0) + export_result = export_raw(allocation, output); + } + if (export_result == CUDA_ERROR_INVALID_HANDLE) { if (error != NULL && error_size != 0) - snprintf(error, error_size, "%s", "creator allocation is unavailable"); - } else if ((export_result = export_raw(allocation, output)) != CUDA_SUCCESS) { + snprintf(error, error_size, "%s", "creator resource is unavailable"); + } else if (export_result != CUDA_SUCCESS) { if (error != NULL && error_size != 0) snprintf(error, error_size, "creator export failed: CUresult=%d", (int)export_result); } else { @@ -386,6 +484,7 @@ inspect_records(uint32_t* count) if (mapping->allocation->shared && mapping->mapped) total++; } + total += cuinterposer_multicast_record_count(); if (total > CUINTERPOSER_MAX_RECORDS) return NULL; records = calloc(total == 0 ? 1 : total, sizeof(*records)); @@ -426,6 +525,10 @@ inspect_records(uint32_t* count) } record++; } + if (cuinterposer_multicast_write_records(record, total - (size_t)(record - records)) != 0) { + free(records); + return NULL; + } *count = (uint32_t)total; return records; } @@ -488,14 +591,13 @@ prepare_topology(void) struct handle* handle; struct mapping* mapping; - if (current_phase != PHASE_CARRIERS || release == NULL || unmap == NULL) + if (current_phase != PHASE_MULTICAST_DETACHED || release == NULL || unmap == NULL) return -1; for (allocation = allocations; allocation != NULL; allocation = allocation->next) { if (allocation->shared && allocation->creator && (live_handle_count(allocation) != 0 || first_mapping(allocation) != NULL) && allocation->carrier == 0) goto failed; } - current_phase = PHASE_PREPARED; for (allocation = allocations; allocation != NULL; allocation = allocation->next) { if (!allocation->shared || (live_handle_count(allocation) == 0 && first_mapping(allocation) == NULL && allocation->carrier == 0)) @@ -531,6 +633,7 @@ prepare_topology(void) if (leave_context(&scope) != 0) goto failed; } + current_phase = PHASE_PREPARED; return 0; failed: set_failure("cannot prepare cuinterposer topology"); @@ -657,6 +760,7 @@ restore_importers(void) memset(&ticket, 0, sizeof(ticket)); ticket.magic = CUINTERPOSER_POSIX_TICKET_MAGIC; ticket.version = CUINTERPOSER_POSIX_TICKET_VERSION; + ticket.resource_kind = CUINTERPOSER_RESOURCE_UNICAST; snprintf( ticket.creator_participant, sizeof(ticket.creator_participant), "%s", allocation->creator_participant); memcpy(ticket.allocation_id, allocation->id, sizeof(ticket.allocation_id)); @@ -714,7 +818,7 @@ restore_importers(void) return -1; } } - current_phase = PHASE_ACTIVE; + current_phase = PHASE_UNICAST_RESTORED; failure[0] = '\0'; return 0; } @@ -770,8 +874,20 @@ serve(int client) (void)write_all(client, records, (size_t)response.payload_size); break; case CUINTERPOSER_PREPARE: - if (create_checkpoint_carriers() != 0 || prepare_topology() != 0) + if (prepare_topology() != 0) + header_error(&response, failure); + (void)send_header(client, &response, -1); + break; + case CUINTERPOSER_PREPARE_MULTICAST: + /* Keep a unicast carrier, then tear down multicast before any rank unmaps UC. */ + if (create_checkpoint_carriers() != 0) { + header_error(&response, failure); + } else if (cuinterposer_multicast_prepare() != 0) { + set_failure(cuinterposer_multicast_error()); header_error(&response, failure); + } else { + current_phase = PHASE_MULTICAST_DETACHED; + } (void)send_header(client, &response, -1); break; case CUINTERPOSER_RESTORE_CREATORS: @@ -784,17 +900,68 @@ serve(int client) header_error(&response, failure); (void)send_header(client, &response, -1); break; + case CUINTERPOSER_RESTORE_MULTICAST_CREATORS: + if (current_phase != PHASE_UNICAST_RESTORED || cuinterposer_multicast_restore_creators() != 0) { + set_failure(cuinterposer_multicast_error()); + header_error(&response, failure); + } else { + current_phase = PHASE_MULTICAST_CREATED; + } + (void)send_header(client, &response, -1); + break; + case CUINTERPOSER_RESTORE_MULTICAST_IMPORTERS: + if (current_phase != PHASE_MULTICAST_CREATED || cuinterposer_multicast_restore_importers() != 0) { + set_failure(cuinterposer_multicast_error()); + header_error(&response, failure); + } else { + current_phase = PHASE_MULTICAST_IMPORTED; + } + (void)send_header(client, &response, -1); + break; + case CUINTERPOSER_RESTORE_MULTICAST_DEVICES: + if (current_phase != PHASE_MULTICAST_IMPORTED || cuinterposer_multicast_restore_devices() != 0) { + set_failure(cuinterposer_multicast_error()); + header_error(&response, failure); + } else { + current_phase = PHASE_MULTICAST_JOINED; + } + (void)send_header(client, &response, -1); + break; + case CUINTERPOSER_RESTORE_MULTICAST: + if (current_phase != PHASE_MULTICAST_JOINED || cuinterposer_multicast_restore_topology() != 0) { + set_failure(cuinterposer_multicast_error()); + header_error(&response, failure); + } else { + current_phase = PHASE_ACTIVE; + failure[0] = '\0'; + } + (void)send_header(client, &response, -1); + break; case CUINTERPOSER_EXPORT: { - struct allocation* allocation = find_allocation(request.allocation_id); CUresult export_result; - if (strcmp(request.participant_id, participant_id) != 0 || allocation == NULL || !allocation->creator || - memcmp(allocation->authorization, request.authorization, sizeof(request.authorization)) != 0 || - (current_phase != PHASE_ACTIVE && current_phase != PHASE_CREATORS_RESTORED)) { - header_error(&response, "creator allocation is unavailable"); + if (strcmp(request.participant_id, participant_id) != 0 || + (request.resource_kind != CUINTERPOSER_RESOURCE_UNICAST && + request.resource_kind != CUINTERPOSER_RESOURCE_MULTICAST) || + (current_phase != PHASE_ACTIVE && current_phase != PHASE_CREATORS_RESTORED && + current_phase != PHASE_UNICAST_RESTORED && current_phase != PHASE_MULTICAST_CREATED && + current_phase != PHASE_MULTICAST_IMPORTED && current_phase != PHASE_MULTICAST_JOINED)) { + header_error(&response, "creator resource is unavailable"); (void)send_header(client, &response, -1); break; } - export_result = export_raw(allocation, &exported_fd); + response.resource_kind = request.resource_kind; + if (request.resource_kind == CUINTERPOSER_RESOURCE_MULTICAST) { + export_result = cuinterposer_multicast_export_raw(request.allocation_id, request.authorization, &exported_fd); + } else { + struct allocation* allocation = find_allocation(request.allocation_id); + if (allocation == NULL || !allocation->creator || + memcmp(allocation->authorization, request.authorization, sizeof(request.authorization)) != 0) { + header_error(&response, "creator allocation is unavailable"); + (void)send_header(client, &response, -1); + break; + } + export_result = export_raw(allocation, &exported_fd); + } if (export_result != CUDA_SUCCESS) { char message[sizeof(response.message)]; snprintf(message, sizeof(message), "creator export failed: CUresult=%d", (int)export_result); @@ -802,7 +969,7 @@ serve(int client) (void)send_header(client, &response, -1); break; } - memcpy(response.allocation_id, allocation->id, sizeof(response.allocation_id)); + memcpy(response.allocation_id, request.allocation_id, sizeof(response.allocation_id)); (void)send_header(client, &response, exported_fd); break; } @@ -820,20 +987,56 @@ serve(int client) free(records); } +static unsigned +control_timeout_seconds(void) +{ + static unsigned cached; + + if (cached == 0) { + const char* value = getenv("DYN_SNAPSHOT_CONTROL_TIMEOUT_SECONDS"); + + cached = bounded_seconds(value, CONTROL_TIMEOUT_SECONDS_DEFAULT); + } + return cached; +} + +static void* +control_connection(void* argument) +{ + int client = (int)(intptr_t)argument; + + if (set_socket_timeouts(client, control_timeout_seconds()) == 0) + serve(client); + close(client); + return NULL; +} + static void* control_agent(void* unused) { (void)unused; for (;;) { int client = accept4(listener, NULL, NULL, SOCK_CLOEXEC); + pthread_t worker; + if (client < 0) { if (errno == EINTR) continue; return NULL; } - if (set_socket_timeouts(client, CONTROL_TIMEOUT_SECONDS) == 0) - serve(client); - close(client); + /* + * Serve each connection on its own detached thread so the accept loop never + * stalls. serve() still takes state_lock for its whole body, so requests + * remain serialised -- the win is that a slow request no longer blocks + * accept(), which previously left peers stuck in connect() burning their + * export deadline before they were ever heard. Participant count is small + * and bounded, so thread churn is not a concern. Fall back to inline + * service if a thread cannot be created. + */ + if (pthread_create(&worker, NULL, control_connection, (void*)(intptr_t)client) == 0) + (void)pthread_detach(worker); + else + (void)control_connection((void*)(intptr_t)client); } } @@ -893,6 +1096,7 @@ fork_child(void) allocations = NULL; handles = NULL; mappings = NULL; + cuinterposer_multicast_reset(); next_logical_handle = 1; current_phase = PHASE_ACTIVE; failure[0] = '\0'; @@ -922,6 +1126,16 @@ ensure_process_endpoint(void) __attribute__((constructor)) static void initialize(void) { + const struct cuinterposer_multicast_callbacks multicast_callbacks = { + .allocate_logical_handle = allocate_logical_handle, + .member_from_handle = multicast_member_from_handle, + .member_from_address = multicast_member_from_address, + .member_from_id = multicast_member_from_id, + .mark_member_shared = multicast_mark_member_shared, + .release_state_lock = release_state_lock, + .acquire_state_lock = acquire_state_lock, + .state_is_active = state_is_active, + }; const char* control; const char* configured_participant; @@ -954,6 +1168,7 @@ initialize(void) set_failure("cannot start cuinterposer control endpoint"); return; } + cuinterposer_multicast_initialize(&multicast_callbacks, participant_id, socket_path); } __attribute__((destructor)) static void @@ -1027,6 +1242,11 @@ cuMemRelease(CUmemGenericAllocationHandle application) pthread_mutex_lock(&state_lock); handle = resolve_managed_handle(application); if (handle == NULL) { + if (cuinterposer_multicast_is_handle(application)) { + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_release(application) : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; + } pthread_mutex_unlock(&state_lock); if (is_logical_handle(application)) return CUDA_ERROR_INVALID_HANDLE; @@ -1060,6 +1280,13 @@ cuMemRetainAllocationHandle(CUmemGenericAllocationHandle* output, void* address) return cuinterposer_unavailable(); if (output == NULL) return CUDA_ERROR_INVALID_VALUE; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_retain(output, address) : CUDA_ERROR_NOT_READY; + if (result != CUDA_ERROR_INVALID_VALUE) { + pthread_mutex_unlock(&state_lock); + return result; + } + pthread_mutex_unlock(&state_lock); result = function(&driver, address); if (result != CUDA_SUCCESS) return result; @@ -1093,6 +1320,12 @@ cuMemMap( pthread_mutex_lock(&state_lock); handle = resolve_managed_handle(application); if (handle == NULL) { + if (cuinterposer_multicast_is_handle(application)) { + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_map(address, size, offset, application, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; + } pthread_mutex_unlock(&state_lock); if (is_logical_handle(application)) return CUDA_ERROR_INVALID_HANDLE; @@ -1136,6 +1369,11 @@ cuMemUnmap(CUdeviceptr address, size_t size) pthread_mutex_lock(&state_lock); mapping = find_mapping(address, size); if (mapping == NULL) { + if (cuinterposer_multicast_has_mapping(address, size)) { + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_unmap(address, size) : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; + } pthread_mutex_unlock(&state_lock); return function != NULL ? function(address, size) : cuinterposer_unavailable(); } @@ -1150,10 +1388,29 @@ cuMemUnmap(CUdeviceptr address, size_t size) return result; } +/* + * Record the access descriptors on every tracked mapping that `[address,size)` + * covers. + * + * Restore replays access only when `mapping->access_count != 0` (see + * restore_mappings). Matching solely on an exact (address, size) pair meant any + * caller that mapped at one granularity and set access at another -- which + * PyTorch's caching allocator routinely does, and which is legal CUDA -- left + * `access_count == 0`. Such a mapping was restored mapped but with NO device + * access, so the VA resolved and the first kernel to touch it died with + * CUDA_ERROR_ILLEGAL_ADDRESS. Observed on an 8-rank GLM-5.2 restore: 192 of 336 + * recorded mappings had access_count == 0. + * + * Returns false if the range partially overlaps a tracked mapping, i.e. the + * request cannot be represented per-mapping. The caller then fails closed + * rather than passing it through and silently losing the access on restore. + * + * When `descriptors` is NULL nothing is mutated; the walk only classifies the + * range. `*matched` receives the number of tracked mappings fully covered. + */ static bool record_access_over_range( - CUdeviceptr address, size_t size, const CUmemAccessDesc* descriptors, size_t count, - size_t* matched) + CUdeviceptr address, size_t size, const CUmemAccessDesc* descriptors, size_t count, size_t* matched) { struct mapping* mapping; CUdeviceptr end; @@ -1171,9 +1428,9 @@ record_access_over_range( return false; mapping_end = mapping->address + mapping->size; if (mapping_end <= address || mapping->address >= end) - continue; + continue; /* disjoint */ if (mapping->address < address || mapping_end > end) - return false; + return false; /* partial overlap: not representable per-mapping */ if (descriptors != NULL) { memcpy(mapping->access, descriptors, count * sizeof(*descriptors)); mapping->access_count = count; @@ -1193,21 +1450,33 @@ cuMemSetAccess(CUdeviceptr address, size_t size, const CUmemAccessDesc* descript if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) return result; pthread_mutex_lock(&state_lock); + + /* Multicast ranges keep their own bookkeeping. */ + if (find_mapping(address, size) == NULL && cuinterposer_multicast_has_mapping(address, size)) { + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_set_access(address, size, descriptors, count) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; + } + + /* + * Classify first (NULL descriptors => no mutation) so a range that touches + * nothing we track stays a pure passthrough, and so an unrepresentable range + * is rejected before the device is modified. + */ if (!record_access_over_range(address, size, NULL, 0, &matched)) { pthread_mutex_unlock(&state_lock); return CUDA_ERROR_NOT_SUPPORTED; } if (matched == 0) { pthread_mutex_unlock(&state_lock); - return function != NULL ? function(address, size, descriptors, count) - : cuinterposer_unavailable(); + return function != NULL ? function(address, size, descriptors, count) : cuinterposer_unavailable(); } if (current_phase != PHASE_ACTIVE || count > CUINTERPOSER_MAX_ACCESS) { pthread_mutex_unlock(&state_lock); return CUDA_ERROR_NOT_SUPPORTED; } - result = function != NULL ? function(address, size, descriptors, count) - : cuinterposer_unavailable(); + result = function != NULL ? function(address, size, descriptors, count) : cuinterposer_unavailable(); if (result == CUDA_SUCCESS) (void)record_access_over_range(address, size, descriptors, count, &matched); pthread_mutex_unlock(&state_lock); @@ -1231,6 +1500,12 @@ cuMemExportToShareableHandle( pthread_mutex_lock(&state_lock); handle = resolve_managed_handle(application); if (handle == NULL) { + if (cuinterposer_multicast_is_handle(application)) { + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_export(shareable, application, type, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; + } pthread_mutex_unlock(&state_lock); if (is_logical_handle(application)) return CUDA_ERROR_INVALID_HANDLE; @@ -1246,6 +1521,7 @@ cuMemExportToShareableHandle( memset(&ticket, 0, sizeof(ticket)); ticket.magic = CUINTERPOSER_POSIX_TICKET_MAGIC; ticket.version = CUINTERPOSER_POSIX_TICKET_VERSION; + ticket.resource_kind = CUINTERPOSER_RESOURCE_UNICAST; snprintf(ticket.creator_participant, sizeof(ticket.creator_participant), "%s", handle->allocation->creator_participant); memcpy(ticket.allocation_id, handle->allocation->id, sizeof(ticket.allocation_id)); snprintf(ticket.creator_endpoint, sizeof(ticket.creator_endpoint), "%s", handle->allocation->creator_endpoint); @@ -1288,6 +1564,19 @@ cuMemImportFromShareableHandle(CUmemGenericAllocationHandle* output, void* os_ha return CUDA_ERROR_INVALID_HANDLE; if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) return result; + if (ticket.resource_kind == CUINTERPOSER_RESOURCE_MULTICAST) { + if (request_export(&ticket, &raw_fd, NULL, 0) != 0) + return CUDA_ERROR_INVALID_HANDLE; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_import(output, &ticket, raw_fd) + : CUDA_ERROR_INVALID_HANDLE; + pthread_mutex_unlock(&state_lock); + if (close(raw_fd) != 0 && result == CUDA_SUCCESS) { + (void)cuMemRelease(*output); + return CUDA_ERROR_UNKNOWN; + } + return result; + } pthread_mutex_lock(&state_lock); if (current_phase != PHASE_ACTIVE) { pthread_mutex_unlock(&state_lock); @@ -1355,14 +1644,142 @@ cuMemGetAllocationPropertiesFromHandle(CUmemAllocationProp* properties, CUmemGen return result; pthread_mutex_lock(&state_lock); handle = resolve_managed_handle(application); - if (handle == NULL) - result = is_logical_handle(application) ? CUDA_ERROR_INVALID_HANDLE - : (function != NULL ? function(properties, application) : cuinterposer_unavailable()); - else + if (handle == NULL) { + if (cuinterposer_multicast_is_handle(application)) + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_get_properties(properties, application) + : CUDA_ERROR_NOT_READY; + else + result = is_logical_handle(application) ? CUDA_ERROR_INVALID_HANDLE + : (function != NULL ? function(properties, application) : cuinterposer_unavailable()); + } else { result = current_phase != PHASE_ACTIVE || handle->driver == 0 ? CUDA_ERROR_NOT_READY : (function != NULL ? function(properties, handle->driver) : cuinterposer_unavailable()); + } + pthread_mutex_unlock(&state_lock); + return result; +} + +CUresult CUDAAPI +cuMulticastCreate(CUmemGenericAllocationHandle* output, const CUmulticastObjectProp* properties) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_create(output, properties) : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} + +CUresult CUDAAPI +cuMulticastAddDevice(CUmemGenericAllocationHandle multicast, CUdevice device) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE ? cuinterposer_multicast_add_device(multicast, device) : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} + +CUresult CUDAAPI +cuMulticastBindMem( + CUmemGenericAllocationHandle multicast, size_t multicast_offset, CUmemGenericAllocationHandle memory, + size_t memory_offset, size_t size, unsigned long long flags) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE + ? cuinterposer_multicast_bind_mem(multicast, 0, false, multicast_offset, memory, memory_offset, size, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} + +#if CUDA_VERSION >= 13010 +CUresult CUDAAPI +cuMulticastBindMem_v2( + CUmemGenericAllocationHandle multicast, CUdevice device, size_t multicast_offset, + CUmemGenericAllocationHandle memory, size_t memory_offset, size_t size, unsigned long long flags) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = + current_phase == PHASE_ACTIVE + ? cuinterposer_multicast_bind_mem(multicast, device, true, multicast_offset, memory, memory_offset, size, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} +#endif + +CUresult CUDAAPI +cuMulticastBindAddr( + CUmemGenericAllocationHandle multicast, size_t multicast_offset, CUdeviceptr memory, size_t size, + unsigned long long flags) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE + ? cuinterposer_multicast_bind_address(multicast, 0, false, multicast_offset, memory, size, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} + +#if CUDA_VERSION >= 13010 +CUresult CUDAAPI +cuMulticastBindAddr_v2( + CUmemGenericAllocationHandle multicast, CUdevice device, size_t multicast_offset, CUdeviceptr memory, size_t size, + unsigned long long flags) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = current_phase == PHASE_ACTIVE + ? cuinterposer_multicast_bind_address(multicast, device, true, multicast_offset, memory, size, flags) + : CUDA_ERROR_NOT_READY; + pthread_mutex_unlock(&state_lock); + return result; +} +#endif + +CUresult CUDAAPI +cuMulticastGetGranularity( + size_t* granularity, const CUmulticastObjectProp* properties, CUmulticastGranularity_flags option) +{ + typedef CUresult(CUDAAPI * function_type)(size_t*, const CUmulticastObjectProp*, CUmulticastGranularity_flags); + function_type function = (function_type)cuinterposer_lookup_real_symbol("cuMulticastGetGranularity"); + + return function != NULL ? function(granularity, properties, option) : cuinterposer_unavailable(); +} + +CUresult CUDAAPI +cuMulticastUnbind(CUmemGenericAllocationHandle multicast, CUdevice device, size_t offset, size_t size) +{ + CUresult result; + + if ((result = ensure_process_endpoint()) != CUDA_SUCCESS) + return result; + pthread_mutex_lock(&state_lock); + result = + current_phase == PHASE_ACTIVE ? cuinterposer_multicast_unbind(multicast, device, offset, size) : CUDA_ERROR_NOT_READY; pthread_mutex_unlock(&state_lock); return result; } diff --git a/agent/cmd/cuinterpose/multicast.c b/agent/cmd/cuinterpose/multicast.c index 1da08c25..74022acf 100644 --- a/agent/cmd/cuinterpose/multicast.c +++ b/agent/cmd/cuinterpose/multicast.c @@ -4,113 +4,1416 @@ * SPDX-License-Identifier: Apache-2.0 */ +#define _GNU_SOURCE + +#include "multicast.h" #include "symbols.h" +#include "util.h" + +#include +#include +#include +#include + +typedef CUresult(CUDAAPI* release_fn)(CUmemGenericAllocationHandle); +typedef CUresult(CUDAAPI* retain_fn)(CUmemGenericAllocationHandle*, void*); +typedef CUresult(CUDAAPI* map_fn)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); +typedef CUresult(CUDAAPI* unmap_fn)(CUdeviceptr, size_t); +typedef CUresult(CUDAAPI* access_fn)(CUdeviceptr, size_t, const CUmemAccessDesc*, size_t); +typedef CUresult(CUDAAPI* export_fn)( + void*, CUmemGenericAllocationHandle, CUmemAllocationHandleType, unsigned long long); +typedef CUresult(CUDAAPI* import_fn)(CUmemGenericAllocationHandle*, void*, CUmemAllocationHandleType); +typedef CUresult(CUDAAPI* properties_fn)(CUmemAllocationProp*, CUmemGenericAllocationHandle); +typedef CUresult(CUDAAPI* context_get_fn)(CUcontext*); +typedef CUresult(CUDAAPI* context_set_fn)(CUcontext); +typedef CUresult(CUDAAPI* context_device_fn)(CUdevice*); +typedef CUresult(CUDAAPI* create_fn)(CUmemGenericAllocationHandle*, const CUmulticastObjectProp*); +typedef CUresult(CUDAAPI* add_device_fn)(CUmemGenericAllocationHandle, CUdevice); +typedef CUresult(CUDAAPI* bind_mem_fn)( + CUmemGenericAllocationHandle, size_t, CUmemGenericAllocationHandle, size_t, size_t, unsigned long long); +typedef CUresult(CUDAAPI* bind_address_fn)( + CUmemGenericAllocationHandle, size_t, CUdeviceptr, size_t, unsigned long long); +#if CUDA_VERSION >= 13010 +typedef CUresult(CUDAAPI* bind_mem_v2_fn)( + CUmemGenericAllocationHandle, CUdevice, size_t, CUmemGenericAllocationHandle, size_t, size_t, unsigned long long); +typedef CUresult(CUDAAPI* bind_address_v2_fn)( + CUmemGenericAllocationHandle, CUdevice, size_t, CUdeviceptr, size_t, unsigned long long); +#endif +typedef CUresult(CUDAAPI* unbind_fn)(CUmemGenericAllocationHandle, CUdevice, size_t, size_t); + +struct multicast; + +struct multicast_handle { + CUmemGenericAllocationHandle logical; + CUmemGenericAllocationHandle driver; + bool live; + struct multicast* multicast; + struct multicast_handle* next; +}; + +struct multicast_device { + CUdevice device; + struct multicast_device* next; +}; + +struct multicast_binding { + uint8_t member_id[CUINTERPOSER_ALLOCATION_ID_SIZE]; + CUdeviceptr member_address; + size_t multicast_offset; + size_t member_offset; + size_t size; + unsigned long long flags; + CUdevice device; + uint8_t kind; + uint8_t api_version; + bool bound; + bool checkpointed; + struct multicast_binding* next; +}; + +struct multicast_mapping { + CUdeviceptr address; + size_t size; + size_t offset; + unsigned long long flags; + CUmemAccessDesc access[CUINTERPOSER_MAX_ACCESS]; + size_t access_count; + bool mapped; + bool checkpointed; + struct multicast_mapping* next; +}; + +struct multicast { + uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]; + uint8_t authorization[CUINTERPOSER_TOKEN_SIZE]; + char creator_participant[CUINTERPOSER_ID_SIZE]; + char creator_endpoint[sizeof(((struct sockaddr_un*)0)->sun_path)]; + CUmulticastObjectProp properties; + size_t effective_size; + CUcontext context; + CUmemGenericAllocationHandle restore_handle; + bool creator; + bool checkpointed; + struct multicast_device* devices; + struct multicast_binding* bindings; + struct multicast_mapping* mappings; + struct multicast* next; +}; + +struct context_scope { + CUcontext previous; + bool changed; +}; + +static struct cuinterposer_multicast_callbacks operations; +static const char* current_participant; +static const char* current_endpoint; +static struct multicast* multicasts; +static struct multicast_handle* handles; +static char failure[128]; + +static int +fail(const char* operation, CUresult result) +{ + if (result == CUDA_SUCCESS) + snprintf(failure, sizeof(failure), "%s", operation); + else + snprintf(failure, sizeof(failure), "%s failed: CUresult=%d", operation, (int)result); + return -1; +} + +static struct multicast* +find_multicast(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]) +{ + struct multicast* multicast; + + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + if (memcmp(multicast->id, id, CUINTERPOSER_ALLOCATION_ID_SIZE) == 0) + return multicast; + } + return NULL; +} + +static struct multicast_handle* +find_handle(CUmemGenericAllocationHandle logical) +{ + struct multicast_handle* handle; + + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle->logical == logical) + return handle; + } + return NULL; +} + +static struct multicast_mapping* +find_mapping(CUdeviceptr address, size_t size) +{ + struct multicast* multicast; -#undef cuMulticastBindAddr -#undef cuMulticastBindMem + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct multicast_mapping* mapping; + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped && mapping->address == address && mapping->size == size) + return mapping; + } + } + return NULL; +} + +static struct multicast_mapping* +find_mapping_at(CUdeviceptr address) +{ + struct multicast* multicast; -#define LOOKUP(name, type) ((type)cuinterposer_lookup_real_symbol(#name)) + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct multicast_mapping* mapping; + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped && address >= mapping->address && address < mapping->address + mapping->size) + return mapping; + } + } + return NULL; +} -CUresult CUDAAPI -cuMulticastCreate(CUmemGenericAllocationHandle* output, const CUmulticastObjectProp* properties) +static struct multicast* +mapping_owner(const struct multicast_mapping* expected) { - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle*, const CUmulticastObjectProp*); - function_type function = LOOKUP(cuMulticastCreate, function_type); - return function != NULL ? function(output, properties) : cuinterposer_unavailable(); + struct multicast* multicast; + + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + const struct multicast_mapping* mapping; + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping == expected) + return multicast; + } + } + return NULL; } -CUresult CUDAAPI -cuMulticastAddDevice(CUmemGenericAllocationHandle multicast, CUdevice device) +static size_t +live_handle_count(const struct multicast* multicast) { - typedef CUresult(CUDAAPI * function_type)(CUmemGenericAllocationHandle, CUdevice); - function_type function = LOOKUP(cuMulticastAddDevice, function_type); - return function != NULL ? function(multicast, device) : cuinterposer_unavailable(); + const struct multicast_handle* handle; + size_t count = 0; + + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle->live && handle->multicast == multicast) + count++; + } + return count; } -CUresult CUDAAPI -cuMulticastBindMem( - CUmemGenericAllocationHandle multicast, size_t multicast_offset, - CUmemGenericAllocationHandle memory, size_t memory_offset, size_t size, - unsigned long long flags) +static CUmemGenericAllocationHandle +current_driver(const struct multicast* multicast) { - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle, size_t, CUmemGenericAllocationHandle, size_t, size_t, - unsigned long long); - function_type function = LOOKUP(cuMulticastBindMem, function_type); - return function != NULL - ? function(multicast, multicast_offset, memory, memory_offset, size, flags) - : cuinterposer_unavailable(); + const struct multicast_handle* handle; + + if (multicast->restore_handle != 0) + return multicast->restore_handle; + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle->live && handle->multicast == multicast && handle->driver != 0) + return handle->driver; + } + return 0; } +static bool +driver_used(const struct multicast_handle* except, CUmemGenericAllocationHandle driver) +{ + const struct multicast_handle* handle; + + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle != except && handle->live && handle->driver == driver) + return true; + } + return false; +} + +static bool +matches_ticket(const struct multicast* multicast, const struct cuinterposer_posix_ticket* ticket) +{ + return memcmp(multicast->authorization, ticket->authorization, sizeof(multicast->authorization)) == 0 && + strcmp(multicast->creator_participant, ticket->creator_participant) == 0 && + strcmp(multicast->creator_endpoint, ticket->creator_endpoint) == 0 && + multicast->properties.numDevices == ticket->num_devices && + multicast->properties.size == ticket->allocation_size && + multicast->properties.handleTypes == ticket->handle_types && + multicast->properties.flags == ticket->object_flags; +} + +static bool +active(const struct multicast* multicast) +{ + const struct multicast_binding* binding; + const struct multicast_mapping* mapping; + + if (live_handle_count(multicast) != 0) + return true; + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + if (binding->bound) + return true; + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped) + return true; + } + return false; +} + +static int +add_handle(struct multicast* multicast, CUmemGenericAllocationHandle driver, CUmemGenericAllocationHandle* logical) +{ + struct multicast_handle* handle = calloc(1, sizeof(*handle)); + + if (handle == NULL || operations.allocate_logical_handle == NULL || + operations.allocate_logical_handle(logical) != 0) { + free(handle); + return -1; + } + handle->logical = *logical; + handle->driver = driver; + handle->live = true; + handle->multicast = multicast; + handle->next = handles; + handles = handle; + return 0; +} + +static void +install_driver(struct multicast* multicast, CUmemGenericAllocationHandle driver) +{ + struct multicast_handle* handle; + + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle->live && handle->multicast == multicast) + handle->driver = driver; + } + multicast->restore_handle = driver; +} + +static int +enter_context(CUcontext context, struct context_scope* scope) +{ + context_get_fn get_current = (context_get_fn)cuinterposer_lookup_real_symbol("cuCtxGetCurrent"); + context_set_fn set_current = (context_set_fn)cuinterposer_lookup_real_symbol("cuCtxSetCurrent"); + + memset(scope, 0, sizeof(*scope)); + if (get_current == NULL || set_current == NULL || get_current(&scope->previous) != CUDA_SUCCESS) + return -1; + if (scope->previous != context) { + if (set_current(context) != CUDA_SUCCESS) + return -1; + scope->changed = true; + } + return 0; +} + +static int +leave_context(const struct context_scope* scope) +{ + context_set_fn set_current = (context_set_fn)cuinterposer_lookup_real_symbol("cuCtxSetCurrent"); + + return !scope->changed || (set_current != NULL && set_current(scope->previous) == CUDA_SUCCESS) ? 0 : -1; +} + +static int +capture_context(CUcontext* context) +{ + context_get_fn get_current = (context_get_fn)cuinterposer_lookup_real_symbol("cuCtxGetCurrent"); + + return get_current != NULL && get_current(context) == CUDA_SUCCESS && *context != NULL ? 0 : -1; +} + +static int +capture_device(CUdevice* device) +{ + context_device_fn get_device = (context_device_fn)cuinterposer_lookup_real_symbol("cuCtxGetDevice"); + + return get_device != NULL && get_device(device) == CUDA_SUCCESS ? 0 : -1; +} + +static int +fill_ticket(const struct multicast* multicast, struct cuinterposer_posix_ticket* ticket) +{ + memset(ticket, 0, sizeof(*ticket)); + ticket->magic = CUINTERPOSER_POSIX_TICKET_MAGIC; + ticket->version = CUINTERPOSER_POSIX_TICKET_VERSION; + ticket->resource_kind = CUINTERPOSER_RESOURCE_MULTICAST; + snprintf( + ticket->creator_participant, sizeof(ticket->creator_participant), "%s", multicast->creator_participant); + memcpy(ticket->allocation_id, multicast->id, sizeof(ticket->allocation_id)); + snprintf(ticket->creator_endpoint, sizeof(ticket->creator_endpoint), "%s", multicast->creator_endpoint); + memcpy(ticket->authorization, multicast->authorization, sizeof(ticket->authorization)); + ticket->allocation_size = multicast->properties.size; + ticket->handle_types = multicast->properties.handleTypes; + ticket->object_flags = multicast->properties.flags; + ticket->num_devices = multicast->properties.numDevices; + return 0; +} + +static CUresult +bind_memory( + CUmemGenericAllocationHandle driver, const struct multicast_binding* binding, CUmemGenericAllocationHandle member) +{ + CUresult result; + + /* Bind waits for the complete multicast team, so it must not block the creator endpoint under state_lock. */ #if CUDA_VERSION >= 13010 -CUresult CUDAAPI -cuMulticastBindMem_v2( - CUmemGenericAllocationHandle multicast, CUdevice device, size_t multicast_offset, - CUmemGenericAllocationHandle memory, size_t memory_offset, size_t size, - unsigned long long flags) -{ - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle, CUdevice, size_t, CUmemGenericAllocationHandle, - size_t, size_t, unsigned long long); - function_type function = LOOKUP(cuMulticastBindMem_v2, function_type); - return function != NULL - ? function( - multicast, device, multicast_offset, memory, memory_offset, size, flags) - : cuinterposer_unavailable(); + if (binding->api_version == 2) { + bind_mem_v2_fn function = (bind_mem_v2_fn)cuinterposer_lookup_real_symbol("cuMulticastBindMem_v2"); + if (function == NULL) + return cuinterposer_unavailable(); + operations.release_state_lock(); + result = function( + driver, binding->device, binding->multicast_offset, member, binding->member_offset, binding->size, + binding->flags); + operations.acquire_state_lock(); + return result; + } +#endif + { + bind_mem_fn function = (bind_mem_fn)cuinterposer_lookup_real_symbol("cuMulticastBindMem"); + if (function == NULL) + return cuinterposer_unavailable(); + operations.release_state_lock(); + result = function(driver, binding->multicast_offset, member, binding->member_offset, binding->size, binding->flags); + operations.acquire_state_lock(); + return result; + } } + +static CUresult +bind_address(CUmemGenericAllocationHandle driver, const struct multicast_binding* binding) +{ + CUresult result; + +#if CUDA_VERSION >= 13010 + if (binding->api_version == 2) { + bind_address_v2_fn function = (bind_address_v2_fn)cuinterposer_lookup_real_symbol("cuMulticastBindAddr_v2"); + if (function == NULL) + return cuinterposer_unavailable(); + operations.release_state_lock(); + result = function( + driver, binding->device, binding->multicast_offset, binding->member_address, binding->size, binding->flags); + operations.acquire_state_lock(); + return result; + } #endif + { + bind_address_fn function = (bind_address_fn)cuinterposer_lookup_real_symbol("cuMulticastBindAddr"); + if (function == NULL) + return cuinterposer_unavailable(); + operations.release_state_lock(); + result = function(driver, binding->multicast_offset, binding->member_address, binding->size, binding->flags); + operations.acquire_state_lock(); + return result; + } +} -CUresult CUDAAPI -cuMulticastBindAddr( - CUmemGenericAllocationHandle multicast, size_t multicast_offset, CUdeviceptr memory, - size_t size, unsigned long long flags) +static void +observe_extent(struct multicast* multicast, size_t offset, size_t size) { - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle, size_t, CUdeviceptr, size_t, unsigned long long); - function_type function = LOOKUP(cuMulticastBindAddr, function_type); - return function != NULL - ? function(multicast, multicast_offset, memory, size, flags) - : cuinterposer_unavailable(); + /* + * CUDA may round a multicast object's physical capacity above the size + * passed to cuMulticastCreate. Record the largest range the driver actually + * accepts, but retain the original properties for export identity and replay. + */ + if (size <= SIZE_MAX - offset && multicast->effective_size < offset + size) + multicast->effective_size = offset + size; } -#if CUDA_VERSION >= 13010 -CUresult CUDAAPI -cuMulticastBindAddr_v2( - CUmemGenericAllocationHandle multicast, CUdevice device, size_t multicast_offset, +void +cuinterposer_multicast_initialize( + const struct cuinterposer_multicast_callbacks* callbacks, const char* participant_id, const char* endpoint) +{ + operations = *callbacks; + current_participant = participant_id; + current_endpoint = endpoint; +} + +void +cuinterposer_multicast_reset(void) +{ + multicasts = NULL; + handles = NULL; + failure[0] = '\0'; +} + +bool +cuinterposer_multicast_is_handle(CUmemGenericAllocationHandle logical) +{ + return find_handle(logical) != NULL; +} + +bool +cuinterposer_multicast_has_mapping(CUdeviceptr address, size_t size) +{ + return find_mapping(address, size) != NULL; +} + +CUresult +cuinterposer_multicast_release(CUmemGenericAllocationHandle logical) +{ + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast_handle* handle = find_handle(logical); + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + result = CUDA_SUCCESS; + if (!driver_used(handle, handle->driver)) + result = release == NULL ? cuinterposer_unavailable() : release(handle->driver); + if (result == CUDA_SUCCESS) + handle->live = false; + return result; +} + +CUresult +cuinterposer_multicast_retain(CUmemGenericAllocationHandle* output, void* address) +{ + retain_fn retain = (retain_fn)cuinterposer_lookup_real_symbol("cuMemRetainAllocationHandle"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast_mapping* mapping = find_mapping_at((CUdeviceptr)(uintptr_t)address); + struct multicast* multicast; + CUmemGenericAllocationHandle driver = 0; + CUmemGenericAllocationHandle existing; + CUmemGenericAllocationHandle logical; + CUresult result; + + if (output == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (mapping == NULL) + return CUDA_ERROR_INVALID_VALUE; + multicast = mapping_owner(mapping); + if (multicast == NULL || retain == NULL) + return cuinterposer_unavailable(); + existing = current_driver(multicast); + result = retain(&driver, address); + if (result != CUDA_SUCCESS) + return result; + if (existing != 0) { + result = release == NULL ? cuinterposer_unavailable() : release(driver); + if (result != CUDA_SUCCESS) + return result; + driver = existing; + } + if (add_handle(multicast, driver, &logical) != 0) { + if (existing == 0 && release != NULL) + (void)release(driver); + return CUDA_ERROR_OUT_OF_MEMORY; + } + *output = logical; + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_map( + CUdeviceptr address, size_t size, size_t offset, CUmemGenericAllocationHandle logical, unsigned long long flags) +{ + map_fn map = (map_fn)cuinterposer_lookup_real_symbol("cuMemMap"); + unmap_fn unmap = (unmap_fn)cuinterposer_lookup_real_symbol("cuMemUnmap"); + struct multicast_handle* handle = find_handle(logical); + struct multicast_mapping* mapping; + CUmemGenericAllocationHandle driver; + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + mapping = calloc(1, sizeof(*mapping)); + if (mapping == NULL) + return CUDA_ERROR_OUT_OF_MEMORY; + /* + * Multicast map waits for the complete team. Keeping state_lock held here + * prevents the creator's control thread from exporting the driver handle + * that the other ranks need to enter the same call. + */ + driver = handle->driver; + operations.release_state_lock(); + result = map == NULL ? cuinterposer_unavailable() : map(address, size, offset, driver, flags); + operations.acquire_state_lock(); + if (result != CUDA_SUCCESS) { + free(mapping); + return result; + } + if (operations.state_is_active != NULL && !operations.state_is_active()) { + if (unmap != NULL) + (void)unmap(address, size); + free(mapping); + return CUDA_ERROR_NOT_READY; + } + handle = find_handle(logical); + if (handle == NULL || !handle->live || handle->driver != driver) { + if (unmap != NULL) + (void)unmap(address, size); + free(mapping); + return CUDA_ERROR_INVALID_HANDLE; + } + mapping->address = address; + mapping->size = size; + mapping->offset = offset; + mapping->flags = flags; + mapping->mapped = true; + mapping->next = handle->multicast->mappings; + handle->multicast->mappings = mapping; + observe_extent(handle->multicast, offset, size); + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_unmap(CUdeviceptr address, size_t size) +{ + unmap_fn unmap = (unmap_fn)cuinterposer_lookup_real_symbol("cuMemUnmap"); + struct multicast_mapping* mapping = find_mapping(address, size); + CUresult result; + + if (mapping == NULL) + return CUDA_ERROR_INVALID_VALUE; + result = unmap == NULL ? cuinterposer_unavailable() : unmap(address, size); + if (result == CUDA_SUCCESS) + mapping->mapped = false; + return result; +} + +CUresult +cuinterposer_multicast_set_access(CUdeviceptr address, size_t size, const CUmemAccessDesc* descriptors, size_t count) +{ + access_fn set_access = (access_fn)cuinterposer_lookup_real_symbol("cuMemSetAccess"); + struct multicast_mapping* mapping = find_mapping(address, size); + CUresult result; + + if (mapping == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (count > CUINTERPOSER_MAX_ACCESS) + return CUDA_ERROR_NOT_SUPPORTED; + result = set_access == NULL ? cuinterposer_unavailable() : set_access(address, size, descriptors, count); + if (result == CUDA_SUCCESS) { + memcpy(mapping->access, descriptors, count * sizeof(*descriptors)); + mapping->access_count = count; + } + return result; +} + +CUresult +cuinterposer_multicast_export( + void* shareable, CUmemGenericAllocationHandle logical, CUmemAllocationHandleType type, unsigned long long flags) +{ + struct cuinterposer_posix_ticket ticket; + struct multicast_handle* handle = find_handle(logical); + int ticket_fd = -1; + + if (shareable == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (type != CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR || flags != 0) + return CUDA_ERROR_NOT_SUPPORTED; + if (handle == NULL || !handle->live || handle->driver == 0) + return CUDA_ERROR_INVALID_HANDLE; + if (!handle->multicast->creator) + return CUDA_ERROR_NOT_SUPPORTED; + fill_ticket(handle->multicast, &ticket); + if (cuinterposer_posix_create_ticket(&ticket, &ticket_fd) != 0) + return CUDA_ERROR_OUT_OF_MEMORY; + *(int*)shareable = ticket_fd; + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_import( + CUmemGenericAllocationHandle* output, const struct cuinterposer_posix_ticket* ticket, int raw_fd) +{ + import_fn import_handle = (import_fn)cuinterposer_lookup_real_symbol("cuMemImportFromShareableHandle"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast; + CUcontext context; + bool created = false; + bool acquired = true; + CUmemGenericAllocationHandle driver = 0; + CUmemGenericAllocationHandle logical; + CUresult result; + + if (output == NULL || ticket == NULL || raw_fd < 0) + return CUDA_ERROR_INVALID_VALUE; + if (ticket->resource_kind != CUINTERPOSER_RESOURCE_MULTICAST || + ticket->handle_types != CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR) + return CUDA_ERROR_INVALID_HANDLE; + result = import_handle == NULL + ? cuinterposer_unavailable() + : import_handle(&driver, (void*)(uintptr_t)raw_fd, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + if (result != CUDA_SUCCESS) + return result; + multicast = find_multicast(ticket->allocation_id); + if (multicast != NULL && !matches_ticket(multicast, ticket)) { + if (release != NULL) + (void)release(driver); + return CUDA_ERROR_INVALID_HANDLE; + } + if (capture_context(&context) != 0 || (multicast != NULL && multicast->context != context)) { + if (release != NULL) + (void)release(driver); + return CUDA_ERROR_NOT_SUPPORTED; + } + if (multicast != NULL) { + CUmemGenericAllocationHandle existing = current_driver(multicast); + if (existing != 0) { + result = release == NULL ? cuinterposer_unavailable() : release(driver); + if (result != CUDA_SUCCESS) + return result; + driver = existing; + acquired = false; + } + } + if (multicast == NULL) { + multicast = calloc(1, sizeof(*multicast)); + if (multicast != NULL) { + memcpy(multicast->id, ticket->allocation_id, sizeof(multicast->id)); + memcpy(multicast->authorization, ticket->authorization, sizeof(multicast->authorization)); + snprintf( + multicast->creator_participant, sizeof(multicast->creator_participant), "%s", + ticket->creator_participant); + snprintf(multicast->creator_endpoint, sizeof(multicast->creator_endpoint), "%s", ticket->creator_endpoint); + multicast->properties.numDevices = ticket->num_devices; + multicast->properties.size = ticket->allocation_size; + multicast->properties.handleTypes = ticket->handle_types; + multicast->properties.flags = ticket->object_flags; + multicast->effective_size = ticket->allocation_size; + multicast->context = context; + multicast->creator = false; + created = true; + } + } + if (multicast == NULL || add_handle(multicast, driver, &logical) != 0) { + if (acquired && release != NULL) + (void)release(driver); + if (created) + free(multicast); + return CUDA_ERROR_OUT_OF_MEMORY; + } + if (created) { + multicast->next = multicasts; + multicasts = multicast; + } + *output = logical; + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_get_properties(CUmemAllocationProp* properties, CUmemGenericAllocationHandle logical) +{ + properties_fn get_properties = (properties_fn)cuinterposer_lookup_real_symbol("cuMemGetAllocationPropertiesFromHandle"); + struct multicast_handle* handle = find_handle(logical); + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + return get_properties == NULL ? cuinterposer_unavailable() : get_properties(properties, handle->driver); +} + +CUresult +cuinterposer_multicast_create(CUmemGenericAllocationHandle* output, const CUmulticastObjectProp* properties) +{ + create_fn create = (create_fn)cuinterposer_lookup_real_symbol("cuMulticastCreate"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast; + CUmemGenericAllocationHandle driver = 0; + CUmemGenericAllocationHandle logical; + CUresult result; + + if (output == NULL || properties == NULL) + return CUDA_ERROR_INVALID_VALUE; + if (properties->handleTypes != CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR) + return CUDA_ERROR_NOT_SUPPORTED; + /* + * cuMulticastCreate contends the NVIDIA RM top lock and can block for a long + * time under concurrent multi-rank startup. Holding state_lock across it + * starves the single-threaded control server (serve()), which needs the same + * lock to answer a peer's export request; the peer then hits the 30 s + * request_export timeout and gets CUDA_ERROR_INVALID_HANDLE while other ranks + * succeed. That per-rank, timing-dependent failure made 2 of 8 ranks diverge + * onto a different allreduce backend and deadlocked the TP group. + * Same reasoning as bind_memory() above. + */ + operations.release_state_lock(); + result = create == NULL ? cuinterposer_unavailable() : create(&driver, properties); + operations.acquire_state_lock(); + if (result != CUDA_SUCCESS) + return result; + /* A checkpoint may have started while the lock was dropped. */ + if (operations.state_is_active != NULL && !operations.state_is_active()) { + if (release != NULL) + (void)release(driver); + return CUDA_ERROR_NOT_READY; + } + multicast = calloc(1, sizeof(*multicast)); + if (multicast == NULL || random_bytes(multicast->id, sizeof(multicast->id)) != 0 || + random_bytes(multicast->authorization, sizeof(multicast->authorization)) != 0 || + capture_context(&multicast->context) != 0 || add_handle(multicast, driver, &logical) != 0) { + if (release != NULL) + (void)release(driver); + free(multicast); + return CUDA_ERROR_OUT_OF_MEMORY; + } + multicast->properties = *properties; + multicast->effective_size = properties->size; + multicast->creator = true; + snprintf(multicast->creator_participant, sizeof(multicast->creator_participant), "%s", current_participant); + snprintf(multicast->creator_endpoint, sizeof(multicast->creator_endpoint), "%s", current_endpoint); + multicast->next = multicasts; + multicasts = multicast; + *output = logical; + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_add_device(CUmemGenericAllocationHandle logical, CUdevice device) +{ + add_device_fn add_device = (add_device_fn)cuinterposer_lookup_real_symbol("cuMulticastAddDevice"); + struct multicast_handle* handle = find_handle(logical); + struct multicast_device* added = calloc(1, sizeof(*added)); + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + if (added == NULL) + return CUDA_ERROR_OUT_OF_MEMORY; + { + /* + * Like cuMulticastCreate, this is an NVSwitch-team collective that contends + * the RM top lock. Holding state_lock across it starves the control server + * and drives peer export timeouts. Capture the driver handle, drop the + * lock, then revalidate both the phase and the handle on reacquire. + */ + CUmemGenericAllocationHandle driver = handle->driver; + + operations.release_state_lock(); + result = add_device == NULL ? cuinterposer_unavailable() : add_device(driver, device); + operations.acquire_state_lock(); + if (result == CUDA_SUCCESS && operations.state_is_active != NULL && !operations.state_is_active()) + result = CUDA_ERROR_NOT_READY; + if (result == CUDA_SUCCESS) { + /* Re-look up: the handle may have been released while unlocked. */ + handle = find_handle(logical); + if (handle == NULL || !handle->live || handle->driver != driver) + result = CUDA_ERROR_INVALID_HANDLE; + } + } + if (result != CUDA_SUCCESS) { + free(added); + return result; + } + added->device = device; + added->next = handle->multicast->devices; + handle->multicast->devices = added; + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_bind_mem( + CUmemGenericAllocationHandle multicast_logical, CUdevice device, bool device_explicit, size_t multicast_offset, + CUmemGenericAllocationHandle memory, size_t memory_offset, size_t size, unsigned long long flags) +{ + struct multicast_handle* handle = find_handle(multicast_logical); + struct cuinterposer_multicast_member member; + struct multicast_binding* binding; + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + memset(&member, 0, sizeof(member)); + if (operations.member_from_handle == NULL || operations.member_from_handle(memory, &member) != 0) + return CUDA_ERROR_NOT_SUPPORTED; + if (!device_explicit) + device = member.device; + binding = calloc(1, sizeof(*binding)); + if (binding == NULL) + return CUDA_ERROR_OUT_OF_MEMORY; + memcpy(binding->member_id, member.id, sizeof(binding->member_id)); + binding->multicast_offset = multicast_offset; + binding->member_offset = memory_offset; + binding->size = size; + binding->flags = flags; + binding->device = device; + binding->kind = CUINTERPOSER_MULTICAST_BIND_MEM; + binding->api_version = device_explicit ? 2 : 1; + result = bind_memory(handle->driver, binding, member.handle); + if (result != CUDA_SUCCESS) { + free(binding); + return result; + } + if (operations.mark_member_shared != NULL) + operations.mark_member_shared(member.id); + binding->bound = true; + binding->next = handle->multicast->bindings; + handle->multicast->bindings = binding; + observe_extent(handle->multicast, multicast_offset, size); + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_bind_address( + CUmemGenericAllocationHandle multicast_logical, CUdevice device, bool device_explicit, size_t multicast_offset, CUdeviceptr memory, size_t size, unsigned long long flags) { - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle, CUdevice, size_t, CUdeviceptr, size_t, - unsigned long long); - function_type function = LOOKUP(cuMulticastBindAddr_v2, function_type); - return function != NULL - ? function(multicast, device, multicast_offset, memory, size, flags) - : cuinterposer_unavailable(); + struct multicast_handle* handle = find_handle(multicast_logical); + struct cuinterposer_multicast_member member; + struct multicast_binding* binding; + bool tracked_member = false; + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + memset(&member, 0, sizeof(member)); + if (operations.member_from_address == NULL || operations.member_from_address(memory, size, &member) != 0) { + if (random_bytes(member.id, sizeof(member.id)) != 0 || + (!device_explicit && capture_device(&device) != 0)) + return CUDA_ERROR_NOT_SUPPORTED; + member.address = memory; + } else { + tracked_member = true; + if (!device_explicit) + device = member.device; + } + binding = calloc(1, sizeof(*binding)); + if (binding == NULL) + return CUDA_ERROR_OUT_OF_MEMORY; + memcpy(binding->member_id, member.id, sizeof(binding->member_id)); + binding->member_address = memory; + binding->multicast_offset = multicast_offset; + binding->member_offset = member.allocation_offset; + binding->size = size; + binding->flags = flags; + binding->device = device; + binding->kind = CUINTERPOSER_MULTICAST_BIND_ADDR; + binding->api_version = device_explicit ? 2 : 1; + result = bind_address(handle->driver, binding); + if (result != CUDA_SUCCESS) { + free(binding); + return result; + } + if (tracked_member && operations.mark_member_shared != NULL) + operations.mark_member_shared(member.id); + binding->bound = true; + binding->next = handle->multicast->bindings; + handle->multicast->bindings = binding; + observe_extent(handle->multicast, multicast_offset, size); + return CUDA_SUCCESS; +} + +CUresult +cuinterposer_multicast_unbind(CUmemGenericAllocationHandle logical, CUdevice device, size_t offset, size_t size) +{ + unbind_fn unbind = (unbind_fn)cuinterposer_lookup_real_symbol("cuMulticastUnbind"); + struct multicast_handle* handle = find_handle(logical); + struct multicast_binding* binding; + CUresult result; + + if (handle == NULL || !handle->live) + return CUDA_ERROR_INVALID_HANDLE; + if (handle->driver == 0) + return CUDA_ERROR_NOT_READY; + for (binding = handle->multicast->bindings; binding != NULL; binding = binding->next) { + if (binding->bound && binding->device == device && binding->multicast_offset == offset && binding->size == size) + break; + } + if (binding == NULL) + return CUDA_ERROR_NOT_SUPPORTED; + result = unbind == NULL ? cuinterposer_unavailable() : unbind(handle->driver, device, offset, size); + if (result != CUDA_SUCCESS) + return result; + binding->bound = false; + return CUDA_SUCCESS; } -#endif -CUresult CUDAAPI -cuMulticastGetGranularity( - size_t* granularity, const CUmulticastObjectProp* properties, - CUmulticastGranularity_flags option) +size_t +cuinterposer_multicast_record_count(void) { - typedef CUresult(CUDAAPI * function_type)( - size_t*, const CUmulticastObjectProp*, CUmulticastGranularity_flags); - function_type function = LOOKUP(cuMulticastGetGranularity, function_type); - return function != NULL ? function(granularity, properties, option) - : cuinterposer_unavailable(); + const struct multicast* multicast; + size_t count = 0; + + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + const struct multicast_binding* binding; + const struct multicast_device* device; + const struct multicast_mapping* mapping; + if (!active(multicast)) + continue; + count++; + for (device = multicast->devices; device != NULL; device = device->next) count++; + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + if (binding->bound) + count++; + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped) + count++; + } + } + return count; } -CUresult CUDAAPI -cuMulticastUnbind( - CUmemGenericAllocationHandle multicast, CUdevice device, size_t offset, size_t size) +int +cuinterposer_multicast_write_records(struct cuinterposer_record* records, size_t count) { - typedef CUresult(CUDAAPI * function_type)( - CUmemGenericAllocationHandle, CUdevice, size_t, size_t); - function_type function = LOOKUP(cuMulticastUnbind, function_type); - return function != NULL ? function(multicast, device, offset, size) - : cuinterposer_unavailable(); + const struct multicast* multicast; + size_t written = 0; + + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + const struct multicast_binding* binding; + const struct multicast_device* device; + const struct multicast_mapping* mapping; + struct cuinterposer_record* record; + size_t index; + + if (!active(multicast)) + continue; + if (written == count) + return -1; + record = &records[written++]; + record->kind = CUINTERPOSER_MULTICAST; + record->flags = multicast->creator ? CUINTERPOSER_CREATOR : 0; + if (live_handle_count(multicast) != 0) + record->flags |= CUINTERPOSER_APPLICATION_HANDLE_LIVE; + memcpy(record->allocation_id, multicast->id, sizeof(record->allocation_id)); + record->allocation_size = multicast->effective_size; + record->application_handle_count = (uint32_t)live_handle_count(multicast); + record->handle_types = multicast->properties.handleTypes; + record->object_flags = multicast->properties.flags; + record->num_devices = multicast->properties.numDevices; + snprintf(record->creator_participant, sizeof(record->creator_participant), "%s", multicast->creator_participant); + for (device = multicast->devices; device != NULL; device = device->next) { + if (written == count) + return -1; + record = &records[written++]; + record->kind = CUINTERPOSER_MULTICAST_DEVICE; + memcpy(record->allocation_id, multicast->id, sizeof(record->allocation_id)); + record->device = device->device; + } + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + if (!binding->bound) + continue; + if (written == count) + return -1; + record = &records[written++]; + record->kind = CUINTERPOSER_MULTICAST_BINDING; + memcpy(record->allocation_id, multicast->id, sizeof(record->allocation_id)); + memcpy(record->member_id, binding->member_id, sizeof(record->member_id)); + record->address = binding->member_address; + record->size = binding->size; + record->offset = binding->multicast_offset; + record->member_offset = binding->member_offset; + record->operation_flags = binding->flags; + record->binding_kind = binding->kind; + record->api_version = binding->api_version; + record->device = binding->device; + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (!mapping->mapped) + continue; + if (written == count) + return -1; + record = &records[written++]; + record->kind = CUINTERPOSER_MULTICAST_MAPPING; + memcpy(record->allocation_id, multicast->id, sizeof(record->allocation_id)); + record->address = mapping->address; + record->size = mapping->size; + record->offset = mapping->offset; + record->operation_flags = mapping->flags; + record->access_count = (uint32_t)mapping->access_count; + for (index = 0; index < mapping->access_count; index++) { + record->access[index].location_type = mapping->access[index].location.type; + record->access[index].location_id = mapping->access[index].location.id; + record->access[index].flags = mapping->access[index].flags; + } + } + } + return written == count ? 0 : -1; } -#undef LOOKUP +int +cuinterposer_multicast_prepare(void) +{ + unmap_fn unmap = (unmap_fn)cuinterposer_lookup_real_symbol("cuMemUnmap"); + unbind_fn unbind = (unbind_fn)cuinterposer_lookup_real_symbol("cuMulticastUnbind"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast; + + failure[0] = '\0'; + if (unmap == NULL || unbind == NULL || release == NULL) + return fail("multicast teardown symbols are unavailable", CUDA_SUCCESS); + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + if (active(multicast) && current_driver(multicast) == 0) { + struct multicast_mapping* mapping; + bool retainable = false; + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped) { + retainable = true; + break; + } + } + if (!retainable) + return fail("multicast object has no handle or mapping", CUDA_SUCCESS); + } + } + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct multicast_binding* binding; + struct multicast_handle* handle; + struct multicast_mapping* mapping; + struct context_scope scope; + CUmemGenericAllocationHandle driver; + bool temporary = false; + + if (!active(multicast)) + continue; + driver = current_driver(multicast); + if (enter_context(multicast->context, &scope) != 0) + return fail("cannot enter multicast context", CUDA_SUCCESS); + if (driver == 0) { + retain_fn retain = (retain_fn)cuinterposer_lookup_real_symbol("cuMemRetainAllocationHandle"); + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped) + break; + } + if (retain == NULL || mapping == NULL || retain(&driver, (void*)(uintptr_t)mapping->address) != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cannot retain multicast teardown handle", CUDA_SUCCESS); + } + temporary = true; + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + CUresult result; + if (!mapping->mapped) + continue; + result = unmap(mapping->address, mapping->size); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMemUnmap multicast mapping", result); + } + mapping->mapped = false; + mapping->checkpointed = true; + } + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + CUresult result; + if (!binding->bound) + continue; + result = unbind(driver, binding->device, binding->multicast_offset, binding->size); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMulticastUnbind", result); + } + binding->bound = false; + binding->checkpointed = true; + } + for (handle = handles; handle != NULL; handle = handle->next) { + CUmemGenericAllocationHandle old; + CUresult result; + if (!handle->live || handle->multicast != multicast) + continue; + old = handle->driver; + if (!driver_used(handle, old)) { + result = release(old); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMemRelease multicast handle", result); + } + } + handle->driver = 0; + } + if (temporary) { + CUresult result = release(driver); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMemRelease multicast teardown handle", result); + } + } + multicast->checkpointed = true; + if (leave_context(&scope) != 0) + return fail("cannot leave multicast context", CUDA_SUCCESS); + } + return 0; +} + +int +cuinterposer_multicast_restore_creators(void) +{ + create_fn create = (create_fn)cuinterposer_lookup_real_symbol("cuMulticastCreate"); + struct multicast* multicast; + + failure[0] = '\0'; + if (create == NULL) + return fail("cuMulticastCreate is unavailable", CUDA_SUCCESS); + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct context_scope scope; + CUmemGenericAllocationHandle driver = 0; + CUresult result; + if (!multicast->checkpointed || !multicast->creator) + continue; + if (enter_context(multicast->context, &scope) != 0) + return fail("cannot enter multicast creator context", CUDA_SUCCESS); + result = create(&driver, &multicast->properties); + if (result == CUDA_SUCCESS) + install_driver(multicast, driver); + if (leave_context(&scope) != 0) + return fail("cannot leave multicast creator context", CUDA_SUCCESS); + if (result != CUDA_SUCCESS) + return fail("cuMulticastCreate", result); + } + return 0; +} + +int +cuinterposer_multicast_restore_importers(void) +{ + import_fn import_handle = (import_fn)cuinterposer_lookup_real_symbol("cuMemImportFromShareableHandle"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast; + + failure[0] = '\0'; + if (import_handle == NULL || release == NULL) + return fail("multicast import symbols are unavailable", CUDA_SUCCESS); + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct context_scope scope; + CUmemGenericAllocationHandle driver = current_driver(multicast); + CUresult result = CUDA_SUCCESS; + + if (!multicast->checkpointed) + continue; + if (enter_context(multicast->context, &scope) != 0) + return fail("cannot enter multicast importer context", CUDA_SUCCESS); + if (!multicast->creator) { + struct cuinterposer_posix_ticket ticket; + char export_error[sizeof(failure)]; + int raw_fd = -1; + + fill_ticket(multicast, &ticket); + if (cuinterposer_posix_request_export(&ticket, &raw_fd, export_error, sizeof(export_error)) != 0) { + (void)leave_context(&scope); + snprintf( + failure, sizeof(failure), "multicast creator export: %.96s", + export_error[0] == '\0' ? "request failed" : export_error); + return -1; + } + result = import_handle(&driver, (void*)(uintptr_t)raw_fd, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + if (close(raw_fd) != 0 && result == CUDA_SUCCESS) { + (void)release(driver); + result = CUDA_ERROR_UNKNOWN; + } + if (result == CUDA_SUCCESS) + install_driver(multicast, driver); + } + if (leave_context(&scope) != 0) + return fail("cannot leave multicast importer context", CUDA_SUCCESS); + if (result != CUDA_SUCCESS) + return fail("multicast import", result); + } + return 0; +} + +int +cuinterposer_multicast_restore_devices(void) +{ + add_device_fn add_device = (add_device_fn)cuinterposer_lookup_real_symbol("cuMulticastAddDevice"); + struct multicast* multicast; + + failure[0] = '\0'; + if (add_device == NULL) + return fail("cuMulticastAddDevice is unavailable", CUDA_SUCCESS); + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct multicast_device* device; + struct context_scope scope; + CUmemGenericAllocationHandle driver; + + if (!multicast->checkpointed) + continue; + driver = current_driver(multicast); + if (driver == 0) + return fail("imported multicast handle is unavailable", CUDA_SUCCESS); + if (enter_context(multicast->context, &scope) != 0) + return fail("cannot enter multicast device context", CUDA_SUCCESS); + for (device = multicast->devices; device != NULL; device = device->next) { + CUresult result = add_device(driver, device->device); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMulticastAddDevice", result); + } + } + if (leave_context(&scope) != 0) + return fail("cannot leave multicast device context", CUDA_SUCCESS); + } + return 0; +} + +int +cuinterposer_multicast_restore_topology(void) +{ + map_fn map = (map_fn)cuinterposer_lookup_real_symbol("cuMemMap"); + access_fn set_access = (access_fn)cuinterposer_lookup_real_symbol("cuMemSetAccess"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast; + + failure[0] = '\0'; + if (map == NULL || set_access == NULL || release == NULL) + return fail("multicast mapping symbols are unavailable", CUDA_SUCCESS); + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + struct multicast_binding* binding; + struct multicast_mapping* mapping; + struct context_scope scope; + CUmemGenericAllocationHandle driver; + + if (!multicast->checkpointed) + continue; + driver = current_driver(multicast); + if (driver == 0) + return fail("restored multicast handle is unavailable", CUDA_SUCCESS); + if (enter_context(multicast->context, &scope) != 0) + return fail("cannot enter multicast restore context", CUDA_SUCCESS); + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + struct cuinterposer_multicast_member member; + CUresult result; + + if (!binding->checkpointed) + continue; + if (binding->kind == CUINTERPOSER_MULTICAST_BIND_MEM) { + memset(&member, 0, sizeof(member)); + if (operations.member_from_id == NULL || operations.member_from_id(binding->member_id, &member) != 0) { + (void)leave_context(&scope); + return fail("restored multicast member is unavailable", CUDA_SUCCESS); + } + result = bind_memory(driver, binding, member.handle); + if (member.temporary_handle) { + CUresult release_result = release(member.handle); + if (result == CUDA_SUCCESS) + result = release_result; + } + } else { + result = bind_address(driver, binding); + } + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("restore multicast binding", result); + } + binding->bound = true; + binding->checkpointed = false; + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + CUresult result; + if (!mapping->checkpointed) + continue; + result = map(mapping->address, mapping->size, mapping->offset, driver, mapping->flags); + if (result == CUDA_SUCCESS && mapping->access_count != 0) + result = set_access(mapping->address, mapping->size, mapping->access, mapping->access_count); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("restore multicast mapping or access", result); + } + mapping->mapped = true; + mapping->checkpointed = false; + } + multicast->checkpointed = false; + multicast->restore_handle = 0; + if (live_handle_count(multicast) == 0) { + CUresult result = release(driver); + if (result != CUDA_SUCCESS) { + (void)leave_context(&scope); + return fail("cuMemRelease restored multicast handle", result); + } + } + if (leave_context(&scope) != 0) + return fail("cannot leave multicast restore context", CUDA_SUCCESS); + } + return cuinterposer_multicast_validate_restored(); +} + +int +cuinterposer_multicast_validate_restored(void) +{ + const struct multicast_handle* handle; + const struct multicast* multicast; + + for (handle = handles; handle != NULL; handle = handle->next) { + if (handle->live && handle->driver == 0) + return fail("multicast logical handle was not restored", CUDA_SUCCESS); + } + for (multicast = multicasts; multicast != NULL; multicast = multicast->next) { + const struct multicast_binding* binding; + const struct multicast_mapping* mapping; + if (multicast->checkpointed) + return fail("multicast object remains checkpointed", CUDA_SUCCESS); + for (binding = multicast->bindings; binding != NULL; binding = binding->next) { + if (binding->checkpointed) + return fail("multicast binding remains checkpointed", CUDA_SUCCESS); + } + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->checkpointed) + return fail("multicast mapping remains checkpointed", CUDA_SUCCESS); + } + } + return 0; +} + +CUresult +cuinterposer_multicast_export_raw( + const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE], const uint8_t authorization[CUINTERPOSER_TOKEN_SIZE], + int* output) +{ + export_fn export_handle = (export_fn)cuinterposer_lookup_real_symbol("cuMemExportToShareableHandle"); + retain_fn retain = (retain_fn)cuinterposer_lookup_real_symbol("cuMemRetainAllocationHandle"); + release_fn release = (release_fn)cuinterposer_lookup_real_symbol("cuMemRelease"); + struct multicast* multicast = find_multicast(id); + struct multicast_mapping* mapping; + struct context_scope scope; + CUmemGenericAllocationHandle driver; + bool temporary = false; + CUresult result; + + *output = -1; + if (multicast == NULL || !multicast->creator || + memcmp(multicast->authorization, authorization, sizeof(multicast->authorization)) != 0) + return CUDA_ERROR_INVALID_HANDLE; + driver = current_driver(multicast); + if (export_handle == NULL || enter_context(multicast->context, &scope) != 0) + return CUDA_ERROR_INVALID_HANDLE; + if (driver == 0) { + for (mapping = multicast->mappings; mapping != NULL; mapping = mapping->next) { + if (mapping->mapped) + break; + } + if (mapping == NULL || retain == NULL || retain(&driver, (void*)(uintptr_t)mapping->address) != CUDA_SUCCESS) { + (void)leave_context(&scope); + return CUDA_ERROR_INVALID_HANDLE; + } + temporary = true; + } + result = export_handle(output, driver, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + if (temporary) { + CUresult release_result = release == NULL ? CUDA_ERROR_NOT_INITIALIZED : release(driver); + if (result == CUDA_SUCCESS) + result = release_result; + } + if (result != CUDA_SUCCESS && *output >= 0) { + close(*output); + *output = -1; + } + if (leave_context(&scope) != 0) { + if (*output >= 0) { + close(*output); + *output = -1; + } + return CUDA_ERROR_UNKNOWN; + } + return result; +} + +const char* +cuinterposer_multicast_error(void) +{ + return failure[0] == '\0' ? "multicast operation failed" : failure; +} diff --git a/agent/cmd/cuinterpose/multicast.h b/agent/cmd/cuinterpose/multicast.h index 390e13f1..ef530e08 100644 --- a/agent/cmd/cuinterpose/multicast.h +++ b/agent/cmd/cuinterpose/multicast.h @@ -7,4 +7,79 @@ #ifndef CUINTERPOSER_MULTICAST_H #define CUINTERPOSER_MULTICAST_H +#include +#include +#include +#include + +#include "posix.h" +#include "protocol.h" + +struct cuinterposer_multicast_member { + uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]; + CUmemGenericAllocationHandle handle; + CUdeviceptr address; + size_t allocation_offset; + CUdevice device; + bool temporary_handle; +}; + +struct cuinterposer_multicast_callbacks { + int (*allocate_logical_handle)(CUmemGenericAllocationHandle* output); + int (*member_from_handle)(CUmemGenericAllocationHandle logical, struct cuinterposer_multicast_member* member); + int (*member_from_address)(CUdeviceptr address, size_t size, struct cuinterposer_multicast_member* member); + int (*member_from_id)(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE], struct cuinterposer_multicast_member* member); + void (*mark_member_shared)(const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE]); + void (*release_state_lock)(void); + void (*acquire_state_lock)(void); + /* + * True while the interposer is in PHASE_ACTIVE. Callers that drop + * state_lock across a long driver call must revalidate with this on + * reacquire: a checkpoint may have started in the window. + */ + bool (*state_is_active)(void); +}; + +void cuinterposer_multicast_initialize( + const struct cuinterposer_multicast_callbacks* callbacks, const char* participant_id, const char* endpoint); +void cuinterposer_multicast_reset(void); + +bool cuinterposer_multicast_is_handle(CUmemGenericAllocationHandle logical); +bool cuinterposer_multicast_has_mapping(CUdeviceptr address, size_t size); +CUresult cuinterposer_multicast_release(CUmemGenericAllocationHandle logical); +CUresult cuinterposer_multicast_retain(CUmemGenericAllocationHandle* output, void* address); +CUresult cuinterposer_multicast_map( + CUdeviceptr address, size_t size, size_t offset, CUmemGenericAllocationHandle logical, unsigned long long flags); +CUresult cuinterposer_multicast_unmap(CUdeviceptr address, size_t size); +CUresult cuinterposer_multicast_set_access( + CUdeviceptr address, size_t size, const CUmemAccessDesc* descriptors, size_t count); +CUresult cuinterposer_multicast_export( + void* shareable, CUmemGenericAllocationHandle logical, CUmemAllocationHandleType type, unsigned long long flags); +CUresult cuinterposer_multicast_import( + CUmemGenericAllocationHandle* output, const struct cuinterposer_posix_ticket* ticket, int raw_fd); +CUresult cuinterposer_multicast_get_properties(CUmemAllocationProp* properties, CUmemGenericAllocationHandle logical); + +CUresult cuinterposer_multicast_create(CUmemGenericAllocationHandle* output, const CUmulticastObjectProp* properties); +CUresult cuinterposer_multicast_add_device(CUmemGenericAllocationHandle logical, CUdevice device); +CUresult cuinterposer_multicast_bind_mem( + CUmemGenericAllocationHandle multicast, CUdevice device, bool device_explicit, size_t multicast_offset, + CUmemGenericAllocationHandle memory, size_t memory_offset, size_t size, unsigned long long flags); +CUresult cuinterposer_multicast_bind_address( + CUmemGenericAllocationHandle multicast, CUdevice device, bool device_explicit, size_t multicast_offset, + CUdeviceptr memory, size_t size, unsigned long long flags); +CUresult cuinterposer_multicast_unbind(CUmemGenericAllocationHandle multicast, CUdevice device, size_t offset, size_t size); + +size_t cuinterposer_multicast_record_count(void); +int cuinterposer_multicast_write_records(struct cuinterposer_record* records, size_t count); +int cuinterposer_multicast_prepare(void); +int cuinterposer_multicast_restore_creators(void); +int cuinterposer_multicast_restore_importers(void); +int cuinterposer_multicast_restore_devices(void); +int cuinterposer_multicast_restore_topology(void); +int cuinterposer_multicast_validate_restored(void); +CUresult cuinterposer_multicast_export_raw( + const uint8_t id[CUINTERPOSER_ALLOCATION_ID_SIZE], const uint8_t authorization[CUINTERPOSER_TOKEN_SIZE], + int* output); +const char* cuinterposer_multicast_error(void); + #endif diff --git a/agent/cmd/cuinterpose/posix.c b/agent/cmd/cuinterpose/posix.c index c2a4360a..0bbf273c 100644 --- a/agent/cmd/cuinterpose/posix.c +++ b/agent/cmd/cuinterpose/posix.c @@ -17,9 +17,33 @@ #include #include +#include + #include "util.h" -#define EXPORT_TIMEOUT_SECONDS 30 +/* + * Deadline for contacting the creating process to export a shared handle. + * + * This was a hard-coded 30 s. Under concurrent multi-rank startup the creator's + * control server can be blocked far longer than that, and the timeout surfaces + * as a per-rank CUDA_ERROR_INVALID_HANDLE -- a non-deterministic, rank-divergent + * failure that frameworks are not built to handle. Default generously and allow + * an operator override. + */ +#define EXPORT_TIMEOUT_SECONDS_DEFAULT 300 + +static unsigned +export_timeout_seconds(void) +{ + static unsigned cached; + + if (cached == 0) { + const char* value = getenv("DYN_SNAPSHOT_EXPORT_TIMEOUT_SECONDS"); + + cached = bounded_seconds(value, EXPORT_TIMEOUT_SECONDS_DEFAULT); + } + return cached; +} static bool zero_bytes(const void* value, size_t size) @@ -69,8 +93,19 @@ cuinterposer_posix_read_ticket(int fd, struct cuinterposer_posix_ticket* ticket) ticket->creator_endpoint[0] != '/' || memchr(ticket->creator_endpoint, '\0', sizeof(ticket->creator_endpoint)) == NULL || zero_bytes(ticket->authorization, sizeof(ticket->authorization)) || + !zero_bytes(ticket->reserved_alignment, sizeof(ticket->reserved_alignment)) || + (ticket->resource_kind != CUINTERPOSER_RESOURCE_UNICAST && + ticket->resource_kind != CUINTERPOSER_RESOURCE_MULTICAST) || !zero_bytes(ticket->reserved_identity, sizeof(ticket->reserved_identity))) return -1; + if (ticket->resource_kind == CUINTERPOSER_RESOURCE_UNICAST && + (ticket->num_devices != 0 || ticket->allocation_size != 0 || ticket->handle_types != 0 || + ticket->object_flags != 0)) + return -1; + if (ticket->resource_kind == CUINTERPOSER_RESOURCE_MULTICAST && + (ticket->num_devices == 0 || ticket->allocation_size == 0 || + ticket->handle_types != CUINTERPOSER_POSIX_HANDLE_TYPE)) + return -1; return 0; } @@ -91,12 +126,13 @@ cuinterposer_posix_request_export( request.magic = CUINTERPOSER_MAGIC; request.version = CUINTERPOSER_VERSION; request.operation = CUINTERPOSER_EXPORT; + request.resource_kind = ticket->resource_kind; snprintf(request.participant_id, sizeof(request.participant_id), "%s", ticket->creator_participant); memcpy(request.authorization, ticket->authorization, sizeof(request.authorization)); memcpy(request.allocation_id, ticket->allocation_id, sizeof(request.allocation_id)); snprintf(address.sun_path, sizeof(address.sun_path), "%s", ticket->creator_endpoint); client = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); - if (client < 0 || set_socket_timeouts(client, EXPORT_TIMEOUT_SECONDS) != 0 || + if (client < 0 || set_socket_timeouts(client, export_timeout_seconds()) != 0 || connect(client, (const struct sockaddr*)&address, sizeof(address)) != 0 || send_header(client, &request, -1) != 0 || receive_header(client, &response, output) != 0) { @@ -106,7 +142,8 @@ cuinterposer_posix_request_export( } if (!header_strings_terminated(&response) || response.magic != CUINTERPOSER_MAGIC || response.version != CUINTERPOSER_VERSION || response.operation != CUINTERPOSER_EXPORT || response.count != 0 || - response.payload_size != 0 || strcmp(response.participant_id, ticket->creator_participant) != 0) { + response.payload_size != 0 || strcmp(response.participant_id, ticket->creator_participant) != 0 || + response.resource_kind != ticket->resource_kind) { if (error != NULL && error_size != 0) snprintf(error, error_size, "%s", "invalid creator export response"); if (*output >= 0) { diff --git a/agent/cmd/cuinterpose/posix.h b/agent/cmd/cuinterpose/posix.h index afc94e3e..99681279 100644 --- a/agent/cmd/cuinterpose/posix.h +++ b/agent/cmd/cuinterpose/posix.h @@ -14,7 +14,7 @@ #include "protocol.h" #define CUINTERPOSER_POSIX_TICKET_MAGIC 0x44564d43U -#define CUINTERPOSER_POSIX_TICKET_VERSION 1U +#define CUINTERPOSER_POSIX_TICKET_VERSION 2U struct cuinterposer_posix_ticket { uint32_t magic; @@ -24,7 +24,13 @@ struct cuinterposer_posix_ticket { uint8_t allocation_id[CUINTERPOSER_ALLOCATION_ID_SIZE]; char creator_endpoint[sizeof(((struct sockaddr_un*)0)->sun_path)]; uint8_t authorization[CUINTERPOSER_TOKEN_SIZE]; - uint8_t reserved_identity[42]; + uint8_t reserved_alignment[2]; + uint32_t resource_kind; + uint32_t num_devices; + uint64_t allocation_size; + uint64_t handle_types; + uint64_t object_flags; + uint8_t reserved_identity[8]; }; _Static_assert(sizeof(struct cuinterposer_posix_ticket) == 256, "cuinterposer POSIX ticket layout changed"); diff --git a/agent/cmd/cuinterpose/protocol.h b/agent/cmd/cuinterpose/protocol.h index 7cb10421..ebb87ece 100644 --- a/agent/cmd/cuinterpose/protocol.h +++ b/agent/cmd/cuinterpose/protocol.h @@ -10,7 +10,7 @@ #include #define CUINTERPOSER_MAGIC 0x44564d4dU -#define CUINTERPOSER_VERSION 1U +#define CUINTERPOSER_VERSION 2U #define CUINTERPOSER_SOCKET_PREFIX "cuinterposer-" #define CUINTERPOSER_ID_SIZE 33U #define CUINTERPOSER_ALLOCATION_ID_SIZE 16U @@ -22,15 +22,24 @@ enum cuinterposer_operation { CUINTERPOSER_IDENTIFY = 1, CUINTERPOSER_INSPECT = 2, - CUINTERPOSER_PREPARE = 3, - CUINTERPOSER_RESTORE_CREATORS = 4, - CUINTERPOSER_RESTORE_IMPORTERS = 5, - CUINTERPOSER_EXPORT = 6, + CUINTERPOSER_PREPARE_MULTICAST = 3, + CUINTERPOSER_PREPARE = 4, + CUINTERPOSER_RESTORE_CREATORS = 5, + CUINTERPOSER_RESTORE_IMPORTERS = 6, + CUINTERPOSER_EXPORT = 7, + CUINTERPOSER_RESTORE_MULTICAST_CREATORS = 8, + CUINTERPOSER_RESTORE_MULTICAST_IMPORTERS = 9, + CUINTERPOSER_RESTORE_MULTICAST_DEVICES = 10, + CUINTERPOSER_RESTORE_MULTICAST = 11, }; enum cuinterposer_record_kind { CUINTERPOSER_ALLOCATION = 1, CUINTERPOSER_MAPPING = 2, + CUINTERPOSER_MULTICAST = 3, + CUINTERPOSER_MULTICAST_DEVICE = 4, + CUINTERPOSER_MULTICAST_BINDING = 5, + CUINTERPOSER_MULTICAST_MAPPING = 6, }; enum cuinterposer_record_flags { @@ -38,6 +47,16 @@ enum cuinterposer_record_flags { CUINTERPOSER_APPLICATION_HANDLE_LIVE = 1U << 1, }; +enum cuinterposer_resource_kind { + CUINTERPOSER_RESOURCE_UNICAST = 1, + CUINTERPOSER_RESOURCE_MULTICAST = 2, +}; + +enum cuinterposer_multicast_binding_kind { + CUINTERPOSER_MULTICAST_BIND_MEM = 1, + CUINTERPOSER_MULTICAST_BIND_ADDR = 2, +}; + struct cuinterposer_header { uint32_t magic; uint16_t version; @@ -49,7 +68,8 @@ struct cuinterposer_header { char message[96]; uint8_t authorization[CUINTERPOSER_TOKEN_SIZE]; uint8_t allocation_id[CUINTERPOSER_ALLOCATION_ID_SIZE]; - uint8_t reserved[71]; + uint32_t resource_kind; + uint8_t reserved[64]; }; struct cuinterposer_access { @@ -73,10 +93,30 @@ struct cuinterposer_record { uint32_t access_count; uint32_t application_handle_count; struct cuinterposer_access access[CUINTERPOSER_MAX_ACCESS]; + uint8_t member_id[CUINTERPOSER_ALLOCATION_ID_SIZE]; + char creator_participant[CUINTERPOSER_ID_SIZE]; + uint8_t binding_kind; + uint8_t api_version; + uint8_t reserved[5]; + uint64_t member_offset; + uint64_t operation_flags; + uint64_t handle_types; + uint64_t object_flags; + uint32_t num_devices; + int32_t device; }; _Static_assert(sizeof(struct cuinterposer_header) == 256, "cuinterposer header layout changed"); _Static_assert(sizeof(struct cuinterposer_access) == 16, "cuinterposer access layout changed"); -_Static_assert(sizeof(struct cuinterposer_record) == 208, "cuinterposer record layout changed"); +_Static_assert(sizeof(struct cuinterposer_record) == 304, "cuinterposer record layout changed"); +_Static_assert( + CUINTERPOSER_ALLOCATION < CUINTERPOSER_MULTICAST_BINDING, "unicast members must sort before multicast bindings"); +_Static_assert( + CUINTERPOSER_MULTICAST < CUINTERPOSER_MULTICAST_DEVICE, "multicast objects must sort before their devices"); +_Static_assert( + CUINTERPOSER_MULTICAST_DEVICE < CUINTERPOSER_MULTICAST_BINDING, + "multicast devices must sort before their bindings"); +_Static_assert( + CUINTERPOSER_MULTICAST < CUINTERPOSER_MULTICAST_MAPPING, "multicast objects must sort before their mappings"); #endif diff --git a/agent/cmd/cuinterpose/tests/fake_cuda.c b/agent/cmd/cuinterpose/tests/fake_cuda.c index 5d9a208e..e4e8c635 100644 --- a/agent/cmd/cuinterpose/tests/fake_cuda.c +++ b/agent/cmd/cuinterpose/tests/fake_cuda.c @@ -8,6 +8,8 @@ #include #include +#include +#include #include #undef cuGetProcAddress @@ -32,14 +34,48 @@ enum { result_multicast_unbind, }; +static pthread_mutex_t multicast_map_lock = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t multicast_map_condition = PTHREAD_COND_INITIALIZER; +static int block_multicast_map; +static int multicast_map_entered; + +void +fakeEnableBlockingMulticastMap(void) +{ + pthread_mutex_lock(&multicast_map_lock); + block_multicast_map = 1; + multicast_map_entered = 0; + pthread_mutex_unlock(&multicast_map_lock); +} + +int +fakeMulticastMapEntered(void) +{ + int entered; + + pthread_mutex_lock(&multicast_map_lock); + entered = multicast_map_entered; + pthread_mutex_unlock(&multicast_map_lock); + return entered; +} + +void +fakeReleaseMulticastMap(void) +{ + pthread_mutex_lock(&multicast_map_lock); + block_multicast_map = 0; + pthread_cond_broadcast(&multicast_map_condition); + pthread_mutex_unlock(&multicast_map_lock); +} + CUresult CUDAAPI fakeCuMemCreateOriginal( CUmemGenericAllocationHandle* output, size_t size, const CUmemAllocationProp* properties, unsigned long long flags) { (void)size; - (void)properties; (void)flags; + (void)properties; *output = 0xabc; return (CUresult)result_create; } @@ -75,8 +111,16 @@ cuMemMap( (void)address; (void)size; (void)offset; - (void)handle; (void)flags; + if (handle == 0x456) { + pthread_mutex_lock(&multicast_map_lock); + multicast_map_entered = 1; + pthread_cond_broadcast(&multicast_map_condition); + while (block_multicast_map) + pthread_cond_wait(&multicast_map_condition, &multicast_map_lock); + pthread_mutex_unlock(&multicast_map_lock); + return CUDA_SUCCESS; + } return (CUresult)result_map; } @@ -135,7 +179,7 @@ fakeCuMulticastCreateOriginal( { (void)properties; *output = 0x456; - return (CUresult)result_multicast_create; + return CUDA_SUCCESS; } CUresult CUDAAPI @@ -203,6 +247,13 @@ cuMulticastUnbind( return (CUresult)result_multicast_unbind; } +CUresult CUDAAPI +cuCtxGetCurrent(CUcontext* context) +{ + *context = (CUcontext)(uintptr_t)1; + return CUDA_SUCCESS; +} + static void* original(const char* symbol) { diff --git a/agent/cmd/cuinterpose/tests/multicast_forward_test.c b/agent/cmd/cuinterpose/tests/multicast_forward_test.c index 5d1f6f90..d62ffce4 100644 --- a/agent/cmd/cuinterpose/tests/multicast_forward_test.c +++ b/agent/cmd/cuinterpose/tests/multicast_forward_test.c @@ -9,8 +9,13 @@ #include #include #include +#include +#include #include #include +#include + +#include "../protocol.h" #undef cuGetProcAddress #undef cuMulticastBindAddr @@ -20,6 +25,22 @@ CUresult CUDAAPI cuGetProcAddress(const char*, void**, int, cuuint64_t); extern CUresult CUDAAPI fakeCuMulticastCreateOriginal( CUmemGenericAllocationHandle*, const CUmulticastObjectProp*); +extern void fakeEnableBlockingMulticastMap(void); +extern int fakeMulticastMapEntered(void); +extern void fakeReleaseMulticastMap(void); + +struct map_call { + CUmemGenericAllocationHandle handle; + CUresult result; +}; + +struct add_device_call { + CUmemGenericAllocationHandle handle; + CUresult result; + int done; +}; + +static CUmulticastObjectProp properties(void); static void require(int condition, const char* message) @@ -31,26 +52,121 @@ require(int condition, const char* message) } static void -test_direct_forwarding(void) +test_multicast_records_driver_accepted_extent(void) +{ + typedef size_t (*record_count_fn)(void); + typedef int (*write_records_fn)(struct cuinterposer_record*, size_t); + record_count_fn record_count = (record_count_fn)dlsym(RTLD_DEFAULT, "cuinterposer_multicast_record_count"); + write_records_fn write_records = + (write_records_fn)dlsym(RTLD_DEFAULT, "cuinterposer_multicast_write_records"); + CUmulticastObjectProp props = properties(); + struct cuinterposer_record* records; + CUmemGenericAllocationHandle handle; + size_t count; + size_t index; + int found = 0; + + props.size = 2048; + require(record_count != NULL && write_records != NULL, "multicast record API"); + require(cuMulticastCreate(&handle, &props) == CUDA_SUCCESS, "create multicast with sub-extent request"); + require(cuMemMap(0x3000, 8192, 0, handle, 0) == CUDA_SUCCESS, "driver accepts larger multicast mapping"); + count = record_count(); + records = calloc(count, sizeof(*records)); + require(records != NULL, "allocate multicast records"); + require(write_records(records, count) == 0, "write multicast records"); + for (index = 0; index < count; index++) { + if (records[index].kind == CUINTERPOSER_MULTICAST && records[index].allocation_size == 8192) + found = 1; + } + free(records); + require(found, "multicast record uses driver-accepted extent"); +} + +static void +require_logical_handle(CUmemGenericAllocationHandle handle, const char* message) +{ + require( + ((uint64_t)handle & UINT64_C(0xffff000000000000)) == + UINT64_C(0xd94d000000000000), + message); +} + +static CUmulticastObjectProp +properties(void) +{ + CUmulticastObjectProp value = {0}; + + value.numDevices = 1; + value.size = 4096; + value.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + return value; +} + +static void* +map_multicast(void* argument) +{ + struct map_call* call = argument; + + call->result = cuMemMap(0x2000, 4096, 0, call->handle, 0); + return NULL; +} + +static void* +add_multicast_device(void* argument) +{ + struct add_device_call* call = argument; + + call->result = cuMulticastAddDevice(call->handle, 1); + __atomic_store_n(&call->done, 1, __ATOMIC_RELEASE); + return NULL; +} + +static void +test_multicast_map_releases_state_lock(void) +{ + struct map_call map_call = {0}; + struct add_device_call add_call = {0}; + CUmulticastObjectProp props = properties(); + pthread_t map_thread; + pthread_t add_thread; + int completed_while_blocked; + int i; + + require(cuMulticastCreate(&map_call.handle, &props) == CUDA_SUCCESS, "cuMulticastCreate for map"); + add_call.handle = map_call.handle; + fakeEnableBlockingMulticastMap(); + require(pthread_create(&map_thread, NULL, map_multicast, &map_call) == 0, "start multicast map"); + for (i = 0; i < 1000 && !fakeMulticastMapEntered(); i++) + usleep(1000); + require(fakeMulticastMapEntered(), "real multicast map entered"); + + require(pthread_create(&add_thread, NULL, add_multicast_device, &add_call) == 0, "start peer state call"); + for (i = 0; i < 1000 && !__atomic_load_n(&add_call.done, __ATOMIC_ACQUIRE); i++) + usleep(1000); + completed_while_blocked = __atomic_load_n(&add_call.done, __ATOMIC_ACQUIRE); + fakeReleaseMulticastMap(); + require(pthread_join(add_thread, NULL) == 0, "join peer state call"); + require(pthread_join(map_thread, NULL) == 0, "join multicast map"); + + require(completed_while_blocked, "multicast map released interposer state lock"); + require(add_call.result == (CUresult)111, "peer state call result"); + require(map_call.result == CUDA_SUCCESS, "multicast map result"); +} + +static void +test_tracked_behavior(void) { CUmemGenericAllocationHandle handle = 0; + CUmulticastObjectProp props = properties(); size_t granularity = 0; - require( - cuMulticastCreate(&handle, NULL) == (CUresult)110 && handle == 0x456, - "cuMulticastCreate"); + require(cuMulticastCreate(&handle, &props) == CUDA_SUCCESS, "cuMulticastCreate"); + require_logical_handle(handle, "cuMulticastCreate logical handle"); require(cuMulticastAddDevice(handle, 1) == (CUresult)111, "cuMulticastAddDevice"); - require( - cuMulticastBindMem(handle, 2, 3, 4, 5, 0) == (CUresult)112, - "cuMulticastBindMem"); - require( - cuMulticastBindAddr(handle, 2, 3, 4, 0) == (CUresult)113, - "cuMulticastBindAddr"); require( cuMulticastGetGranularity(&granularity, NULL, 0) == (CUresult)114 && granularity == 4096, "cuMulticastGetGranularity"); - require(cuMulticastUnbind(handle, 1, 2, 3) == (CUresult)115, "cuMulticastUnbind"); } static void @@ -61,6 +177,7 @@ test_resolvers(void) void* cuda = dlopen("libcuda.so.1", RTLD_NOW); void* symbol = NULL; CUmemGenericAllocationHandle handle = 0; + CUmulticastObjectProp props = properties(); enum cudaDriverEntryPointQueryResult runtime_status; require(cuda != NULL, "dlopen libcuda"); @@ -69,8 +186,9 @@ test_resolvers(void) symbol != NULL && symbol != (void*)&fakeCuMulticastCreateOriginal, "dlsym substitution"); require( - ((create_type)symbol)(&handle, NULL) == (CUresult)110, - "dlsym forwarding"); + ((create_type)symbol)(&handle, &props) == CUDA_SUCCESS, + "dlsym behavior"); + require_logical_handle(handle, "dlsym logical handle"); symbol = NULL; require( @@ -79,6 +197,11 @@ test_resolvers(void) require( symbol != NULL && symbol != (void*)&fakeCuMulticastCreateOriginal, "cuGetProcAddress substitution"); + handle = 0; + require( + ((create_type)symbol)(&handle, &props) == CUDA_SUCCESS, + "cuGetProcAddress behavior"); + require_logical_handle(handle, "cuGetProcAddress logical handle"); symbol = NULL; require( @@ -89,14 +212,21 @@ test_resolvers(void) runtime_status == cudaDriverEntryPointSuccess && symbol != (void*)&fakeCuMulticastCreateOriginal, "runtime substitution"); + handle = 0; + require( + ((create_type)symbol)(&handle, &props) == CUDA_SUCCESS, + "runtime behavior"); + require_logical_handle(handle, "runtime logical handle"); dlclose(cuda); } int main(void) { - test_direct_forwarding(); + test_tracked_behavior(); + test_multicast_map_releases_state_lock(); test_resolvers(); - puts("multicast forwarding OK"); + test_multicast_records_driver_accepted_extent(); + puts("multicast behavior OK"); return 0; } diff --git a/agent/cmd/cuinterpose/tests/test_cucheckpoint.py b/agent/cmd/cuinterpose/tests/test_cucheckpoint.py index 1769493e..1dfeb3e6 100644 --- a/agent/cmd/cuinterpose/tests/test_cucheckpoint.py +++ b/agent/cmd/cuinterpose/tests/test_cucheckpoint.py @@ -6,6 +6,7 @@ import ctypes import os import queue +import re import shutil import signal import socket @@ -136,7 +137,9 @@ def _worker( restore_fds: tuple[int, int], sync_dir: Path, store_path: Path, + multicast: bool, ) -> None: + _require_launch_job() _cuda_call(driver.cuInit, 0) device = _cuda_call(driver.cuDeviceGet, rank) properties = _allocation_properties(device) @@ -202,16 +205,26 @@ def _worker( input_tensor = symm_mem.empty(NUMEL, dtype=torch.float32, device="cuda") input_tensor.fill_(rank + 1) symm_handle = symm_mem.rendezvous(input_tensor, group=group_name) + if multicast: + if not symm_handle.has_multicast_support: + raise AssertionError("PyTorch silently fell back from CUDA multicast") + if int(symm_handle.multicast_ptr) == 0: + raise AssertionError("PyTorch selected multicast without a multicast VA") + _replace_local_binding_with_address( + rank, + input_tensor, + symm_handle, + properties, + ) + dist.barrier() output = torch.empty_like(input_tensor) - torch.ops.symm_mem.one_shot_all_reduce_out(input_tensor, "sum", group_name, output) + _collective(input_tensor, group_name, output, multicast) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - torch.ops.symm_mem.one_shot_all_reduce_out( - input_tensor, "sum", group_name, output - ) + _collective(input_tensor, group_name, output, multicast) graph.replay() torch.cuda.synchronize() _assert_exact_result(output, "before checkpoint") @@ -266,6 +279,67 @@ def _worker( _destroy_mapped_allocation(private_address, private_size, private_handle) +def _collective( + input_tensor: torch.Tensor, + group_name: str, + output: torch.Tensor, + multicast: bool, +) -> None: + operation = ( + torch.ops.symm_mem.multimem_one_shot_all_reduce_out + if multicast + else torch.ops.symm_mem.one_shot_all_reduce_out + ) + operation(input_tensor, "sum", group_name, output) + + +def _replace_local_binding_with_address( + rank: int, + input_tensor: torch.Tensor, + symm_handle, + properties: driver.CUmemAllocationProp, +) -> None: + granularity = int( + _cuda_call( + driver.cuMemGetAllocationGranularity, + properties, + driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ) + ) + buffer_size = input_tensor.numel() * input_tensor.element_size() + signal_offset = (buffer_size + 15) // 16 * 16 + unrounded_size = signal_offset + symm_mem.get_signal_pad_size() + block_size = (unrounded_size + granularity - 1) // granularity * granularity + multicast_handle = _cuda_call( + driver.cuMemRetainAllocationHandle, + int(symm_handle.multicast_ptr), + ) + _assert_handle_namespace( + multicast_handle, + True, + "retained multicast handle", + ) + device = _cuda_call(driver.cuDeviceGet, rank) + try: + _cuda_call( + driver.cuMulticastUnbind, + multicast_handle, + device, + 0, + block_size, + ) + _cuda_call( + driver.cuMulticastBindAddr, + multicast_handle, + 0, + int(symm_handle.buffer_ptrs[rank]), + block_size, + 0, + ) + finally: + _cuda_call(driver.cuMemRelease, multicast_handle) + + def _assert_exact_result(output: torch.Tensor, stage: str) -> None: expected = torch.full((NUMEL,), 3.0, dtype=torch.float32) actual = output.cpu() @@ -401,6 +475,7 @@ def _fork_workers( restore_fds: tuple[int, int], sync_dir: Path, store_path: Path, + multicast: bool, ) -> None: if torch.cuda.is_initialized(): raise RuntimeError("parent initialized CUDA before forking workers") @@ -411,7 +486,14 @@ def _fork_workers( child = os.fork() if child == 0: try: - _worker(rank, raw_fds, restore_fds, sync_dir, store_path) + _worker( + rank, + raw_fds, + restore_fds, + sync_dir, + store_path, + multicast, + ) except BaseException: # noqa: BLE001 -- report child failures to parent traceback.print_exc() os._exit(1) @@ -455,8 +537,10 @@ def _start_parent( store_path: Path, raw_fds: tuple[int, int], restore_fds: tuple[int, int], + multicast: bool, ) -> subprocess.Popen[str]: environment = os.environ.copy() + launch_job_fds = _require_launch_job() environment.update( { "CUDA_VISIBLE_DEVICES": ",".join(gpus), @@ -466,9 +550,12 @@ def _start_parent( "PYTHONFAULTHANDLER": "1", "PYTHONUNBUFFERED": "1", "TORCH_SYMMEM_IMPLICIT_POOL": "0", - "TORCH_SYMM_MEM_DISABLE_MULTICAST": "1", } ) + if multicast: + environment.pop("TORCH_SYMM_MEM_DISABLE_MULTICAST", None) + else: + environment["TORCH_SYMM_MEM_DISABLE_MULTICAST"] = "1" return subprocess.Popen( [ sys.executable, @@ -481,16 +568,36 @@ def _start_parent( *(str(fd) for fd in restore_fds), str(sync_dir), str(store_path), + "multicast" if multicast else "unicast", ], env=environment, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True, - pass_fds=raw_fds + restore_fds, + pass_fds=raw_fds + restore_fds + launch_job_fds, ) +def _require_launch_job() -> tuple[int, ...]: + job_file = os.environ.get("CUDA_CHECKPOINT_JOB_FILE") + if not job_file: + raise RuntimeError( + "CUDA_CHECKPOINT_JOB_FILE is missing; run this test through " + "cuda-checkpoint --launch-job" + ) + if not Path(job_file).is_file(): + raise RuntimeError( + f"CUDA checkpoint launch-job file is unavailable: {job_file}" + ) + match = re.fullmatch(r"/proc/self/fd/([0-9]+)", job_file) + if match is None: + return () + descriptor = int(match.group(1)) + os.fstat(descriptor) + return (descriptor,) + + def _wait_for_child_pids( parent: subprocess.Popen[str], sync_dir: Path ) -> tuple[int, int]: @@ -654,6 +761,17 @@ def run() -> None: def test_cucheckpoint_preserves_symmetric_memory_cuda_graph(tmp_path: Path) -> None: + _run_cucheckpoint_test(tmp_path, multicast=False) + + +def test_cucheckpoint_preserves_multicast_symmetric_memory_cuda_graph( + tmp_path: Path, +) -> None: + _run_cucheckpoint_test(tmp_path, multicast=True) + + +def _run_cucheckpoint_test(tmp_path: Path, multicast: bool) -> None: + _require_launch_job() interposer, coordinator = _build_native_tools(tmp_path) gpus = _visible_gpus() control_dir = tmp_path / "control" @@ -666,7 +784,7 @@ def test_cucheckpoint_preserves_symmetric_memory_cuda_graph(tmp_path: Path) -> N parent: subprocess.Popen[str] | None = None child_pids: tuple[int, int] = () - output = ("", "") + output: tuple[str | None, str | None] = ("", "") output_collected = False failure: Exception | None = None external_allocations: list[_ExternalAllocation] = [] @@ -682,6 +800,7 @@ def test_cucheckpoint_preserves_symmetric_memory_cuda_graph(tmp_path: Path) -> N store_path, tuple(allocation.fd for allocation in external_allocations), tuple(receiver.fileno() for _, receiver in restore_channels), + multicast, ) for _, receiver in restore_channels: receiver.close() @@ -752,20 +871,24 @@ def test_cucheckpoint_preserves_symmetric_memory_cuda_graph(tmp_path: Path) -> N f"parent PID/return code: {parent_pid}/{parent_returncode}\n" f"forked child PIDs: {child_pids}\n" f"control sockets: {sockets}\n" - f"parent and worker stdout:\n{output[0]}\n" - f"parent and worker stderr:\n{output[1]}" + f"parent and worker stdout:\n{output[0] or ''}\n" + f"parent and worker stderr:\n{output[1] or ''}" ) from failure if __name__ == "__main__": - if len(sys.argv) != 8 or sys.argv[1] != "--parent": + if len(sys.argv) != 9 or sys.argv[1] != "--parent": raise SystemExit( "usage: test_cucheckpoint.py --parent RAW_FD_0 RAW_FD_1 " - "RESTORE_FD_0 RESTORE_FD_1 SYNC_DIR STORE_PATH" + "RESTORE_FD_0 RESTORE_FD_1 SYNC_DIR STORE_PATH " + "(unicast|multicast)" ) + if sys.argv[8] not in {"unicast", "multicast"}: + raise SystemExit("worker mode must be unicast or multicast") _fork_workers( (int(sys.argv[2]), int(sys.argv[3])), (int(sys.argv[4]), int(sys.argv[5])), Path(sys.argv[6]), Path(sys.argv[7]), + sys.argv[8] == "multicast", ) diff --git a/agent/cmd/cuinterpose/util.c b/agent/cmd/cuinterpose/util.c index 32331a91..3adef1d1 100644 --- a/agent/cmd/cuinterpose/util.c +++ b/agent/cmd/cuinterpose/util.c @@ -139,6 +139,43 @@ set_socket_timeouts(int fd, int seconds) : -1; } +/* + * Parses a whole-second timeout override, returning fallback for anything that + * is not a plain positive decimal integer of at most a day. Surrounding + * whitespace is tolerated because config sources often add a trailing newline. + * + * Hand-rolled rather than strtol because every caller is compiled with + * _GNU_SOURCE, which forces _ISOC23_SOURCE on and makes redirect + * the whole strtol family to __isoc23_*@GLIBC_2.38. That would push the shim + * above the MIN_GLIBC floor asserted in the Makefile. atol is not an escape: + * glibc defines it as an inline wrapper around strtol. + * + * Trailing garbage is rejected rather than ignored as strtol would. Reading + * "3x" as 3 seconds would silently arm a timeout far shorter than the operator + * asked for, and the default is the safer reading of a malformed override. + */ +unsigned +bounded_seconds(const char* value, unsigned fallback) +{ + unsigned long parsed = 0; + const char* digits; + + if (value == NULL) + return fallback; + while (*value == ' ' || *value == '\t' || *value == '\n' || *value == '\r') + value++; + for (digits = value; *value >= '0' && *value <= '9'; value++) { + parsed = parsed * 10 + (unsigned long)(*value - '0'); + if (parsed > 86400) + return fallback; + } + if (value == digits || parsed == 0) + return fallback; + while (*value == ' ' || *value == '\t' || *value == '\n' || *value == '\r') + value++; + return *value == '\0' ? (unsigned)parsed : fallback; +} + int send_header(int fd, const struct cuinterposer_header* header, int passed_fd) { diff --git a/agent/cmd/cuinterpose/util.h b/agent/cmd/cuinterpose/util.h index 07970b81..724668d9 100644 --- a/agent/cmd/cuinterpose/util.h +++ b/agent/cmd/cuinterpose/util.h @@ -21,6 +21,7 @@ bool is_lower_hex_id(const char value[CUINTERPOSER_ID_SIZE]); bool header_strings_terminated(const struct cuinterposer_header* header); void header_error(struct cuinterposer_header* header, const char* message); int set_socket_timeouts(int fd, int seconds); +unsigned bounded_seconds(const char* value, unsigned fallback); int send_header(int fd, const struct cuinterposer_header* header, int passed_fd); int receive_header(int fd, struct cuinterposer_header* header, int* passed_fd);