diff --git a/csrc/bindings.cu b/csrc/bindings.cu index 6f00f4a..fa915ed 100644 --- a/csrc/bindings.cu +++ b/csrc/bindings.cu @@ -53,7 +53,8 @@ PYBIND11_MODULE(_C, m) { pybind11::arg("size_bytes"), pybind11::arg("num_devices"), "Bind the local chunk's owned physical memory (nvl_dist_alloc handle) " "to the multicast object and map a multicast VA; returns a keepalive " - "int32 tensor whose data_ptr is the multimem address."); + "int32 tensor whose data_ptr is the multimem address. Consumes " + "mc_handle and borrows owned_mem_handle."); m.def("nvl_release_mem_handle", &nvl_release_mem_handle, pybind11::arg("mem_handle"), "Release an owned mem handle returned by nvl_dist_alloc."); diff --git a/csrc/nvl_shared_buffer.cuh b/csrc/nvl_shared_buffer.cuh index 0b7747c..f38ea76 100644 --- a/csrc/nvl_shared_buffer.cuh +++ b/csrc/nvl_shared_buffer.cuh @@ -1,39 +1,41 @@ #pragma once #include -#include -#include #include #include #include +#include + #include #include #include #include +#include #ifndef CUCHECK -#define CUCHECK(cmd) do { \ - CUresult err = cmd; \ - if (err != CUDA_SUCCESS) { \ - const char *errStr; \ - cuGetErrorString(err, &errStr); \ - fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ - __FILE__, __LINE__, errStr); \ - exit(EXIT_FAILURE); \ - } \ +#define CUCHECK(cmd) do { \ + const CUresult err = (cmd); \ + if (err != CUDA_SUCCESS) { \ + const char *err_name = nullptr; \ + const char *err_string = nullptr; \ + cuGetErrorName(err, &err_name); \ + cuGetErrorString(err, &err_string); \ + TORCH_CHECK(false, "CUDA driver call `", #cmd, "` failed with ", \ + err_name ? err_name : "unknown error", " (code ", \ + static_cast(err), "): ", \ + err_string ? err_string : "unknown error"); \ + } \ } while(0) #endif #ifndef CUDACHECK -#define CUDACHECK(cmd) do { \ - cudaError_t err = cmd; \ - if (err != cudaSuccess) { \ - fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n", \ - __FILE__, __LINE__, cudaGetErrorString(err)); \ - exit(EXIT_FAILURE); \ - } \ +#define CUDACHECK(cmd) do { \ + const cudaError_t err = (cmd); \ + TORCH_CHECK(err == cudaSuccess, "CUDA runtime call `", #cmd, \ + "` failed with ", cudaGetErrorName(err), " (code ", \ + static_cast(err), "): ", cudaGetErrorString(err)); \ } while(0) #endif @@ -66,36 +68,18 @@ static inline std::tuple nvl_prepare( return {size, allocated_size, device_id}; } -static inline at::Tensor make_vmm_tensor( - void *ptr, - size_t nbytes, - size_t allocated_size, +static inline auto make_vmm_deleter( int device_id, - const std::vector &shape, - at::ScalarType dtype + CUdeviceptr dptr, + size_t allocated_size ) { - auto deleter = [device_id, ptr, allocated_size](void *p) { + return [device_id, dptr, allocated_size](void *p) { if (!p) return; c10::cuda::CUDAGuard guard(device_id); cudaStreamSynchronize(c10::cuda::getCurrentCUDAStream().stream()); - cuMemUnmap(reinterpret_cast(ptr), allocated_size); - cuMemAddressFree(reinterpret_cast(ptr), allocated_size); + cuMemUnmap(dptr, allocated_size); + cuMemAddressFree(dptr, allocated_size); }; - - auto storage = c10::Storage( - c10::Storage::use_byte_size_t(), - static_cast(nbytes), - at::InefficientStdFunctionContext::makeDataPtr( - ptr, std::move(deleter), c10::Device(c10::kCUDA, device_id)), - /*allocator=*/nullptr, - /*resizable=*/false); - - auto impl = c10::make_intrusive( - std::move(storage), - c10::DispatchKeySet(c10::DispatchKey::CUDA), - c10::scalarTypeToTypeMeta(dtype)); - impl->set_sizes_contiguous(shape); - return at::Tensor(std::move(impl)); } // Returns (keepalive VA tensor, exported POSIX fd, owned mem handle as int64). @@ -117,12 +101,34 @@ inline std::tuple nvl_dist_alloc( prop.location.id = device_id; prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; - CUmemGenericAllocationHandle mem_handle; + CUmemGenericAllocationHandle mem_handle = {}; + bool mem_handle_owned = false; + CUdeviceptr dptr = 0; + bool va_reserved = false; + bool va_mapped = false; + int fd = -1; + auto cleanup = c10::make_scope_exit([&]() { + if (fd >= 0) { + close(fd); + } + if (va_mapped) { + cuMemUnmap(dptr, allocated_size); + } + if (va_reserved) { + cuMemAddressFree(dptr, allocated_size); + } + if (mem_handle_owned) { + cuMemRelease(mem_handle); + } + }); + CUCHECK(cuMemCreate(&mem_handle, allocated_size, &prop, 0)); + mem_handle_owned = true; - CUdeviceptr dptr; CUCHECK(cuMemAddressReserve(&dptr, allocated_size, 0, 0, 0)); + va_reserved = true; CUCHECK(cuMemMap(dptr, allocated_size, 0, mem_handle, 0)); + va_mapped = true; CUmemAccessDesc access_desc = {}; access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; @@ -135,17 +141,39 @@ inline std::tuple nvl_dist_alloc( // caller is responsible for nvl_release_mem_handle. The physical memory is // referenced by the unicast map (keepalive) and (optionally) multicast; // it is only truly freed after all unmaps once the handle is released. - int fd = -1; CUCHECK(cuMemExportToShareableHandle( &fd, mem_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd"); - auto keepalive = make_vmm_tensor( - reinterpret_cast(dptr), nbytes, allocated_size, - device_id, chunk_shape, dtype); + auto data_ptr = at::InefficientStdFunctionContext::makeDataPtr( + reinterpret_cast(dptr), + make_vmm_deleter(device_id, dptr, allocated_size), + c10::Device(c10::kCUDA, device_id)); - return {std::move(keepalive), static_cast(fd), - static_cast(mem_handle)}; + // data_ptr now owns the mapping and VA. + va_mapped = false; + va_reserved = false; + auto storage = c10::Storage( + c10::Storage::use_byte_size_t(), + static_cast(nbytes), + std::move(data_ptr), + /*allocator=*/nullptr, + /*resizable=*/false); + auto impl = c10::make_intrusive( + std::move(storage), + c10::DispatchKeySet(c10::DispatchKey::CUDA), + c10::scalarTypeToTypeMeta(dtype)); + impl->set_sizes_contiguous(chunk_shape); + auto keepalive = at::Tensor(std::move(impl)); + + // Construct the complete return value before transferring the fd and + // allocation handle to the caller. + auto result = std::make_tuple( + std::move(keepalive), static_cast(fd), + static_cast(mem_handle)); + fd = -1; + mem_handle_owned = false; + return result; } // Release the owned mem handle returned by nvl_dist_alloc (the physical memory @@ -183,19 +211,39 @@ inline at::Tensor nvl_dist_map( size_t total_allocated_size = chunk_allocated_size * world_size; size_t total_nbytes = chunk_nbytes * world_size; - CUdeviceptr dptr; + CUdeviceptr dptr = 0; + bool va_reserved = false; + size_t mapped_chunks = 0; + auto cleanup = c10::make_scope_exit([&]() { + for (size_t i = 0; i < mapped_chunks; i++) { + cuMemUnmap(dptr + i * chunk_allocated_size, chunk_allocated_size); + } + if (va_reserved) { + cuMemAddressFree(dptr, total_allocated_size); + } + }); + CUCHECK(cuMemAddressReserve(&dptr, total_allocated_size, 0, 0, 0)); + va_reserved = true; for (int64_t i = 0; i < world_size; i++) { int fd = static_cast(fds[i]); - CUmemGenericAllocationHandle mem_handle; + CUmemGenericAllocationHandle mem_handle = {}; + bool mem_handle_owned = false; + auto release_handle = c10::make_scope_exit([&]() { + if (mem_handle_owned) { + cuMemRelease(mem_handle); + } + }); CUCHECK(cuMemImportFromShareableHandle( &mem_handle, reinterpret_cast(static_cast(fd)), CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); + mem_handle_owned = true; CUdeviceptr chunk_va = dptr + i * chunk_allocated_size; CUCHECK(cuMemMap(chunk_va, chunk_allocated_size, 0, mem_handle, 0)); + mapped_chunks++; // All chunks get RW access for dispatch writes CUmemAccessDesc access = {}; @@ -204,6 +252,9 @@ inline at::Tensor nvl_dist_map( access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; CUCHECK(cuMemSetAccess(chunk_va, chunk_allocated_size, &access, 1)); + // Release exactly once. A failed driver call may report an earlier + // asynchronous error, so retrying this opaque handle is not safe. + mem_handle_owned = false; CUCHECK(cuMemRelease(mem_handle)); } @@ -223,12 +274,14 @@ inline at::Tensor nvl_dist_map( cuMemAddressFree(dptr, tas); }; + auto data_ptr = at::InefficientStdFunctionContext::makeDataPtr( + reinterpret_cast(dptr), std::move(deleter), + c10::Device(c10::kCUDA, device_id)); + cleanup.release(); auto storage = c10::Storage( c10::Storage::use_byte_size_t(), static_cast(total_nbytes), - at::InefficientStdFunctionContext::makeDataPtr( - reinterpret_cast(dptr), std::move(deleter), - c10::Device(c10::kCUDA, device_id)), + std::move(data_ptr), /*allocator=*/nullptr, /*resizable=*/false); @@ -237,7 +290,8 @@ inline at::Tensor nvl_dist_map( c10::DispatchKeySet(c10::DispatchKey::CUDA), c10::scalarTypeToTypeMeta(dtype)); impl->set_sizes_contiguous(full_shape); - return at::Tensor(std::move(impl)); + auto result = at::Tensor(std::move(impl)); + return result; } // ============================================================ @@ -315,15 +369,30 @@ inline std::tuple nvl_multicast_create( prop.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; prop.flags = 0; - CUmemGenericAllocationHandle mc_handle; + CUmemGenericAllocationHandle mc_handle = {}; + bool mc_handle_owned = false; + int fd = -1; + auto cleanup = c10::make_scope_exit([&]() { + if (fd >= 0) { + close(fd); + } + if (mc_handle_owned) { + cuMemRelease(mc_handle); + } + }); + CUCHECK(cuMulticastCreate(&mc_handle, &prop)); + mc_handle_owned = true; - int fd = -1; CUCHECK(cuMemExportToShareableHandle( &fd, mc_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd"); - return {static_cast(mc_handle), static_cast(fd)}; + auto result = std::make_tuple( + static_cast(mc_handle), static_cast(fd)); + mc_handle_owned = false; + fd = -1; + return result; } // non-root: import the multicast object from the POSIX fd sent by root. The @@ -357,12 +426,42 @@ inline at::Tensor nvl_multicast_bind_map( int64_t mc_handle_u64, int64_t owned_mem_handle_u64, int64_t size_bytes, int64_t num_devices ) { - int device_id; - CUDACHECK(cudaGetDevice(&device_id)); + // Ownership transfers here: _create_nvl_multicast_view has no outer + // release for the created/imported multicast handle. On success the + // returned tensor's deleter releases it; this guard covers failure. CUmemGenericAllocationHandle mc_handle = static_cast(mc_handle_u64); + // Borrowed: create_nvl_dist_multicast_tensor separately releases the + // physical allocation handle in its outer finally block, including when + // this function throws. CUmemGenericAllocationHandle mem_handle = static_cast(owned_mem_handle_u64); + bool mc_handle_owned = true; + CUdevice cu_device = 0; + bool mem_bound = false; + CUdeviceptr mc_va = 0; + bool va_reserved = false; + bool va_mapped = false; + auto cleanup = c10::make_scope_exit([&]() { + if (va_mapped) { + cuMemUnmap(mc_va, static_cast(size_bytes)); + } + if (va_reserved) { + cuMemAddressFree(mc_va, static_cast(size_bytes)); + } + if (mem_bound) { + cuMulticastUnbind( + mc_handle, cu_device, /*mcOffset=*/0, + static_cast(size_bytes)); + } + if (mc_handle_owned) { + cuMemRelease(mc_handle); + } + }); + + int device_id; + CUDACHECK(cudaGetDevice(&device_id)); + CUCHECK(cuDeviceGet(&cu_device, device_id)); const size_t gran = nvl_multicast_granularity(static_cast(num_devices)); TORCH_CHECK(static_cast(size_bytes) % gran == 0, @@ -373,14 +472,16 @@ inline at::Tensor nvl_multicast_bind_map( // Bind this rank's owned physical allocation to offset 0 of the multicast object. CUCHECK(cuMulticastBindMem( mc_handle, /*mcOffset=*/0, mem_handle, /*memOffset=*/0, size, 0)); + mem_bound = true; // Reserve a separate VA for the multicast handle and map it. Pass 0 for // alignment (the driver aligns to granularity by default, same as // nvl_dist_alloc/nvl_dist_map; a non-zero alignment is rejected by // cuMemAddressReserve as invalid argument). - CUdeviceptr mc_va; CUCHECK(cuMemAddressReserve(&mc_va, size, 0, 0, 0)); + va_reserved = true; CUCHECK(cuMemMap(mc_va, size, 0, mc_handle, 0)); + va_mapped = true; CUmemAccessDesc access = {}; access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; @@ -400,12 +501,17 @@ inline at::Tensor nvl_multicast_bind_map( }; std::vector shape = {static_cast(size / sizeof(int32_t))}; + auto data_ptr = at::InefficientStdFunctionContext::makeDataPtr( + reinterpret_cast(mc_va), std::move(deleter), + c10::Device(c10::kCUDA, device_id)); + va_mapped = false; + va_reserved = false; + mem_bound = false; + mc_handle_owned = false; auto storage = c10::Storage( c10::Storage::use_byte_size_t(), static_cast(size), - at::InefficientStdFunctionContext::makeDataPtr( - reinterpret_cast(mc_va), std::move(deleter), - c10::Device(c10::kCUDA, device_id)), + std::move(data_ptr), /*allocator=*/nullptr, /*resizable=*/false); @@ -414,5 +520,6 @@ inline at::Tensor nvl_multicast_bind_map( c10::DispatchKeySet(c10::DispatchKey::CUDA), c10::scalarTypeToTypeMeta(at::kInt)); impl->set_sizes_contiguous(shape); - return at::Tensor(std::move(impl)); + auto result = at::Tensor(std::move(impl)); + return result; } diff --git a/tests/test_cpp_error_propagation.py b/tests/test_cpp_error_propagation.py new file mode 100644 index 0000000..37dccf2 --- /dev/null +++ b/tests/test_cpp_error_propagation.py @@ -0,0 +1,39 @@ +"""Regression tests for errors crossing the CUDA extension boundary.""" + +import subprocess +import sys +import textwrap + + +def test_non_cuda_posix_fd_raises_without_terminating_process(): + """A driver error must be catchable instead of calling exit().""" + script = textwrap.dedent( + """ + import os + + from moonep._C import nvl_multicast_import + + fd = os.open(os.devnull, os.O_RDONLY) + try: + try: + nvl_multicast_import(fd) + except RuntimeError as error: + assert "cuMemImportFromShareableHandle" in str(error), str(error) + else: + raise AssertionError("non-CUDA POSIX fd was accepted") + finally: + os.close(fd) + + print("process continued after CUDA driver error") + """ + ) + + result = subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + assert "process continued after CUDA driver error" in result.stdout