derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1995 lines
85 KiB
C++
1995 lines
85 KiB
C++
/*************************************************************************
|
|
* Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
|
|
*
|
|
* See LICENSE.txt for license information
|
|
************************************************************************/
|
|
|
|
/**
|
|
* @file SuspendResumeMPITests.cpp
|
|
* @brief Multi-process tests for the suspend/resume feature -- both the
|
|
* internal ncclCommMemSuspend/Resume/Stats path and the public
|
|
* ncclCommSuspend/Resume/MemStats API.
|
|
*
|
|
* The single-process unit-test suite in `test/mem_manager/MemManagerTests.cpp`
|
|
* covers all early-return error paths, `ncclCommMemStats`, and the manager
|
|
* state machine via fake / real-VMM entries that don't need a bootstrap.
|
|
*
|
|
* This file extends coverage to the bootstrap-coupled paths in
|
|
* `ncclCommMemSuspend` / `ncclCommMemResume`:
|
|
* - Cross-rank `bootstrapBarrier(0xBEEF)` on entry to Suspend.
|
|
* - Cross-rank `bootstrapBarrier(0xBEEF)` after local restore in Resume.
|
|
* - `bootstrapAllGather` of per-rank export counts in Resume Step 3.
|
|
* - `bootstrapSend`/`bootstrapRecv(0xFEED)` handle-info exchange.
|
|
* - Final `bootstrapBarrier(0xCAFE)` after peer re-import.
|
|
* - The peer-import matching loop in Step 4 (both positive and negative).
|
|
* - `ncclProxyClientGetFdBlocking` round-trip used by Step 4 to convert
|
|
* a peer's CUmemGenericAllocationHandle into a POSIX FD via UDS.
|
|
* - End-to-end data preservation: rank 0 fills an `ncclMemOffload` buffer
|
|
* with a canary, both ranks Suspend (CPU backup) + Resume (restore from
|
|
* backup, re-export, peer re-import), rank 1 reads the canary back
|
|
* through its imported VA.
|
|
*
|
|
* The tests allocate VMM directly via raw HIP (hipMemCreate / Reserve / Map /
|
|
* SetAccess) instead of going through ncclCuMemAlloc - the wrapper has its
|
|
* own issues on ROCm 7.0/7.2 and the goal is to exercise mem_manager logic,
|
|
* not the allocator. If the runtime really lacks raw VMM, hipMemCreate fails
|
|
* with an explicit HIP error code instead of a misleading version skip.
|
|
*
|
|
* Suites:
|
|
* - MemManagerAnyRanks - each rank runs the same local Suspend/Resume
|
|
* scenario in lockstep. The cross-rank bootstrapBarrier inside
|
|
* ncclCommMemSuspend / ncclCommMemResume keeps all ranks in sync, so
|
|
* every test in this suite passes on any np >= 1.
|
|
* - MemManagerTwoRank - peer-coupled scenarios that need exactly 2 ranks
|
|
* (rank 0 exporter, rank 1 importer).
|
|
* - SuspendResumeBasic / MemStats / CollectiveIntegrity / Lifecycle / Group
|
|
* / ArgValidation / MemStatsValidation - public API tests for
|
|
* ncclCommSuspend / ncclCommResume / ncclCommMemStats. Need np >= 2.
|
|
*
|
|
* Run (a single mpirun -np 2 invocation runs the entire file):
|
|
* mpirun -np 2 ./rccl-UnitTestsMPI --gtest_filter='MemManager*:SuspendResume*'
|
|
*
|
|
* Or, if you want to focus on one group of suites:
|
|
* mpirun -np 1 ./rccl-UnitTestsMPI --gtest_filter='MemManagerAnyRanks.*'
|
|
* mpirun -np 2 ./rccl-UnitTestsMPI --gtest_filter='MemManagerTwoRank.*'
|
|
* mpirun -np 2 ./rccl-UnitTestsMPI --gtest_filter='SuspendResume*'
|
|
*/
|
|
|
|
#include "DeviceBufferHelpers.hpp"
|
|
#include "MPITestBase.hpp"
|
|
#include "ResourceGuards.hpp"
|
|
#include "TestChecks.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <iomanip>
|
|
#include <iostream>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <unistd.h>
|
|
#include <vector>
|
|
|
|
#ifdef MPI_TESTS_ENABLED
|
|
|
|
#include "bootstrap.h"
|
|
#include "comm.h"
|
|
#include "mem_manager.h"
|
|
#include "proxy.h"
|
|
|
|
using namespace MPITestConstants;
|
|
using namespace RCCLTestGuards;
|
|
using namespace RCCLTestHelpers;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// VMM helpers
|
|
//
|
|
// We talk to the HIP driver directly (hipMemCreate / Reserve / Map / SetAccess)
|
|
// instead of going through ncclCuMemAlloc. The wrapper has rollback paths and
|
|
// a ROCm-2550 zero-fill memset that hide the actual VMM call site, and it
|
|
// segfaults inside cuMemCreate on ROCm 7.0/7.2 even when raw VMM works.
|
|
// Same approach the standalone unit test takes (Track_RealVmm_PosixFd).
|
|
// ---------------------------------------------------------------------------
|
|
|
|
namespace
|
|
{
|
|
struct VmmAlloc
|
|
{
|
|
void* ptr = nullptr;
|
|
hipDeviceptr_t pdev = 0;
|
|
size_t size = 0;
|
|
hipMemGenericAllocationHandle_t handle = 0;
|
|
};
|
|
|
|
void allocateVmmPosixFd(int dev, size_t requestedSize, VmmAlloc* out)
|
|
{
|
|
ASSERT_NE(out, nullptr);
|
|
|
|
hipMemAllocationProp prop = {};
|
|
prop.type = hipMemAllocationTypePinned;
|
|
prop.location.type = hipMemLocationTypeDevice;
|
|
prop.location.id = dev;
|
|
prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
|
|
prop.allocFlags.gpuDirectRDMACapable = 1;
|
|
|
|
size_t granularity = 0;
|
|
ASSERT_EQ(hipMemGetAllocationGranularity(&granularity, &prop,
|
|
hipMemAllocationGranularityMinimum),
|
|
hipSuccess);
|
|
ASSERT_GT(granularity, 0u);
|
|
size_t size = ((requestedSize + granularity - 1) / granularity) * granularity;
|
|
|
|
hipMemGenericAllocationHandle_t handle = 0;
|
|
ASSERT_EQ(hipMemCreate(&handle, size, &prop, 0), hipSuccess);
|
|
|
|
hipDeviceptr_t pdev = 0;
|
|
ASSERT_EQ(hipMemAddressReserve(&pdev, size, granularity, 0, 0), hipSuccess);
|
|
ASSERT_EQ(hipMemMap(pdev, size, 0, handle, 0), hipSuccess);
|
|
|
|
hipMemAccessDesc accessDesc = {};
|
|
accessDesc.location.type = hipMemLocationTypeDevice;
|
|
accessDesc.location.id = dev;
|
|
accessDesc.flags = hipMemAccessFlagsProtReadWrite;
|
|
ASSERT_EQ(hipMemSetAccess(pdev, size, &accessDesc, 1), hipSuccess);
|
|
|
|
out->ptr = reinterpret_cast<void*>(pdev);
|
|
out->pdev = pdev;
|
|
out->size = size;
|
|
out->handle = handle;
|
|
}
|
|
|
|
// Reserve only (no physical handle / no map). Used to set up a "released" VA
|
|
// region as if Suspend had already torn the physical backing down. Lets a test
|
|
// drive Resume's restore path without needing a Suspend/Resume roundtrip.
|
|
void reserveVaOnly(int dev, size_t requestedSize, VmmAlloc* out)
|
|
{
|
|
ASSERT_NE(out, nullptr);
|
|
|
|
hipMemAllocationProp prop = {};
|
|
prop.type = hipMemAllocationTypePinned;
|
|
prop.location.type = hipMemLocationTypeDevice;
|
|
prop.location.id = dev;
|
|
prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor;
|
|
prop.allocFlags.gpuDirectRDMACapable = 1;
|
|
|
|
size_t granularity = 0;
|
|
ASSERT_EQ(hipMemGetAllocationGranularity(&granularity, &prop,
|
|
hipMemAllocationGranularityMinimum),
|
|
hipSuccess);
|
|
size_t size = ((requestedSize + granularity - 1) / granularity) * granularity;
|
|
|
|
hipDeviceptr_t pdev = 0;
|
|
ASSERT_EQ(hipMemAddressReserve(&pdev, size, granularity, 0, 0), hipSuccess);
|
|
|
|
out->ptr = reinterpret_cast<void*>(pdev);
|
|
out->pdev = pdev;
|
|
out->size = size;
|
|
out->handle = 0;
|
|
}
|
|
|
|
// Mirrors ncclCuMemFree: recover the live handle from the VA (cached
|
|
// a->handle is stale after Suspend/Resume). Retain bumps refcount by 1, so
|
|
// release twice. Reserved-but-unmapped VA falls through to AddressFree.
|
|
void releaseVmm(VmmAlloc* a)
|
|
{
|
|
if (a == nullptr || a->pdev == 0) return;
|
|
|
|
hipMemGenericAllocationHandle_t live = 0;
|
|
if (hipMemRetainAllocationHandle(&live, a->ptr) == hipSuccess) {
|
|
(void)hipMemRelease(live);
|
|
(void)hipMemUnmap(a->pdev, a->size);
|
|
(void)hipMemRelease(live);
|
|
} else {
|
|
(void)hipMemUnmap(a->pdev, a->size);
|
|
}
|
|
(void)hipMemAddressFree(a->pdev, a->size);
|
|
a->pdev = 0;
|
|
a->handle = 0;
|
|
a->ptr = nullptr;
|
|
a->size = 0;
|
|
}
|
|
|
|
// Grant local-device RW access to a previously reserved + mapped VA range.
|
|
// Used after hipMemMap on an imported handle to make the VA usable.
|
|
void applyDeviceRwAccess(int dev, hipDeviceptr_t pdev, size_t size)
|
|
{
|
|
hipMemAccessDesc accessDesc = {};
|
|
accessDesc.location.type = hipMemLocationTypeDevice;
|
|
accessDesc.location.id = dev;
|
|
accessDesc.flags = hipMemAccessFlagsProtReadWrite;
|
|
ASSERT_EQ(hipMemSetAccess(pdev, size, &accessDesc, 1), hipSuccess);
|
|
}
|
|
|
|
// Pour `pattern` over the entire VA on the device, then DToH + verify every
|
|
// byte equals `pattern`. Uses bulk std::all_of so failures don't spam gtest
|
|
// with millions of EXPECT_EQ lines on big buffers. `what` is taken by const-
|
|
// reference so callers can pass either a string literal or a `std::to_string`-
|
|
// composed message without forcing a `.c_str()` everywhere.
|
|
void verifyCanaryReadback(void* ptr, size_t size, uint8_t pattern, const std::string& what)
|
|
{
|
|
std::vector<uint8_t> host(size, static_cast<uint8_t>(~pattern));
|
|
ASSERT_EQ(hipMemcpy(host.data(), ptr, size, hipMemcpyDeviceToHost), hipSuccess) << what;
|
|
bool ok = std::all_of(host.begin(), host.end(),
|
|
[pattern](uint8_t b) { return b == pattern; });
|
|
EXPECT_TRUE(ok) << what << ": canary 0x" << std::hex << static_cast<int>(pattern)
|
|
<< " not preserved";
|
|
}
|
|
|
|
// ncclMemUntrack + releaseVmm. Mirrors the ncclCudaFree pattern.
|
|
void untrackAndRelease(struct ncclComm* comm, void* ptr, VmmAlloc* a)
|
|
{
|
|
ASSERT_NE(comm, nullptr);
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, ptr, a->size), ncclSuccess);
|
|
releaseVmm(a);
|
|
}
|
|
|
|
// Snapshot of manager state taken before the test adds its own entries.
|
|
// With NCCL_CUMEM_ENABLE=1 a freshly created comm already tracks internal
|
|
// RCCL allocations (proxy P2P buffers, channel buffers etc.), so absolute
|
|
// counters/sizes can't be asserted directly - tests must check deltas.
|
|
struct MgrBaseline
|
|
{
|
|
int numEntries = 0;
|
|
uint64_t cpuBackupUsage = 0;
|
|
uint64_t total = 0;
|
|
uint64_t persist = 0;
|
|
uint64_t suspendBytes = 0;
|
|
};
|
|
|
|
MgrBaseline captureBaseline(struct ncclComm* comm)
|
|
{
|
|
MgrBaseline b{};
|
|
EXPECT_NE(comm, nullptr);
|
|
EXPECT_NE(comm->memManager, nullptr);
|
|
b.numEntries = comm->memManager->numEntries;
|
|
b.cpuBackupUsage = comm->memManager->cpuBackupUsage;
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemTotal, &b.total), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemPersist, &b.persist), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspend, &b.suspendBytes), ncclSuccess);
|
|
return b;
|
|
}
|
|
|
|
// Run one Suspend/Resume cycle to learn how much CPU backup the comm's
|
|
// internal Offload entries consume - tests with their own Offload Track must
|
|
// add their size on top of this baseline when checking cpuBackupUsage.
|
|
uint64_t captureInternalOffloadCpuBackup(struct ncclComm* comm)
|
|
{
|
|
EXPECT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
uint64_t v = comm->memManager->cpuBackupUsage;
|
|
EXPECT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
return v;
|
|
}
|
|
|
|
// Find the entry tracking `ptr`. With internals tracked it's no longer
|
|
// guaranteed to be at manager->entries[0].
|
|
ncclDynMemEntry* findEntryByPtr(struct ncclComm* comm, void* ptr)
|
|
{
|
|
for (auto* e = comm->memManager->entries; e != nullptr; e = e->next) {
|
|
if (e->ptr == ptr) return e;
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
// Snapshot of GPU + CPU memory usage at one point in a Suspend/Resume cycle.
|
|
// Combines what the manager tracks (ncclCommMemStats + cpuBackupUsage) with
|
|
// what the OS reports for the process (hipMemGetInfo) so the caller can see
|
|
// both bookkeeping and actual physical impact of Suspend/Resume.
|
|
struct MemFootprint
|
|
{
|
|
uint64_t mgrTotal = 0; // ncclStatGpuMemTotal (bytes)
|
|
uint64_t mgrPersist = 0; // ncclStatGpuMemPersist (bytes)
|
|
uint64_t mgrSuspendable = 0; // ncclStatGpuMemSuspend (bytes)
|
|
uint64_t mgrSuspended = 0; // ncclStatGpuMemSuspended (0 / 1 flag)
|
|
uint64_t cpuBackupUsage = 0; // manager->cpuBackupUsage (bytes)
|
|
size_t gpuFree = 0; // hipMemGetInfo free (bytes)
|
|
size_t gpuTotal = 0; // hipMemGetInfo total (bytes)
|
|
size_t gpuUsedByProc = 0; // gpuTotal - gpuFree (bytes)
|
|
};
|
|
|
|
MemFootprint captureMemFootprint(struct ncclComm* comm)
|
|
{
|
|
MemFootprint f{};
|
|
EXPECT_NE(comm, nullptr);
|
|
EXPECT_NE(comm->memManager, nullptr);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemTotal, &f.mgrTotal), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemPersist, &f.mgrPersist), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspend, &f.mgrSuspendable), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspended, &f.mgrSuspended), ncclSuccess);
|
|
f.cpuBackupUsage = comm->memManager->cpuBackupUsage;
|
|
EXPECT_EQ(hipMemGetInfo(&f.gpuFree, &f.gpuTotal), hipSuccess);
|
|
f.gpuUsedByProc = (f.gpuTotal > f.gpuFree) ? (f.gpuTotal - f.gpuFree) : 0;
|
|
return f;
|
|
}
|
|
|
|
// Pretty-print one footprint line. Kept on std::cout so the host test wrapper
|
|
// captures it the same way as gtest's own output. Drop these prints (and the
|
|
// helpers above) once we no longer need to eyeball the numbers.
|
|
void printMemFootprint(int rank, const char* tag, const MemFootprint& f)
|
|
{
|
|
constexpr double MiB = 1024.0 * 1024.0;
|
|
std::cout << "[MemFootprint rank=" << rank << "] " << tag << '\n'
|
|
<< " mgr.Total = " << std::fixed << std::setprecision(2)
|
|
<< (f.mgrTotal / MiB) << " MiB\n"
|
|
<< " mgr.Persist = " << (f.mgrPersist / MiB) << " MiB\n"
|
|
<< " mgr.Suspendable = " << (f.mgrSuspendable / MiB) << " MiB\n"
|
|
<< " mgr.Suspended = " << f.mgrSuspended << '\n'
|
|
<< " cpuBackupUsage = " << (f.cpuBackupUsage / MiB) << " MiB\n"
|
|
<< " gpu.UsedByProc = " << (f.gpuUsedByProc / MiB) << " MiB"
|
|
<< " (free " << (f.gpuFree / MiB)
|
|
<< " / total " << (f.gpuTotal / MiB) << " MiB)" << std::endl;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Peer export/import helpers - shared by the multi-rank tests.
|
|
//
|
|
// PeerMeta is what the owner ships over bootstrapSend so the importer can
|
|
// reconstruct the matching ncclMemTrackImportFromPeer call: handle value
|
|
// (round-tripped through the proxy GetFd UDS for FD conversion), VA address
|
|
// (matched verbatim by Resume Step 4's matching loop), aligned size, and the
|
|
// owner's cudaDev.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct PeerMeta
|
|
{
|
|
uint64_t handleVal; // owner's hipMemGenericAllocationHandle_t value
|
|
uint64_t ownerPtr; // owner's VA (uintptr_t) - matched verbatim in Step 4
|
|
uint64_t size; // owner's aligned size - matched verbatim
|
|
int ownerDev; // owner's cudaDev - stored in entry.desc.imported
|
|
};
|
|
|
|
// Owner side: Track the buffer, MarkExportToPeer, and ship the metadata over
|
|
// bootstrapSend. The owner keeps `vmm` alive through Suspend/Resume.
|
|
// Each (peerRank, tag) pair must be unique across concurrent calls.
|
|
void exportAndShipMeta(struct ncclComm* comm, const VmmAlloc& vmm, int peerRank,
|
|
ncclMemType_t memType, int tag)
|
|
{
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, vmm.ptr, vmm.size, vmm.handle,
|
|
hipMemHandleTypePosixFileDescriptor, memType),
|
|
ncclSuccess);
|
|
ASSERT_EQ(ncclDynMemMarkExportToPeer(comm->memManager, vmm.ptr, peerRank), ncclSuccess);
|
|
|
|
PeerMeta meta{};
|
|
meta.handleVal = (uint64_t)(uintptr_t)vmm.handle;
|
|
meta.ownerPtr = reinterpret_cast<uintptr_t>(vmm.ptr);
|
|
meta.size = static_cast<uint64_t>(vmm.size);
|
|
meta.ownerDev = comm->cudaDev;
|
|
ASSERT_EQ(bootstrapSend(comm->bootstrap, peerRank, tag, &meta, sizeof(meta)),
|
|
ncclSuccess);
|
|
}
|
|
|
|
// Importer side: bootstrapRecv the metadata, ask owner's proxy for a POSIX
|
|
// FD, import + map + setAccess into a fresh local VA, then Track as imported.
|
|
// Fills `out` with VA + size + the imported handle so untrackAndRelease() can
|
|
// free everything later.
|
|
void recvAndImportPeerBuffer(struct ncclComm* comm, int ownerRank, ncclMemType_t memType,
|
|
int tag, VmmAlloc* out)
|
|
{
|
|
PeerMeta meta{};
|
|
ASSERT_EQ(bootstrapRecv(comm->bootstrap, ownerRank, tag, &meta, sizeof(meta)),
|
|
ncclSuccess);
|
|
ASSERT_GT(meta.size, 0u);
|
|
ASSERT_NE(meta.ownerPtr, 0u);
|
|
|
|
int fd = -1;
|
|
hipMemGenericAllocationHandle_t peerHval =
|
|
(hipMemGenericAllocationHandle_t)(uintptr_t)meta.handleVal;
|
|
ASSERT_EQ(ncclProxyClientGetFdBlocking(comm, ownerRank, &peerHval, &fd), ncclSuccess);
|
|
ASSERT_GE(fd, 0) << "proxy GetFd from rank " << ownerRank << " returned invalid FD";
|
|
|
|
hipMemGenericAllocationHandle_t imported = 0;
|
|
ASSERT_EQ(hipMemImportFromShareableHandle(&imported, (void*)(uintptr_t)fd,
|
|
hipMemHandleTypePosixFileDescriptor),
|
|
hipSuccess);
|
|
close(fd);
|
|
|
|
reserveVaOnly(comm->cudaDev, static_cast<size_t>(meta.size), out);
|
|
ASSERT_EQ(hipMemMap(out->pdev, out->size, 0, imported, 0), hipSuccess);
|
|
applyDeviceRwAccess(comm->cudaDev, out->pdev, out->size);
|
|
out->handle = imported;
|
|
|
|
ASSERT_EQ(ncclMemTrackImportFromPeer(comm->memManager, out->ptr, out->size, imported,
|
|
hipMemHandleTypePosixFileDescriptor, memType,
|
|
ownerRank, meta.ownerDev,
|
|
reinterpret_cast<void*>(meta.ownerPtr)),
|
|
ncclSuccess);
|
|
}
|
|
} // namespace
|
|
|
|
// ===========================================================================
|
|
// AnyRanks suite (np >= 1) - local-only Suspend/Resume scenarios.
|
|
//
|
|
// Every test here only touches its own rank's manager (no cross-rank exports).
|
|
// The bootstrapBarrier(0xBEEF/0xCAFE) inside ncclCommMemSuspend/Resume keeps
|
|
// all ranks in sync, AllGather/Send-Recv in Resume Step 3 short-circuits when
|
|
// nobody has exports, so the same body works identically on np=1, np=2, np=N.
|
|
// ===========================================================================
|
|
|
|
class MemManagerAnyRanks : public MPITestBase
|
|
{
|
|
protected:
|
|
void SetUp() override
|
|
{
|
|
MPITestBase::SetUp();
|
|
if (!validateTestPrerequisites(/*min_processes=*/1, kNoProcessLimit)) {
|
|
GTEST_SKIP() << "AnyRanks suite needs at least 1 rank";
|
|
}
|
|
ASSERT_EQ(createTestCommunicator(), ncclSuccess);
|
|
}
|
|
};
|
|
|
|
TEST_F(MemManagerAnyRanks, SuspendResume_LocalScratch)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm, nullptr);
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
VmmAlloc a{};
|
|
allocateVmmPosixFd(comm->cudaDev, 1 << 20, &a);
|
|
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, a.ptr, a.size, a.handle,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemScratch),
|
|
ncclSuccess);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
|
|
// The single tracked entry is the one we just added.
|
|
ncclDynMemEntry* entry = comm->memManager->entries;
|
|
ASSERT_NE(entry, nullptr);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateReleased);
|
|
EXPECT_EQ(static_cast<void*>(entry->handle), nullptr);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 0);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateActive);
|
|
EXPECT_NE(static_cast<void*>(entry->handle), nullptr);
|
|
|
|
// Verify the VA is mapped again - a roundtrip memset/D2H must succeed.
|
|
ASSERT_EQ(hipMemset(a.ptr, 0xAA, a.size), hipSuccess);
|
|
std::vector<uint8_t> host(a.size, 0);
|
|
ASSERT_EQ(hipMemcpy(host.data(), a.ptr, a.size, hipMemcpyDeviceToHost), hipSuccess);
|
|
EXPECT_EQ(host[0], 0xAA);
|
|
EXPECT_EQ(host[a.size - 1], 0xAA);
|
|
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, a.ptr, a.size), ncclSuccess);
|
|
releaseVmm(&a);
|
|
}
|
|
|
|
TEST_F(MemManagerAnyRanks, SuspendResume_LocalOffload_DataPreserved)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const uint64_t internalOffload = captureInternalOffloadCpuBackup(comm);
|
|
|
|
VmmAlloc a{};
|
|
allocateVmmPosixFd(comm->cudaDev, 64 * 1024, &a);
|
|
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, a.ptr, a.size, a.handle,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemOffload),
|
|
ncclSuccess);
|
|
|
|
std::vector<uint8_t> pattern(a.size);
|
|
for (size_t i = 0; i < a.size; i++) pattern[i] = static_cast<uint8_t>(i * 31u);
|
|
ASSERT_EQ(hipMemcpy(a.ptr, pattern.data(), a.size, hipMemcpyHostToDevice), hipSuccess);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
|
|
ncclDynMemEntry* entry = findEntryByPtr(comm, a.ptr);
|
|
ASSERT_NE(entry, nullptr);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateReleased);
|
|
EXPECT_NE(entry->cpuBackup, nullptr);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, internalOffload + a.size);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateActive);
|
|
EXPECT_EQ(entry->cpuBackup, nullptr);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, 0u);
|
|
|
|
std::vector<uint8_t> roundtrip(a.size, 0);
|
|
ASSERT_EQ(hipMemcpy(roundtrip.data(), a.ptr, a.size, hipMemcpyDeviceToHost), hipSuccess);
|
|
EXPECT_EQ(std::memcmp(roundtrip.data(), pattern.data(), a.size), 0);
|
|
|
|
untrackAndRelease(comm, a.ptr, &a);
|
|
}
|
|
|
|
TEST_F(MemManagerAnyRanks, SuspendResume_PersistUntouched)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
|
|
constexpr size_t persistSize = 8192;
|
|
void* persistPtr = reinterpret_cast<void*>(0xDEADBEEF00ULL);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, persistPtr, persistSize,
|
|
/*handle=*/0, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemPersist),
|
|
ncclSuccess);
|
|
|
|
uint64_t persistBefore = 0;
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemPersist, &persistBefore), ncclSuccess);
|
|
EXPECT_EQ(persistBefore - base.persist, persistSize);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
|
|
uint64_t persistAfter = 0;
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemPersist, &persistAfter), ncclSuccess);
|
|
EXPECT_EQ(persistAfter - base.persist, persistSize);
|
|
|
|
uint64_t suspended = 0;
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspended, &suspended), ncclSuccess);
|
|
EXPECT_EQ(suspended, 1u);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, persistPtr, persistSize), ncclSuccess);
|
|
}
|
|
|
|
TEST_F(MemManagerAnyRanks, SuspendResume_DoubleCalls_Rejected)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemSuspend(comm), ncclInvalidUsage);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(ncclCommMemResume(comm), ncclInvalidUsage);
|
|
}
|
|
|
|
TEST_F(MemManagerAnyRanks, MemStats_FlipsAcrossSuspendResume)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
uint64_t suspended = ~0ULL;
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspended, &suspended), ncclSuccess);
|
|
EXPECT_EQ(suspended, 0u);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspended, &suspended), ncclSuccess);
|
|
EXPECT_EQ(suspended, 1u);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemStats(comm, ncclStatGpuMemSuspended, &suspended), ncclSuccess);
|
|
EXPECT_EQ(suspended, 0u);
|
|
}
|
|
|
|
// Capture and print the memory footprint at three checkpoints: before Suspend,
|
|
// after Suspend, after Resume. Allocates a handful of real VMM buffers as
|
|
// ncclMemOffload so Suspend actually moves bytes to CPU and frees the GPU
|
|
// backing - otherwise the gpu.UsedByProc delta would be invisible.
|
|
//
|
|
// The test is non-strict on absolute numbers (other ranks share the GPU on
|
|
// shared boxes, internal RCCL allocations vary per build) and instead checks
|
|
// the invariants we care about across the cycle:
|
|
// - mgr.Suspended flips 0 -> 1 -> 0.
|
|
// - cpuBackupUsage rises during Suspend, drops to baseline on Resume.
|
|
// - gpu.UsedByProc is lower after Suspend than before, and recovers on
|
|
// Resume to within a small slack of the pre-Suspend value.
|
|
TEST_F(MemManagerAnyRanks, MemoryFootprint_BeforeSuspendAfterSuspendAfterResume)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
// Allocate enough Offload buffers to make the GPU-memory drop on Suspend
|
|
// visible against general device noise (other processes, driver pools).
|
|
constexpr int kNumBuffers = 4;
|
|
constexpr size_t kBufBytes = 4u * 1024u * 1024u; // 4 MiB each, 16 MiB total
|
|
constexpr uint8_t kPattern = 0xC7;
|
|
|
|
std::vector<VmmAlloc> bufs(kNumBuffers);
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
allocateVmmPosixFd(comm->cudaDev, kBufBytes, &bufs[i]);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, bufs[i].ptr, bufs[i].size, bufs[i].handle,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemOffload),
|
|
ncclSuccess);
|
|
ASSERT_EQ(hipMemset(bufs[i].ptr, kPattern, bufs[i].size), hipSuccess);
|
|
}
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
MemFootprint before = captureMemFootprint(comm);
|
|
//printMemFootprint(comm->rank, "BEFORE SUSPEND", before);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
|
|
MemFootprint suspended = captureMemFootprint(comm);
|
|
//printMemFootprint(comm->rank, "AFTER SUSPEND ", suspended);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
MemFootprint resumed = captureMemFootprint(comm);
|
|
//printMemFootprint(comm->rank, "AFTER RESUME ", resumed);
|
|
|
|
EXPECT_EQ(before.mgrSuspended, 0u);
|
|
EXPECT_EQ(suspended.mgrSuspended, 1u);
|
|
EXPECT_EQ(resumed.mgrSuspended, 0u);
|
|
|
|
EXPECT_GE(suspended.cpuBackupUsage,
|
|
before.cpuBackupUsage + static_cast<uint64_t>(kNumBuffers) * kBufBytes)
|
|
<< "Suspend must back up at least our Offload buffers to CPU";
|
|
EXPECT_EQ(resumed.cpuBackupUsage, before.cpuBackupUsage);
|
|
|
|
EXPECT_LT(suspended.gpuUsedByProc, before.gpuUsedByProc)
|
|
<< "Suspend should free GPU backing for Scratch + Offload entries";
|
|
|
|
// Resume must restore the exact pre-Suspend footprint: manager-tracked
|
|
// totals don't change at all across the cycle (entries stay in the list,
|
|
// only `state` toggles), and real GPU usage measured by hipMemGetInfo
|
|
// returns to baseline once cuMemCreate + cuMemMap re-back the released VAs.
|
|
EXPECT_EQ(resumed.mgrTotal, before.mgrTotal);
|
|
EXPECT_EQ(resumed.mgrPersist, before.mgrPersist);
|
|
EXPECT_EQ(resumed.mgrSuspendable, before.mgrSuspendable);
|
|
EXPECT_EQ(resumed.gpuUsedByProc, before.gpuUsedByProc)
|
|
<< "Resume must restore the exact GPU memory footprint";
|
|
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
untrackAndRelease(comm, bufs[i].ptr, &bufs[i]);
|
|
}
|
|
}
|
|
|
|
// ===========================================================================
|
|
// Two-rank suite (-np 2) - exercises bootstrap-coupled paths
|
|
// ===========================================================================
|
|
|
|
class MemManagerTwoRank : public MPITestBase
|
|
{
|
|
protected:
|
|
void SetUp() override
|
|
{
|
|
MPITestBase::SetUp();
|
|
if (!validateTestPrerequisites(/*min_processes=*/2, /*max_processes=*/2)) {
|
|
GTEST_SKIP() << "Two-rank suite, run with mpirun -np 2";
|
|
}
|
|
ASSERT_EQ(createTestCommunicator(), ncclSuccess);
|
|
}
|
|
};
|
|
|
|
// Both ranks own only local entries. No exports, no peer imports.
|
|
// Verifies:
|
|
// - bootstrapBarrier(0xBEEF) on Suspend entry succeeds across ranks.
|
|
// - bootstrapBarrier(0xBEEF) after Resume Step 1 succeeds.
|
|
// - Resume Step 3 short-circuits when nobody has exports
|
|
// (allCounts all zero, totalInfoCount == 0, no Send/Recv).
|
|
// - Final bootstrapBarrier(0xCAFE) succeeds.
|
|
TEST_F(MemManagerTwoRank, SuspendResume_LocalOnly)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
VmmAlloc a{};
|
|
allocateVmmPosixFd(comm->cudaDev, 1 << 20, &a);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, a.ptr, a.size, a.handle,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemScratch),
|
|
ncclSuccess);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 0);
|
|
|
|
// VA must be live again on both ranks.
|
|
ASSERT_EQ(hipMemset(a.ptr, 0x55, a.size), hipSuccess);
|
|
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, a.ptr, a.size), ncclSuccess);
|
|
releaseVmm(&a);
|
|
}
|
|
|
|
// Inject a synthetic peer-imported entry on rank 1 whose ownerPtr has no
|
|
// matching local export on rank 0. Resume Step 4 must surface this as
|
|
// ncclSystemError on rank 1 (RCCL divergence from NCCL upstream silent skip)
|
|
// and leave the entry in Released state. Rank 0 has no entries -> ncclSuccess.
|
|
TEST_F(MemManagerTwoRank, Resume_PeerImport_NoMatch_ReturnsErrorEntryStaysReleased)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
int rank = comm->rank;
|
|
|
|
// Rank 1 reserves a VA region (no physical backing) and registers it as
|
|
// a peer import that points at a fake ownerPtr on rank 0. Rank 0 does
|
|
// nothing - it'll have no exports to broadcast.
|
|
VmmAlloc fakeImport{};
|
|
if (rank == 1) {
|
|
reserveVaOnly(comm->cudaDev, 64 * 1024, &fakeImport);
|
|
void* fakeOwnerPtr = reinterpret_cast<void*>(0xFEEDFACE0000ULL);
|
|
ASSERT_EQ(ncclMemTrackImportFromPeer(
|
|
comm->memManager, fakeImport.ptr, fakeImport.size,
|
|
/*handle=*/0, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemScratch, /*ownerRank=*/0, /*ownerDev=*/0, fakeOwnerPtr),
|
|
ncclSuccess);
|
|
|
|
// Pretend Suspend already ran for this entry so Resume Step 4 considers
|
|
// it a candidate for re-import (which is then expected to no-match).
|
|
ncclDynMemEntry* entry = comm->memManager->entries;
|
|
ASSERT_NE(entry, nullptr);
|
|
entry->state = ncclDynMemStateReleased;
|
|
}
|
|
|
|
// Drive the full Suspend / Resume cycle. We expect:
|
|
// - Suspend on both ranks succeeds (rank 1's "released" entry is skipped
|
|
// in Pass 1 since state already == Released; rank 0 has no entries).
|
|
// - Resume on rank 0 succeeds (no peer-imports to re-map).
|
|
// - Resume on rank 1 returns ncclSystemError because the fake entry has
|
|
// no matching broadcast info; entry stays Released so caller can clean
|
|
// up via Destroy.
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ncclResult_t resumeRet = ncclCommMemResume(comm);
|
|
|
|
if (rank == 0) {
|
|
ASSERT_EQ(resumeRet, ncclSuccess) << "Rank 0 has no entries, Resume must succeed";
|
|
} else {
|
|
ASSERT_EQ(resumeRet, ncclSystemError)
|
|
<< "Rank 1 has a no-match peer-import; Resume must surface the error";
|
|
|
|
ncclDynMemEntry* entry = comm->memManager->entries;
|
|
ASSERT_NE(entry, nullptr);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateReleased)
|
|
<< "Resume must leave a no-match peer entry in Released state";
|
|
|
|
// released flag must remain set so the manager is in a clearly
|
|
// suspended state and the caller knows to retry or Destroy.
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
|
|
// Cleanup: drop the fake entry and free the VA reservation.
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, fakeImport.ptr, fakeImport.size),
|
|
ncclSuccess);
|
|
releaseVmm(&fakeImport);
|
|
}
|
|
}
|
|
|
|
// End-to-end happy path of the Suspend/Resume P2P protocol with a real,
|
|
// mutual peer import driven through the proxy:
|
|
//
|
|
// Setup (before Suspend):
|
|
// rank 0: hipMemCreate(POSIX_FD) -> hipMemMap -> hipMemSetAccess
|
|
// -> ncclMemTrack(..., ncclMemOffload)
|
|
// -> ncclDynMemMarkExportToPeer(peer=1)
|
|
// -> hipMemset(canary)
|
|
// -> bootstrapSend(meta = {handleVal, ownerPtr, size, ownerDev}, tag=0xC0FFEE)
|
|
//
|
|
// rank 1: bootstrapRecv(meta, tag=0xC0FFEE)
|
|
// -> ncclProxyClientGetFdBlocking(rank 0, &handleVal, &fd) <-- UDS round-trip
|
|
// -> hipMemImportFromShareableHandle(fd) -> hipMemMap + SetAccess
|
|
// -> ncclMemTrackImportFromPeer(..., ncclMemOffload, owner=rank 0)
|
|
// -> (sanity) hipMemcpy DToH and verify canary on imported VA
|
|
//
|
|
// Drive Suspend then Resume on both ranks. This exercises:
|
|
// - Suspend Step 1: rank 1 unmaps the imported entry, releases its handle.
|
|
// - Suspend Step 2: rank 0 copies the offload buffer to CPU backup,
|
|
// closes the shareable FD, unmaps + releases the physical handle.
|
|
// - Resume Step 1: rank 0 re-creates a fresh physical handle, re-maps to
|
|
// the same VA, re-applies POSIX-FD path's empty same-process peer-ACL
|
|
// loop (cross-process here, so no cuMemSetAccess), and restores data
|
|
// from the CPU backup. rank 1 has no local entries to restore.
|
|
// - Resume Step 3: AllGather counts (rank 0 reports 1 export, rank 1
|
|
// reports 0), Send/Recv the new ncclDynMemP2pHandleInfo, both ranks
|
|
// end up with allInfos[0] == rank 0's new handle.
|
|
// - Resume Step 4 (positive matching): rank 1 finds the matching entry
|
|
// (ownerRank/ownerPtr/size all match), calls ncclProxyClientGetFdBlocking
|
|
// on the new handle, hipMemImportFromShareableHandle, ncclCuMemMapAndSetAccess
|
|
// to the same imported VA. Entry transitions Released -> Active.
|
|
// - Final 0xCAFE barrier completes.
|
|
//
|
|
// Verify: rank 1 hipMemcpy DToH from imported VA -> canary preserved.
|
|
//
|
|
// Failure of any step short-circuits and asserts so we get a precise diagnosis
|
|
// instead of a silent stale-data success.
|
|
TEST_F(MemManagerTwoRank, SuspendResume_PeerExportImport_RoundtripCanary)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
constexpr size_t kRequestedSize = 1u << 20;
|
|
constexpr uint8_t kCanary = 0xAB;
|
|
constexpr int kMetaTag = 0xC0FFEE;
|
|
|
|
if (comm->rank == 0) {
|
|
VmmAlloc owner{};
|
|
allocateVmmPosixFd(comm->cudaDev, kRequestedSize, &owner);
|
|
ASSERT_EQ(hipMemset(owner.ptr, kCanary, owner.size), hipSuccess);
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
exportAndShipMeta(comm, owner, /*peerRank=*/1, ncclMemOffload, kMetaTag);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 0);
|
|
|
|
untrackAndRelease(comm, owner.ptr, &owner);
|
|
return;
|
|
}
|
|
|
|
ASSERT_EQ(comm->rank, 1);
|
|
VmmAlloc imp{};
|
|
recvAndImportPeerBuffer(comm, /*ownerRank=*/0, ncclMemOffload, kMetaTag, &imp);
|
|
|
|
// Sanity-check the primary import path BEFORE Suspend/Resume - if this
|
|
// fails, the GetFd / import / map chain is broken and any post-Resume
|
|
// verification would be a stale-data success.
|
|
verifyCanaryReadback(imp.ptr, imp.size, kCanary, "primary peer import");
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
{
|
|
ncclDynMemEntry* e = comm->memManager->entries;
|
|
ASSERT_NE(e, nullptr);
|
|
EXPECT_EQ(e->state, ncclDynMemStateReleased);
|
|
EXPECT_EQ(static_cast<void*>(e->handle), nullptr)
|
|
<< "Suspend must release imported handle";
|
|
}
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 0);
|
|
|
|
ncclDynMemEntry* entry = comm->memManager->entries;
|
|
ASSERT_NE(entry, nullptr);
|
|
EXPECT_EQ(entry->state, ncclDynMemStateActive)
|
|
<< "Resume Step 4 must transition the matched peer entry back to Active";
|
|
EXPECT_NE(static_cast<void*>(entry->handle), nullptr)
|
|
<< "Resume must install the re-imported handle";
|
|
|
|
verifyCanaryReadback(imp.ptr, imp.size, kCanary,
|
|
"post-Resume imported VA (rank 0's offload data)");
|
|
|
|
untrackAndRelease(comm, imp.ptr, &imp);
|
|
}
|
|
|
|
// ===========================================================================
|
|
// Stress / coverage tests for AnyRanks (np >= 1)
|
|
// ===========================================================================
|
|
|
|
// A: no user-tracked entries. Suspend/Resume must drive the bootstrap barriers
|
|
// + flip released without inventing or losing entries. Pre-CUMEM the manager
|
|
// was always empty here; with NCCL_CUMEM_ENABLE=1 internal entries exist, so
|
|
// we check that the entry count and cpuBackup are preserved across the cycle.
|
|
TEST_F(MemManagerAnyRanks, EmptyManager_SuspendResume_NoOp)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 1);
|
|
EXPECT_EQ(comm->memManager->numEntries, base.numEntries)
|
|
<< "Suspend must not invent or lose entries";
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->released, 0);
|
|
EXPECT_EQ(comm->memManager->numEntries, base.numEntries)
|
|
<< "Resume must not invent or lose entries";
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, base.cpuBackupUsage);
|
|
}
|
|
|
|
// B: Persist + Scratch + Offload tracked together. Suspend / Resume releases
|
|
// + restores Scratch and Offload but Persist is invisible to it (lives outside
|
|
// the linked list). Verifies stat invariants:
|
|
// - Total = Persist + Scratch + Offload (constant before and after)
|
|
// - Persist (constant)
|
|
// - Suspend = Scratch + Offload (constant - tracked, not freed)
|
|
// - Suspended toggles 0 -> 1 -> 0
|
|
TEST_F(MemManagerAnyRanks, MixedMemTypes_StatsConsistent)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
|
|
constexpr size_t kPersistReq = 8 * 1024;
|
|
constexpr size_t kScratchReq = 256 * 1024;
|
|
constexpr size_t kOffloadReq = 128 * 1024;
|
|
constexpr uint8_t kPersistByte = 0xC1;
|
|
|
|
VmmAlloc persistVmm{};
|
|
allocateVmmPosixFd(comm->cudaDev, kPersistReq, &persistVmm);
|
|
ASSERT_EQ(hipMemset(persistVmm.ptr, kPersistByte, persistVmm.size), hipSuccess);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, persistVmm.ptr, persistVmm.size,
|
|
persistVmm.handle, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemPersist),
|
|
ncclSuccess);
|
|
|
|
VmmAlloc scratchVmm{};
|
|
allocateVmmPosixFd(comm->cudaDev, kScratchReq, &scratchVmm);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, scratchVmm.ptr, scratchVmm.size,
|
|
scratchVmm.handle, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemScratch),
|
|
ncclSuccess);
|
|
|
|
VmmAlloc offloadVmm{};
|
|
allocateVmmPosixFd(comm->cudaDev, kOffloadReq, &offloadVmm);
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, offloadVmm.ptr, offloadVmm.size,
|
|
offloadVmm.handle, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemOffload),
|
|
ncclSuccess);
|
|
|
|
auto readStat = [&](ncclCommMemStat_t s) {
|
|
uint64_t v = ~0ULL;
|
|
EXPECT_EQ(ncclCommMemStats(comm, s, &v), ncclSuccess);
|
|
return v;
|
|
};
|
|
|
|
const uint64_t deltaTotal = persistVmm.size + scratchVmm.size + offloadVmm.size;
|
|
const uint64_t deltaPersist = persistVmm.size;
|
|
const uint64_t deltaDynamic = scratchVmm.size + offloadVmm.size;
|
|
|
|
EXPECT_EQ(readStat(ncclStatGpuMemTotal) - base.total, deltaTotal);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemPersist) - base.persist, deltaPersist);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspend) - base.suspendBytes, deltaDynamic);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspended), 0u);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
|
|
EXPECT_EQ(readStat(ncclStatGpuMemTotal) - base.total, deltaTotal);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemPersist) - base.persist, deltaPersist);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspend) - base.suspendBytes, deltaDynamic);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspended), 1u);
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemTotal) - base.total, deltaTotal);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemPersist) - base.persist, deltaPersist);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspend) - base.suspendBytes, deltaDynamic);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspended), 0u);
|
|
|
|
verifyCanaryReadback(persistVmm.ptr, persistVmm.size, kPersistByte,
|
|
"persist after Suspend/Resume");
|
|
|
|
untrackAndRelease(comm, scratchVmm.ptr, &scratchVmm);
|
|
untrackAndRelease(comm, offloadVmm.ptr, &offloadVmm);
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, persistVmm.ptr, persistVmm.size),
|
|
ncclSuccess);
|
|
releaseVmm(&persistVmm);
|
|
|
|
EXPECT_EQ(readStat(ncclStatGpuMemTotal), base.total);
|
|
EXPECT_EQ(readStat(ncclStatGpuMemPersist), base.persist);
|
|
}
|
|
|
|
// C: multiple Suspend/Resume cycles on the same Offload buffer. Verifies that
|
|
// the CPU-backup -> GPU restore loop is idempotent across many cycles - no
|
|
// cpuBackup leaks, no handle leaks, canary preserved every iteration.
|
|
TEST_F(MemManagerAnyRanks, MultipleSuspendResumeCycles_DataPreserved)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const uint64_t internalOffload = captureInternalOffloadCpuBackup(comm);
|
|
|
|
constexpr size_t kSize = 64 * 1024;
|
|
constexpr int kCycles = 5;
|
|
constexpr uint8_t kInitial = 0x5A;
|
|
|
|
VmmAlloc a{};
|
|
allocateVmmPosixFd(comm->cudaDev, kSize, &a);
|
|
|
|
std::vector<uint8_t> seed(a.size);
|
|
for (size_t i = 0; i < a.size; ++i) {
|
|
seed[i] = static_cast<uint8_t>(kInitial ^ (i & 0xFFu));
|
|
}
|
|
ASSERT_EQ(hipMemcpy(a.ptr, seed.data(), a.size, hipMemcpyHostToDevice), hipSuccess);
|
|
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, a.ptr, a.size, a.handle,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemOffload),
|
|
ncclSuccess);
|
|
|
|
for (int cycle = 0; cycle < kCycles; ++cycle) {
|
|
SCOPED_TRACE("cycle=" + std::to_string(cycle));
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, internalOffload + a.size)
|
|
<< "Suspend must allocate one CPU backup per Offload entry";
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, 0u)
|
|
<< "Resume must free every CPU backup after restore";
|
|
|
|
std::vector<uint8_t> readback(a.size, 0u);
|
|
ASSERT_EQ(hipMemcpy(readback.data(), a.ptr, a.size, hipMemcpyDeviceToHost),
|
|
hipSuccess);
|
|
EXPECT_EQ(std::memcmp(readback.data(), seed.data(), a.size), 0)
|
|
<< "data corrupted after cycle " << cycle;
|
|
}
|
|
|
|
untrackAndRelease(comm, a.ptr, &a);
|
|
}
|
|
|
|
// D: stress - many small Scratch + Offload entries tracked together, single
|
|
// Suspend/Resume cycle. Catches O(n^2) regressions in the linked-list scan,
|
|
// cumulative leaks, granularity-padding bugs, and quietly-skipped entries.
|
|
TEST_F(MemManagerAnyRanks, Suspend_LargeNumberOfEntries_AllRestored)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
const uint64_t internalOffload = captureInternalOffloadCpuBackup(comm);
|
|
|
|
constexpr int kNumEntries = 32;
|
|
constexpr int kNumOffload = 8; // first kNumOffload are Offload, rest are Scratch
|
|
|
|
std::vector<VmmAlloc> entries(kNumEntries);
|
|
std::vector<uint8_t> canaries(kNumEntries);
|
|
uint64_t testDynamicBytes = 0;
|
|
uint64_t testOffloadBytes = 0;
|
|
|
|
for (int i = 0; i < kNumEntries; ++i) {
|
|
size_t requested = (4 * 1024) + (i * 8 * 1024);
|
|
allocateVmmPosixFd(comm->cudaDev, requested, &entries[i]);
|
|
|
|
canaries[i] = static_cast<uint8_t>(0x40 + i);
|
|
ASSERT_EQ(hipMemset(entries[i].ptr, canaries[i], entries[i].size), hipSuccess);
|
|
|
|
bool isOffload = (i < kNumOffload);
|
|
ncclMemType_t type = isOffload ? ncclMemOffload : ncclMemScratch;
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, entries[i].ptr, entries[i].size,
|
|
entries[i].handle,
|
|
hipMemHandleTypePosixFileDescriptor, type),
|
|
ncclSuccess);
|
|
testDynamicBytes += entries[i].size;
|
|
if (isOffload) testOffloadBytes += entries[i].size;
|
|
}
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
EXPECT_EQ(comm->memManager->numEntries - base.numEntries, kNumEntries);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, internalOffload + testOffloadBytes)
|
|
<< "CPU backup pool must equal sum of all Offload sizes (internal + test)";
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, 0u);
|
|
|
|
for (int i = 0; i < kNumEntries; ++i) {
|
|
SCOPED_TRACE("entry=" + std::to_string(i));
|
|
if (i < kNumOffload) {
|
|
verifyCanaryReadback(entries[i].ptr, entries[i].size, canaries[i],
|
|
"offload after Resume");
|
|
} else {
|
|
ASSERT_EQ(hipMemset(entries[i].ptr, 0xAA, entries[i].size), hipSuccess);
|
|
}
|
|
}
|
|
|
|
auto readStat = [&](ncclCommMemStat_t s) {
|
|
uint64_t v = ~0ULL;
|
|
EXPECT_EQ(ncclCommMemStats(comm, s, &v), ncclSuccess);
|
|
return v;
|
|
};
|
|
EXPECT_EQ(readStat(ncclStatGpuMemSuspend) - base.suspendBytes, testDynamicBytes);
|
|
|
|
for (auto& e : entries) untrackAndRelease(comm, e.ptr, &e);
|
|
EXPECT_EQ(comm->memManager->numEntries, base.numEntries);
|
|
}
|
|
|
|
// E: Persist physical memory must survive Suspend/Resume completely untouched
|
|
// (it lives outside the linked list, so neither Pass nor Step iterates over
|
|
// it). PersistUntouched (counter-only) leaves room for a regression where the
|
|
// VA gets accidentally unmapped; this catches it via canary roundtrip.
|
|
TEST_F(MemManagerAnyRanks, Persist_PhysicalMemoryUntouched)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
const uint64_t internalOffload = captureInternalOffloadCpuBackup(comm);
|
|
|
|
constexpr size_t kSize = 32 * 1024;
|
|
constexpr uint8_t kCanary = 0xE7;
|
|
|
|
VmmAlloc persistVmm{};
|
|
allocateVmmPosixFd(comm->cudaDev, kSize, &persistVmm);
|
|
ASSERT_EQ(hipMemset(persistVmm.ptr, kCanary, persistVmm.size), hipSuccess);
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
ASSERT_EQ(ncclMemTrack(comm->memManager, persistVmm.ptr, persistVmm.size,
|
|
persistVmm.handle, hipMemHandleTypePosixFileDescriptor,
|
|
ncclMemPersist),
|
|
ncclSuccess);
|
|
|
|
EXPECT_EQ(comm->memManager->numEntries, base.numEntries)
|
|
<< "Persist must not enter the linked list (only counter)";
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
EXPECT_EQ(comm->memManager->cpuBackupUsage, internalOffload)
|
|
<< "Persist must NOT contribute to CPU backup pool";
|
|
|
|
verifyCanaryReadback(persistVmm.ptr, persistVmm.size, kCanary, "persist mid-suspend");
|
|
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
verifyCanaryReadback(persistVmm.ptr, persistVmm.size, kCanary, "persist post-resume");
|
|
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, persistVmm.ptr, persistVmm.size),
|
|
ncclSuccess);
|
|
releaseVmm(&persistVmm);
|
|
}
|
|
|
|
// ===========================================================================
|
|
// Stress / coverage tests for TwoRank (np == 2)
|
|
// ===========================================================================
|
|
|
|
// F: rank 0 exports N=3 distinct buffers, rank 1 imports all of them, full
|
|
// Suspend/Resume cycle, every canary preserved. Exercises:
|
|
// - Suspend Step 1: rank 1's Pass-1 unmap loop iterates over multiple
|
|
// peer-imported entries.
|
|
// - Suspend Step 2: rank 0's Pass-2 offload loop runs CPU-backup over N
|
|
// distinct entries.
|
|
// - Resume Step 3: rank 0's localBroadcastCount = N, localInfos buffer
|
|
// holds N elements, AllGather counts come out as [N, 0].
|
|
// - Resume Step 4: rank 1's matching loop matches N times in the same
|
|
// allInfos array (positive branch, multiple iterations).
|
|
TEST_F(MemManagerTwoRank, Multiple_PeerEntries_RoundtripCanary)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
constexpr int kNumBuffers = 3;
|
|
constexpr int kBaseTag = 0xC0FFEE;
|
|
const size_t kSizes[kNumBuffers] = { 64 * 1024, 96 * 1024, 128 * 1024 };
|
|
const uint8_t kCanaries[kNumBuffers] = { 0xA1, 0xB2, 0xC3 };
|
|
|
|
if (comm->rank == 0) {
|
|
std::vector<VmmAlloc> owners(kNumBuffers);
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
allocateVmmPosixFd(comm->cudaDev, kSizes[i], &owners[i]);
|
|
ASSERT_EQ(hipMemset(owners[i].ptr, kCanaries[i], owners[i].size), hipSuccess);
|
|
}
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
exportAndShipMeta(comm, owners[i], /*peerRank=*/1,
|
|
ncclMemOffload, kBaseTag + i);
|
|
}
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
for (auto& o : owners) untrackAndRelease(comm, o.ptr, &o);
|
|
return;
|
|
}
|
|
|
|
ASSERT_EQ(comm->rank, 1);
|
|
std::vector<VmmAlloc> imports(kNumBuffers);
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
recvAndImportPeerBuffer(comm, /*ownerRank=*/0, ncclMemOffload, kBaseTag + i,
|
|
&imports[i]);
|
|
verifyCanaryReadback(imports[i].ptr, imports[i].size, kCanaries[i],
|
|
"primary import of buffer #" + std::to_string(i));
|
|
}
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
// Look up each test-owned import explicitly. With NCCL_CUMEM_ENABLE=1 the
|
|
// manager also tracks ~64 internal peer-imports from src/transport/p2p.cc
|
|
// (P2P proxy setup), so we cannot count by `isImportedFromPeer` alone.
|
|
int activeCount = 0;
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
auto* e = findEntryByPtr(comm, imports[i].ptr);
|
|
ASSERT_NE(e, nullptr) << "Import #" << i << " missing from manager";
|
|
EXPECT_TRUE(e->isImportedFromPeer);
|
|
EXPECT_EQ(e->state, ncclDynMemStateActive);
|
|
++activeCount;
|
|
}
|
|
EXPECT_EQ(activeCount, kNumBuffers);
|
|
|
|
for (int i = 0; i < kNumBuffers; ++i) {
|
|
verifyCanaryReadback(imports[i].ptr, imports[i].size, kCanaries[i],
|
|
"post-Resume import of buffer #" + std::to_string(i));
|
|
}
|
|
for (auto& imp : imports) untrackAndRelease(comm, imp.ptr, &imp);
|
|
}
|
|
|
|
// G: bidirectional - both ranks export to each other. Each rank has BOTH a
|
|
// local exported buffer AND an imported one. Exercises Resume Step 3 with
|
|
// localBroadcastCount > 0 on EVERY rank simultaneously: localInfos contains
|
|
// our own export, allInfos receives both ranks' contributions, allCounts
|
|
// comes out as [1, 1] (not [N, 0] like in F).
|
|
TEST_F(MemManagerTwoRank, Bidirectional_PeerExchange_RoundtripCanary)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
int rank = comm->rank;
|
|
int peer = 1 - rank;
|
|
|
|
const MgrBaseline base = captureBaseline(comm);
|
|
|
|
constexpr size_t kSize = 64 * 1024;
|
|
const uint8_t myCanary = (rank == 0) ? 0xD1 : 0xD2;
|
|
const uint8_t peerCanary = (rank == 0) ? 0xD2 : 0xD1;
|
|
const int myExportTag = 0xC0FFEE0 + rank;
|
|
const int peerExportTag = 0xC0FFEE0 + peer;
|
|
|
|
VmmAlloc myExport{};
|
|
allocateVmmPosixFd(comm->cudaDev, kSize, &myExport);
|
|
ASSERT_EQ(hipMemset(myExport.ptr, myCanary, myExport.size), hipSuccess);
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
exportAndShipMeta(comm, myExport, peer, ncclMemOffload, myExportTag);
|
|
|
|
VmmAlloc peerImport{};
|
|
recvAndImportPeerBuffer(comm, peer, ncclMemOffload, peerExportTag, &peerImport);
|
|
verifyCanaryReadback(peerImport.ptr, peerImport.size, peerCanary,
|
|
"primary import of peer's buffer");
|
|
|
|
EXPECT_EQ(comm->memManager->numEntries - base.numEntries, 2)
|
|
<< "Each rank should add 1 local + 1 imported entry on top of baseline";
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
verifyCanaryReadback(peerImport.ptr, peerImport.size, peerCanary,
|
|
"post-Resume import of peer's buffer");
|
|
|
|
untrackAndRelease(comm, myExport.ptr, &myExport);
|
|
untrackAndRelease(comm, peerImport.ptr, &peerImport);
|
|
}
|
|
|
|
// Rank 1 has 1 real peer-import + 1 stale (fake ownerPtr). Per RCCL divergence
|
|
// (mem_manager.cc:905), Step 4 partial-success: real -> Active, stale stays
|
|
// Released, ncclCommMemResume returns ncclSystemError.
|
|
TEST_F(MemManagerTwoRank, PartialPeerImport_MixedMatch)
|
|
{
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
constexpr size_t kSize = 64 * 1024;
|
|
constexpr uint8_t kCanary = 0x9E;
|
|
constexpr int kRealTag = 0xC0FFEE;
|
|
|
|
if (comm->rank == 0) {
|
|
VmmAlloc realOwner{};
|
|
allocateVmmPosixFd(comm->cudaDev, kSize, &realOwner);
|
|
ASSERT_EQ(hipMemset(realOwner.ptr, kCanary, realOwner.size), hipSuccess);
|
|
ASSERT_EQ(hipDeviceSynchronize(), hipSuccess);
|
|
|
|
exportAndShipMeta(comm, realOwner, /*peerRank=*/1, ncclMemOffload, kRealTag);
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
ASSERT_EQ(ncclCommMemResume(comm), ncclSuccess);
|
|
|
|
untrackAndRelease(comm, realOwner.ptr, &realOwner);
|
|
return;
|
|
}
|
|
|
|
ASSERT_EQ(comm->rank, 1);
|
|
VmmAlloc realImport{};
|
|
recvAndImportPeerBuffer(comm, /*ownerRank=*/0, ncclMemOffload, kRealTag, &realImport);
|
|
|
|
// Inject a synthetic "stale" peer-imported entry that points at an
|
|
// ownerPtr rank 0 doesn't have. Step 4 will gather rank 0's only
|
|
// localInfo (matching realImport), the matching loop will then hit our
|
|
// stale entry, fail to find a match, and WARN+skip it.
|
|
VmmAlloc stale{};
|
|
reserveVaOnly(comm->cudaDev, kSize, &stale);
|
|
void* fakeOwnerPtr = reinterpret_cast<void*>(0xDEADBEEF000ULL);
|
|
ASSERT_EQ(ncclMemTrackImportFromPeer(
|
|
comm->memManager, stale.ptr, stale.size, /*handle=*/0,
|
|
hipMemHandleTypePosixFileDescriptor, ncclMemOffload,
|
|
/*ownerRank=*/0, /*ownerDev=*/0, fakeOwnerPtr),
|
|
ncclSuccess);
|
|
|
|
// Pretend Suspend already ran for the stale entry so Step 4 considers it
|
|
// a candidate. realImport will go through the normal Suspend/Resume path.
|
|
for (auto* e = comm->memManager->entries; e != nullptr; e = e->next) {
|
|
if (e->ptr == stale.ptr) {
|
|
e->state = ncclDynMemStateReleased;
|
|
break;
|
|
}
|
|
}
|
|
|
|
ASSERT_EQ(ncclCommMemSuspend(comm), ncclSuccess);
|
|
// Per RCCL divergence: stale's no-match -> ncclSystemError. Real import is
|
|
// still re-imported successfully in the same pass (partial-success).
|
|
EXPECT_EQ(ncclCommMemResume(comm), ncclSystemError)
|
|
<< "Stale no-match peer-import must surface ncclSystemError";
|
|
EXPECT_EQ(comm->memManager->released, 1)
|
|
<< "Partial-failure must leave manager in suspended state";
|
|
|
|
// Look up the two test-owned entries explicitly. With NCCL_CUMEM_ENABLE=1
|
|
// the manager also tracks ~64 internal peer-imports (proxy P2P from
|
|
// src/transport/p2p.cc:325,380) so we cannot count by isImportedFromPeer.
|
|
auto* realEntry = findEntryByPtr(comm, realImport.ptr);
|
|
auto* staleEntry = findEntryByPtr(comm, stale.ptr);
|
|
ASSERT_NE(realEntry, nullptr);
|
|
ASSERT_NE(staleEntry, nullptr);
|
|
EXPECT_EQ(realEntry->state, ncclDynMemStateActive)
|
|
<< "Real import must be re-imported despite stale's failure";
|
|
EXPECT_EQ(staleEntry->state, ncclDynMemStateReleased)
|
|
<< "Stale entry must stay Released (no-match branch)";
|
|
|
|
verifyCanaryReadback(realImport.ptr, realImport.size, kCanary,
|
|
"real import survived mixed-match Resume");
|
|
|
|
ASSERT_EQ(ncclMemUntrack(comm->memManager, stale.ptr, stale.size), ncclSuccess);
|
|
releaseVmm(&stale);
|
|
untrackAndRelease(comm, realImport.ptr, &realImport);
|
|
}
|
|
|
|
// ============================================================================
|
|
// Public API tests (ncclCommSuspend / ncclCommResume / ncclCommMemStats).
|
|
//
|
|
// The internal ncclCommMemSuspend/Resume/Stats path is exercised in detail by
|
|
// the MemManagerAnyRanks / MemManagerTwoRank suites above (peer export/import
|
|
// canary, manager state machine, etc). The suites below target the public API:
|
|
//
|
|
// - ncclCommSuspend(comm, NCCL_SUSPEND_MEM)
|
|
// - ncclCommResume(comm)
|
|
// - ncclCommMemStats(comm, stat, &value)
|
|
//
|
|
// Coverage matrix:
|
|
//
|
|
// Happy paths
|
|
// - SuspendResumeBasic.BasicCycle one Suspend -> Resume round-trip
|
|
// - SuspendResumeBasic.MultipleCycles 5 back-to-back cycles
|
|
// - SuspendResumeBasic.BarrierSync rank 0 sleeps before suspend;
|
|
// every rank should observe the
|
|
// cross-rank bootstrap barrier
|
|
// - SuspendResumeMemStats.BeforeAndAfter ncclStatGpuMem* across the
|
|
// active/suspended/resumed
|
|
// transitions
|
|
// - SuspendResumeMemStats.Conservation total + suspended bytes is
|
|
// conserved across a Suspend
|
|
// - SuspendResumeCollectiveIntegrity.AllReduceAfterResume
|
|
// AllReduce produces correct
|
|
// result after Suspend/Resume
|
|
// - SuspendResumeCollectiveIntegrity.AllReduceTwoCycles
|
|
// same after multiple cycles
|
|
// - SuspendResumeCollectiveIntegrity.AllReduceP2pAutoMark
|
|
// verifies that p2p collectives
|
|
// auto-mark exported buffers via
|
|
// ncclDynMemMarkExportToPeer
|
|
// (Bucket 1)
|
|
// - SuspendResumeLifecycle.DestroyWhileSuspended
|
|
// ncclCommDestroy succeeds on a
|
|
// still-suspended comm (drains
|
|
// pending tasks + manager
|
|
// destruction frees released VAs)
|
|
// - SuspendResumeGroup.AtomicSuspend ncclGroupStart() + Suspend +
|
|
// ncclGroupEnd() drains the
|
|
// pending suspend task
|
|
// - SuspendResumeGroup.AtomicResume same for Resume
|
|
// - SuspendResumeGroup.MixedSuspendResume Suspend + Resume in the same
|
|
// group leaves the comm active
|
|
//
|
|
// Argument validation (no bootstrap barrier reached, so individual ranks can
|
|
// fail without deadlocking peers)
|
|
// - SuspendResumeArgValidation.ZeroFlags Suspend(comm, 0) is rejected
|
|
// - SuspendResumeArgValidation.UnknownFlagsRejected bogus high bits rejected
|
|
// - SuspendResumeArgValidation.DoubleSuspend Suspend twice -> 2nd rejected
|
|
// - SuspendResumeArgValidation.ResumeWhenActive Resume on active comm rejected
|
|
// - SuspendResumeArgValidation.DoubleResume Resume twice -> 2nd rejected
|
|
// - SuspendResumeMemStatsValidation.NullValuePtr NULL output pointer rejected
|
|
// - SuspendResumeMemStatsValidation.UnknownStat bogus enum value rejected
|
|
// ============================================================================
|
|
|
|
namespace SuspendResumeTestConfig
|
|
{
|
|
constexpr int kMinRanks = 2; // multi-rank tests need >= 2
|
|
constexpr size_t kAllReduceCount = 1024; // 4 KiB float buffer
|
|
constexpr float kEpsilon = 1e-3f;
|
|
constexpr int kMultiCycleCount = 5;
|
|
constexpr int kBarrierSleepMs = 200; // rank-0 stagger time
|
|
} // namespace SuspendResumeTestConfig
|
|
|
|
using namespace SuspendResumeTestConfig;
|
|
|
|
/**
|
|
* @class SuspendResumePublicAPITestBase
|
|
* @brief Shared fixture for all public-API suspend/resume MPI tests.
|
|
*
|
|
* Wraps MPITestBase and adds a convenience accessor for ncclCommMemStats and
|
|
* a self-checking AllReduce helper.
|
|
*/
|
|
class SuspendResumePublicAPITestBase : public MPITestBase
|
|
{
|
|
protected:
|
|
// Read one ncclCommMemStat as uint64_t. Fails the test on any non-success
|
|
// return so callers can use the value directly.
|
|
uint64_t readStat(ncclComm_t comm, ncclCommMemStat_t stat)
|
|
{
|
|
uint64_t v = ~uint64_t{0};
|
|
ncclResult_t r = ncclCommMemStats(comm, stat, &v);
|
|
EXPECT_EQ(r, ncclSuccess) << "ncclCommMemStats failed for stat " << (int)stat;
|
|
return v;
|
|
}
|
|
|
|
// Run an AllReduce with float-sum and verify each element equals
|
|
// sum(1..nranks) = nranks * (nranks + 1) / 2. Returns true on success.
|
|
bool runAllReduceAndVerify(ncclComm_t comm, hipStream_t stream)
|
|
{
|
|
const size_t n = kAllReduceCount;
|
|
const int rank = MPIEnvironment::world_rank;
|
|
const int nrnks = MPIEnvironment::world_size;
|
|
|
|
void* sendbuf = nullptr;
|
|
void* recvbuf = nullptr;
|
|
if (hipMalloc(&sendbuf, n * sizeof(float)) != hipSuccess) return false;
|
|
if (hipMalloc(&recvbuf, n * sizeof(float)) != hipSuccess)
|
|
{
|
|
(void)hipFree(sendbuf);
|
|
return false;
|
|
}
|
|
auto sendGuard = makeDeviceBufferAutoGuard(sendbuf);
|
|
auto recvGuard = makeDeviceBufferAutoGuard(recvbuf);
|
|
|
|
if (initializeBufferWithPattern<float>(
|
|
sendbuf, n,
|
|
[rank](size_t) { return static_cast<float>(rank + 1); }) != hipSuccess)
|
|
return false;
|
|
if (zeroInitializeBuffer<float>(recvbuf, n) != hipSuccess) return false;
|
|
|
|
if (ncclAllReduce(sendbuf, recvbuf, n, ncclFloat32, ncclSum, comm, stream)
|
|
!= ncclSuccess)
|
|
return false;
|
|
if (hipStreamSynchronize(stream) != hipSuccess) return false;
|
|
|
|
const float expected = static_cast<float>(nrnks * (nrnks + 1) / 2);
|
|
return verifyBufferData<float>(
|
|
recvbuf, n,
|
|
[expected](size_t) { return expected; },
|
|
0,
|
|
static_cast<double>(kEpsilon * expected));
|
|
}
|
|
};
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 1. Basic happy-path tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeBasic : public SuspendResumePublicAPITestBase {};
|
|
|
|
TEST_F(SuspendResumeBasic, BasicCycle)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "Newly initialized comm must report not-suspended";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
EXPECT_EQ(1u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "After Suspend, ncclStatGpuMemSuspended must be 1";
|
|
|
|
// ncclStatGpuMemSuspend == totalScratch + totalOffload of *tracked* entries
|
|
// With NCCL_CUMEM_ENABLE=0 the comm's internal allocations are not
|
|
// VMM-backed so the count is legitimately zero; with CUMEM=1 it is > 0.
|
|
uint64_t suspBytes = readStat(comm, ncclStatGpuMemSuspend);
|
|
TEST_INFO("Rank %d ncclStatGpuMemSuspend after Suspend = %llu bytes",
|
|
MPIEnvironment::world_rank, (unsigned long long)suspBytes);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "After Resume, ncclStatGpuMemSuspended must be 0";
|
|
}
|
|
|
|
TEST_F(SuspendResumeBasic, MultipleCycles)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
for (int i = 0; i < kMultiCycleCount; ++i)
|
|
{
|
|
TEST_INFO("Cycle %d: Suspend", i);
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
EXPECT_EQ(1u, readStat(comm, ncclStatGpuMemSuspended)) << "cycle " << i;
|
|
|
|
TEST_INFO("Cycle %d: Resume", i);
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended)) << "cycle " << i;
|
|
}
|
|
}
|
|
|
|
TEST_F(SuspendResumeBasic, BarrierSync)
|
|
{
|
|
// Rank 0 stages a deliberate delay before calling Suspend. Without the
|
|
// bootstrapBarrier inside ncclCommMemSuspend, the other ranks would return
|
|
// immediately and the elapsed time on those ranks would be much smaller
|
|
// than rank 0's. With the barrier wired up correctly, every rank must
|
|
// observe at least kBarrierSleepMs of elapsed wall-clock between the
|
|
// (cross-rank) start of the call and its return.
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
MPI_Barrier(MPI_COMM_WORLD);
|
|
auto t0 = std::chrono::steady_clock::now();
|
|
|
|
if (MPIEnvironment::world_rank == 0)
|
|
{
|
|
std::this_thread::sleep_for(std::chrono::milliseconds(kBarrierSleepMs));
|
|
}
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
auto t1 = std::chrono::steady_clock::now();
|
|
|
|
auto elapsed_ms =
|
|
std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
|
|
|
|
// Generous slack: this is a synchronization assertion, not a microbenchmark.
|
|
// Without a barrier, non-zero ranks would typically observe < 10 ms.
|
|
EXPECT_GE(elapsed_ms, kBarrierSleepMs / 2)
|
|
<< "Rank " << MPIEnvironment::world_rank
|
|
<< ": Suspend returned in " << elapsed_ms
|
|
<< " ms (rank 0 staged a " << kBarrierSleepMs
|
|
<< " ms delay; expected the barrier to hold every rank)";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 2. Memory-statistics correctness
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeMemStats : public SuspendResumePublicAPITestBase {};
|
|
|
|
TEST_F(SuspendResumeMemStats, BeforeAndAfter)
|
|
{
|
|
// ncclCommMemStats reads per-type counters on the memory manager.
|
|
// The counters cover *tracked* allocations only. With NCCL_CUMEM_ENABLE=0
|
|
// RCCL skips tracking on most internal allocations so all counters
|
|
// legitimately read zero; the suspended boolean still toggles. With
|
|
// CUMEM=1, counters are non-zero and persist across the suspend boundary
|
|
// (Suspend doesn't change totalScratch/totalOffload -- it only flips
|
|
// released).
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
uint64_t totalActive = readStat(comm, ncclStatGpuMemTotal);
|
|
uint64_t suspActive = readStat(comm, ncclStatGpuMemSuspend);
|
|
uint64_t persActive = readStat(comm, ncclStatGpuMemPersist);
|
|
uint64_t isSusActive = readStat(comm, ncclStatGpuMemSuspended);
|
|
EXPECT_EQ(isSusActive, 0u) << "fresh comm must not be suspended";
|
|
EXPECT_EQ(totalActive, persActive + suspActive)
|
|
<< "Total = persist + (scratch+offload) invariant";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
|
|
uint64_t totalSusp = readStat(comm, ncclStatGpuMemTotal);
|
|
uint64_t suspSusp = readStat(comm, ncclStatGpuMemSuspend);
|
|
uint64_t isSusSusp = readStat(comm, ncclStatGpuMemSuspended);
|
|
EXPECT_EQ(isSusSusp, 1u) << "after Suspend, suspended boolean must be 1";
|
|
EXPECT_EQ(totalSusp, totalActive)
|
|
<< "Suspend does not change tracked byte counters (mirrors NCCL)";
|
|
EXPECT_EQ(suspSusp, suspActive)
|
|
<< "Suspend does not change scratch/offload byte counters";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
uint64_t totalResumed = readStat(comm, ncclStatGpuMemTotal);
|
|
uint64_t isSusResumed = readStat(comm, ncclStatGpuMemSuspended);
|
|
EXPECT_EQ(isSusResumed, 0u);
|
|
EXPECT_EQ(totalResumed, totalActive)
|
|
<< "Resume does not change tracked byte counters either";
|
|
|
|
TEST_INFO("rank %d: total=%llu persist=%llu suspendable=%llu "
|
|
"[scratch+offload] (CUMEM gates whether non-zero)",
|
|
MPIEnvironment::world_rank,
|
|
(unsigned long long)totalActive,
|
|
(unsigned long long)persActive,
|
|
(unsigned long long)suspActive);
|
|
}
|
|
|
|
TEST_F(SuspendResumeMemStats, Conservation)
|
|
{
|
|
// ncclCommMemStats reports per-type byte counters; Suspend doesn't move
|
|
// bytes between persist / scratch / offload buckets, it only flips the
|
|
// released boolean. So Total = Persist + (Scratch+Offload) holds in every
|
|
// state, and Total/Persist/Suspend are constant across a Suspend->Resume
|
|
// round-trip. (CPU backup memory used by ncclMemOffload entries is tracked
|
|
// separately on manager->cpuBackupUsage and not part of the public stats.)
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
uint64_t totalA = readStat(comm, ncclStatGpuMemTotal);
|
|
uint64_t persA = readStat(comm, ncclStatGpuMemPersist);
|
|
uint64_t suspA = readStat(comm, ncclStatGpuMemSuspend);
|
|
EXPECT_EQ(totalA, persA + suspA);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
uint64_t totalS = readStat(comm, ncclStatGpuMemTotal);
|
|
uint64_t persS = readStat(comm, ncclStatGpuMemPersist);
|
|
uint64_t suspS = readStat(comm, ncclStatGpuMemSuspend);
|
|
EXPECT_EQ(totalS, totalA);
|
|
EXPECT_EQ(persS, persA);
|
|
EXPECT_EQ(suspS, suspA);
|
|
EXPECT_EQ(totalS, persS + suspS);
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
uint64_t totalR = readStat(comm, ncclStatGpuMemTotal);
|
|
EXPECT_EQ(totalR, totalA);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 3. Collective integrity across a suspend/resume cycle
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeCollectiveIntegrity : public SuspendResumePublicAPITestBase
|
|
{
|
|
protected:
|
|
void SetUp() override
|
|
{
|
|
SuspendResumePublicAPITestBase::SetUp();
|
|
if (!ncclCuMemEnable()) {
|
|
GTEST_SKIP() << "NCCL_CUMEM_ENABLE=0 — skipped";
|
|
}
|
|
}
|
|
};
|
|
|
|
TEST_F(SuspendResumeCollectiveIntegrity, AllReduceAfterResume)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "AllReduce baseline failed before any Suspend";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "AllReduce produced wrong values after Suspend/Resume cycle";
|
|
}
|
|
|
|
TEST_F(SuspendResumeCollectiveIntegrity, AllReduceTwoCycles)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream)) << "baseline AllReduce";
|
|
|
|
for (int i = 0; i < 2; ++i)
|
|
{
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "AllReduce wrong after cycle " << i;
|
|
}
|
|
}
|
|
|
|
// Regression: ncclP2pImportShareableBuffer (and the intra-PID different-GPU
|
|
// branch in p2pMap) must call ncclMemTrackImportFromPeer so the importer side
|
|
// of every cuMem-backed proxy buffer is visible to ncclMemManager. Without
|
|
// this wiring Resume Step 4 has nothing to re-import: after Suspend the owner
|
|
// rebuilds its physical pages, but the importer's local VA stays mapped to
|
|
// the OLD physical pages, owner/importer drift apart, and the next collective
|
|
// hangs in hipStreamSynchronize waiting on flags that never converge across
|
|
// the diverged buffers.
|
|
//
|
|
// We catch the bug structurally instead of as a hang:
|
|
// 1. Force p2p proxy connections via a warm-up AllReduce.
|
|
// 2. Walk manager->entries and require >= 1 entry with isImportedFromPeer
|
|
// and a non-zero live handle. Pre-fix this set is empty under CUMEM=1;
|
|
// post-fix every cross-PID p2p connection adds one.
|
|
// 3. Suspend/Resume and verify that those imported entries return Active
|
|
// with rebuilt (still non-zero) handles and a non-stale ownerPtr.
|
|
TEST_F(SuspendResumeCollectiveIntegrity, P2pPeerImports_TrackedAndRebuilt)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
ASSERT_NE(comm->memManager, nullptr);
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "warm-up AllReduce must succeed before peer-import inspection";
|
|
|
|
auto countPeerImports = [&](int* nActive, int* nReleased) {
|
|
*nActive = *nReleased = 0;
|
|
std::lock_guard<std::mutex> lk(comm->memManager->lock);
|
|
for (auto* e = comm->memManager->entries; e != nullptr; e = e->next) {
|
|
if (!e->isImportedFromPeer) continue;
|
|
if (e->state == ncclDynMemStateActive) (*nActive)++;
|
|
else if (e->state == ncclDynMemStateReleased) (*nReleased)++;
|
|
}
|
|
};
|
|
|
|
int activeBefore = 0, releasedBefore = 0;
|
|
countPeerImports(&activeBefore, &releasedBefore);
|
|
EXPECT_EQ(releasedBefore, 0)
|
|
<< "no entry should be Released while comm is unsuspended";
|
|
EXPECT_GT(activeBefore, 0)
|
|
<< "ncclP2pImportShareableBuffer must register every cross-PID p2p "
|
|
"import via ncclMemTrackImportFromPeer (regression: 0 imports "
|
|
"tracked => Resume Step 4 is a no-op => post-Resume collectives "
|
|
"hang on stale peer mappings)";
|
|
|
|
if (activeBefore == 0) return;
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
|
|
int activeMid = 0, releasedMid = 0;
|
|
countPeerImports(&activeMid, &releasedMid);
|
|
EXPECT_EQ(activeMid, 0)
|
|
<< "Suspend must move every peer-imported entry to Released";
|
|
EXPECT_EQ(releasedMid, activeBefore)
|
|
<< "Suspend must Release exactly the entries that were Active before";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
int activeAfter = 0, releasedAfter = 0;
|
|
countPeerImports(&activeAfter, &releasedAfter);
|
|
EXPECT_EQ(releasedAfter, 0)
|
|
<< "Resume must rebuild every peer-imported entry";
|
|
EXPECT_EQ(activeAfter, activeBefore)
|
|
<< "Resume must restore the same number of peer-imported entries";
|
|
|
|
{
|
|
std::lock_guard<std::mutex> lk(comm->memManager->lock);
|
|
for (auto* e = comm->memManager->entries; e != nullptr; e = e->next) {
|
|
if (!e->isImportedFromPeer) continue;
|
|
EXPECT_NE(e->ptr, nullptr) << "imported VA must be preserved";
|
|
EXPECT_NE(e->desc.imported.ownerPtr, nullptr)
|
|
<< "ownerPtr must survive Suspend/Resume so Step 4 can match";
|
|
// handle == 0 is legal for the intra-PID different-GPU branch
|
|
// (handle pre-released by p2pMap), so we don't assert non-zero
|
|
// here. The cross-PID branch will have a fresh non-zero handle.
|
|
}
|
|
}
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "post-Resume AllReduce must complete (no hang on stale imports)";
|
|
}
|
|
|
|
TEST_F(SuspendResumeCollectiveIntegrity, AllReduceP2pAutoMark)
|
|
{
|
|
// Bucket 1 invariant: p2pSendProxySetup / p2pRecvProxySetup pass req->peerRank
|
|
// into ncclP2pAllocateShareableBuffer, which in turn calls
|
|
// ncclDynMemMarkExportToPeer on every cuMem-allocated proxy buffer. After
|
|
// Suspend, the manager's Step 3 in ncclCommMemResume re-exports those
|
|
// buffers to the same peer using the recorded peerRank, so the post-Resume
|
|
// AllReduce can rebuild peer mappings without manual MarkExport calls.
|
|
//
|
|
// This test exercises the full path end-to-end: an AllReduce that uses the
|
|
// p2p transport must continue to work after a full Suspend/Resume cycle.
|
|
// If peerRank is lost this test still passes when CUMEM=0
|
|
// (no peer marking happens at all) but exposes the regression
|
|
// as a SIGSEGV during Resume Step 4 with CUMEM=1 + nRanks >= 2.
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
// Warm-up AllReduce so the p2p proxy buffers are actually allocated and
|
|
// marked for export before we suspend.
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream)) << "warm-up AllReduce";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
EXPECT_TRUE(runAllReduceAndVerify(comm, stream))
|
|
<< "AllReduce broken after Suspend/Resume - peer-export marking "
|
|
"regression?";
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 4. Lifecycle: destroy while suspended
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeLifecycle : public SuspendResumePublicAPITestBase
|
|
{
|
|
protected:
|
|
// We tear down our own comm in this test rather than relying on the base
|
|
// class, so that we drive ncclCommDestroy from a suspended state.
|
|
void TearDown() override
|
|
{
|
|
// Skip MPITestBase::TearDown's cleanupTestCommunicator since the test
|
|
// already destroyed it.
|
|
test_comm_ = nullptr;
|
|
test_stream_ = nullptr;
|
|
}
|
|
};
|
|
|
|
TEST_F(SuspendResumeLifecycle, DestroyWhileSuspended)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
hipStream_t stream = getActiveStream();
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
EXPECT_EQ(1u, readStat(comm, ncclStatGpuMemSuspended));
|
|
|
|
// Tear the comm down without resuming. commFree drains pending suspend/resume
|
|
// tasks and ncclMemManagerDestroy reclaims VAs of Released entries, so the
|
|
// destructor walk must not SIGSEGV.
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommDestroy(comm));
|
|
test_comm_ = nullptr;
|
|
if (stream != nullptr)
|
|
{
|
|
(void)hipStreamDestroy(stream);
|
|
test_stream_ = nullptr;
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 5. Group context (Bucket 3): Suspend/Resume inside ncclGroupStart/End
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeGroup : public SuspendResumePublicAPITestBase {};
|
|
|
|
TEST_F(SuspendResumeGroup, AtomicSuspend)
|
|
{
|
|
// ncclCommSuspend always wraps the body in
|
|
// ncclGroupStartInternal/EndInternal, so a user-level
|
|
// ncclGroupStart/Suspend/ncclGroupEnd nesting must still drain the pending
|
|
// suspend task and leave the comm in the suspended state.
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
|
|
EXPECT_EQ(1u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "Suspend inside a group must drain in groupLaunch";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended));
|
|
}
|
|
|
|
TEST_F(SuspendResumeGroup, AtomicResume)
|
|
{
|
|
// Same as AtomicSuspend but for Resume.
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
EXPECT_EQ(1u, readStat(comm, ncclStatGpuMemSuspended));
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "Resume inside a group must drain in groupLaunch";
|
|
}
|
|
|
|
TEST_F(SuspendResumeGroup, MixedSuspendResume)
|
|
{
|
|
// Suspend + Resume in the same group: groupLaunch drains suspendTaskQueue
|
|
// first, then resumeTaskQueue, so the net effect is "active" again.
|
|
// The RCCL pre-check in ncclCommResume_impl accepts this case because at
|
|
// the moment Resume is called there is still a Suspend task pending in
|
|
// suspendTaskQueue (i.e. the comm is "going to be suspended" — not
|
|
// currently active in the sense of the validation).
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended));
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupStart());
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclGroupEnd());
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "After Suspend+Resume in the same group, comm must be active";
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// 6. Argument validation / corner cases
|
|
//
|
|
// Each of these returns from inside ncclCommSuspend/Resume/MemStats BEFORE the
|
|
// bootstrapBarrier is reached, so individual ranks failing the call do not
|
|
// leave their peers blocked on a barrier.
|
|
// ----------------------------------------------------------------------------
|
|
|
|
class SuspendResumeArgValidation : public SuspendResumePublicAPITestBase {};
|
|
|
|
TEST_F(SuspendResumeArgValidation, ZeroFlags)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclInvalidArgument, ncclCommSuspend(comm, 0));
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "rejected Suspend should not transition the comm";
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
}
|
|
|
|
TEST_F(SuspendResumeArgValidation, UnknownFlagsRejected)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclInvalidArgument, ncclCommSuspend(comm, 0xffff));
|
|
ASSERT_MPI_EQ(ncclInvalidArgument,
|
|
ncclCommSuspend(comm, NCCL_SUSPEND_MEM | 0x80));
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
}
|
|
|
|
TEST_F(SuspendResumeArgValidation, DoubleSuspend)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
|
|
// Second suspend rejects on every rank before reaching the barrier
|
|
// (state check sees comm->memManager->released==1).
|
|
ASSERT_MPI_EQ(ncclInvalidUsage, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended));
|
|
}
|
|
|
|
TEST_F(SuspendResumeArgValidation, ResumeWhenActive)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
// Resume on an already-active comm rejects before any barrier.
|
|
ASSERT_MPI_EQ(ncclInvalidUsage, ncclCommResume(comm));
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended))
|
|
<< "rejected Resume should not transition the comm";
|
|
}
|
|
|
|
TEST_F(SuspendResumeArgValidation, DoubleResume)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommSuspend(comm, NCCL_SUSPEND_MEM));
|
|
ASSERT_MPI_EQ(ncclSuccess, ncclCommResume(comm));
|
|
|
|
// Second resume rejects: comm is no longer suspended.
|
|
ASSERT_MPI_EQ(ncclInvalidUsage, ncclCommResume(comm));
|
|
|
|
EXPECT_EQ(0u, readStat(comm, ncclStatGpuMemSuspended));
|
|
}
|
|
|
|
// ----- ncclCommMemStats validation --------------------------------------
|
|
|
|
class SuspendResumeMemStatsValidation : public SuspendResumePublicAPITestBase {};
|
|
|
|
TEST_F(SuspendResumeMemStatsValidation, NullValuePtr)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
EXPECT_EQ(ncclInvalidArgument,
|
|
ncclCommMemStats(comm, ncclStatGpuMemTotal, nullptr));
|
|
EXPECT_EQ(ncclInvalidArgument,
|
|
ncclCommMemStats(comm, ncclStatGpuMemSuspend, nullptr));
|
|
}
|
|
|
|
TEST_F(SuspendResumeMemStatsValidation, UnknownStat)
|
|
{
|
|
ASSERT_TRUE(validateTestPrerequisites(kMinRanks)) << "Need >= 2 ranks";
|
|
ASSERT_MPI_EQ(ncclSuccess, createTestCommunicator());
|
|
ncclComm_t comm = getActiveCommunicator();
|
|
|
|
uint64_t v = 0;
|
|
EXPECT_EQ(ncclInvalidArgument,
|
|
ncclCommMemStats(comm, static_cast<ncclCommMemStat_t>(999), &v));
|
|
EXPECT_EQ(ncclInvalidArgument,
|
|
ncclCommMemStats(comm, static_cast<ncclCommMemStat_t>(-1), &v));
|
|
}
|
|
|
|
#endif // MPI_TESTS_ENABLED
|