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
1484 lines
63 KiB
C++
1484 lines
63 KiB
C++
/*************************************************************************
|
||
* Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
|
||
*
|
||
* See LICENSE.txt for license information
|
||
************************************************************************/
|
||
|
||
#include "NetIbMPITestBase.hpp"
|
||
#include "NetIbCastInspect.hpp"
|
||
#include <initializer_list>
|
||
|
||
#ifdef MPI_TESTS_ENABLED
|
||
|
||
// =============================================================================
|
||
// Test: CastEqualWeightsTwoQPsTokenCounts
|
||
//
|
||
// White-box: request 2 QPs, equal weights (schedWeight=0), WRR+split mode.
|
||
// Actual nqps determined at runtime (ionic caps at 1, mlx5 supports 2+).
|
||
// Verifies: initTokens.totTokens=100, per-QP tokens equal, sum invariant.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastEqualWeightsTwoQPsTokenCounts) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(/*dev=*/0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSize = 1024;
|
||
char sendBuf[kMsgSize], recvBuf[kMsgSize];
|
||
for (size_t i = 0; i < kMsgSize; i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* buf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kMsgSize, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps (NCCL_PARAM cache may cap below requested 2).
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, buf, kMsgSize, 122, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Arm equal-weight tokens.
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSize, 123, mhandle);
|
||
|
||
if (rank == 0)
|
||
EXPECT_EQ(memcmp(sendBuf, recvBuf, kMsgSize), 0) << "data mismatch";
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
ExpectEqualWeightInitTokens(state, actualNqps);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastWeightsDistributionOneRound
|
||
//
|
||
// White-box: Verifies that exactly totTokens sends exhaust all WRR tokens.
|
||
// Two phases on the same connection:
|
||
// Phase 1: equal weights (100/nqps each) — totTokens sends → activeTokens == 0
|
||
// Phase 2: unequal weights (triangular: higher-indexed QPs get fewer tokens)
|
||
// — totTokens sends → activeTokens == 0
|
||
// SetTokens resets both init and active tokens, so no reconnect needed.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastWeightsDistributionOneRound) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
CAST_REQUIRE_UPDATE_INTERVAL_OR_SKIP(10000000);
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr int kTotTokens = 100;
|
||
constexpr size_t kMsgSz = 64;
|
||
|
||
// Buffer sized for kTotTokens messages.
|
||
char sendBuf[kTotTokens * kMsgSz], recvBuf[kTotTokens * kMsgSz];
|
||
for (size_t i = 0; i < sizeof(sendBuf); i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, sizeof(sendBuf), NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps before arming tokens.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, kMsgSz, 199, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Phase 1: equal weights, kTotTokens sends exhaust one full WRR round.
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps, kTotTokens);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm, sendBuf, recvBuf, kMsgSz, kTotTokens, 200, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
ASSERT_TRUE(state.schedInit);
|
||
EXPECT_EQ(state.activeTotTokens, 0) << "equal weights: after one full round active tokens must be 0";
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
|
||
// Phase 2: unequal weights (triangular: token[i] ∝ nqps-i), kTotTokens sends.
|
||
if (rank == 1) {
|
||
const int weightSum = actualNqps * (actualNqps + 1) / 2;
|
||
std::vector<int> tokens(actualNqps);
|
||
int allocated = 0;
|
||
for (int i = 0; i < actualNqps - 1; i++) {
|
||
tokens[i] = kTotTokens * (actualNqps - i) / weightSum;
|
||
allocated += tokens[i];
|
||
}
|
||
tokens[actualNqps - 1] = kTotTokens - allocated;
|
||
// For nqps=1 the sequence is trivially {100}; strictly-decreasing check only for nqps>1.
|
||
if (actualNqps > 1) {
|
||
for (int i = 0; i < actualNqps - 1; i++)
|
||
ASSERT_GT(tokens[i], tokens[i + 1])
|
||
<< "triangular sequence not strictly decreasing for nqps=" << actualNqps;
|
||
}
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm, sendBuf, recvBuf, kMsgSz, kTotTokens, 300, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
ASSERT_TRUE(state.schedInit);
|
||
EXPECT_EQ(state.activeTotTokens, 0) << "unequal weights: after one full round active tokens must be 0";
|
||
}
|
||
|
||
// Forward rank-1 gtest failures to rank 0 so the test's exit code reflects them.
|
||
// (rank 1 has its default gtest listeners detached by main_mpi.cpp.)
|
||
{
|
||
int total_failed = 0; std::string msg;
|
||
if (rank == 1) {
|
||
auto* r = ::testing::UnitTest::GetInstance()->current_test_info()->result();
|
||
for (int i = 0; i < r->total_part_count(); i++) {
|
||
const auto& p = r->GetTestPartResult(i);
|
||
if (p.failed()) { total_failed++; msg += std::string(" - ") + (p.summary() ? p.summary() : "") + "\n"; }
|
||
}
|
||
int len = (int)msg.size();
|
||
MPI_Send(&total_failed, 1, MPI_INT, 0, 9620, MPI_COMM_WORLD);
|
||
MPI_Send(&len, 1, MPI_INT, 0, 9621, MPI_COMM_WORLD);
|
||
if (len > 0) MPI_Send(msg.data(), len, MPI_CHAR, 0, 9622, MPI_COMM_WORLD);
|
||
} else if (rank == 0) {
|
||
int len = 0;
|
||
MPI_Recv(&total_failed, 1, MPI_INT, 1, 9620, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
MPI_Recv(&len, 1, MPI_INT, 1, 9621, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
if (len > 0) {
|
||
msg.assign(len, '\0');
|
||
MPI_Recv(msg.data(), len, MPI_CHAR, 1, 9622, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
ADD_FAILURE() << "rank 1 reported " << total_failed << " failure(s):\n" << msg;
|
||
}
|
||
}
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastTokenSumInvariantAfterConsumption
|
||
//
|
||
// White-box: Equal tokens. After 10 sends: initTokens immutable,
|
||
// sum(activeQpTokens)==activeTotTokens.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastTokenSumInvariantAfterConsumption) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSize = 128;
|
||
constexpr int kNMsgs = 10;
|
||
char sendBuf[kMsgSize * kNMsgs], recvBuf[kMsgSize * kNMsgs];
|
||
for (size_t i = 0; i < sizeof(sendBuf); i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, sizeof(sendBuf), NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, kMsgSize, 699, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm, sendBuf, recvBuf, kMsgSize, kNMsgs, 700, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
ExpectEqualWeightInitTokens(state, actualNqps);
|
||
ExpectActiveTokenSumInvariant(state);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastSingleQPBypassesWrr
|
||
//
|
||
// White-box: nqps=1. WRR must be bypassed (schedInit stays false).
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastSingleQPBypassesWrr) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSize = 256;
|
||
char sendBuf[kMsgSize], recvBuf[kMsgSize];
|
||
for (size_t i = 0; i < kMsgSize; i++) sendBuf[i] = static_cast<char>((i * 3) & 0xFF);
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* buf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kMsgSize, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSize, 800, mhandle);
|
||
|
||
if (rank == 0)
|
||
EXPECT_EQ(memcmp(sendBuf, recvBuf, kMsgSize), 0) << "data mismatch";
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
if (state.nqps <= 1) {
|
||
EXPECT_FALSE(state.schedInit) << "WRR must be bypassed for nqps=1";
|
||
}
|
||
// else: nqps was cached to a higher value; skip the bypass assertion.
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastSchedParmsReflectEnvVars
|
||
//
|
||
// White-box: schedParms inside sendComm must match env vars:
|
||
// enable=true, doWrr=true, splitData=false, splitDataMin from env.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastSchedParmsReflectEnvVars) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
if (GetSplitDataMin() == 0) {
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: library ignores zero, "
|
||
"state.splitDataMin would not match";
|
||
}
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSize = 64;
|
||
char sendBuf[kMsgSize], recvBuf[kMsgSize];
|
||
memset(sendBuf, 0xAB, sizeof(sendBuf));
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* buf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kMsgSize, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSize, 900, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
EXPECT_TRUE(state.schedEnable);
|
||
EXPECT_TRUE(state.doWrr);
|
||
EXPECT_EQ(state.splitDataMin, GetSplitDataMin());
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastCursorWrapsAtNqpsBoundary
|
||
//
|
||
// White-box: actual nqps determined at runtime. Tokens={0,...,0,1} (only last QP).
|
||
// The WRR while(1) loop must skip QP0..nqps-2 and select QP[nqps-1].
|
||
// After selection cursor advances to 0.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastCursorWrapsAtNqpsBoundary) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSize = 128;
|
||
char sendBuf[kMsgSize], recvBuf[kMsgSize];
|
||
memset(sendBuf, 0xCC, kMsgSize);
|
||
memset(recvBuf, 0, kMsgSize);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* buf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kMsgSize, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, buf, kMsgSize, 999, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Set single token on the last QP; all others get 0.
|
||
if (rank == 1) {
|
||
std::vector<int> tokens(actualNqps, 0);
|
||
tokens[actualNqps - 1] = 1;
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// The send must land on QP[nqps-1] and wrap cursor to 0.
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSize, 1000, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
|
||
ASSERT_TRUE(state.schedInit);
|
||
EXPECT_EQ(state.activeTotTokens, 0);
|
||
EXPECT_EQ(state.qpIndex, 0);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastMaxQPCount128
|
||
//
|
||
// Run exactly actualNqps sends — one full WRR round.
|
||
// Verify every QP was selected exactly once: activeQpTokens[i]==0 for all i,
|
||
// activeTotTokens==0, and qpIndex==0 (cursor wrapped back to start).
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastMaxQPCount128) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSz = 32;
|
||
constexpr size_t kBufSz = (NCCL_IB_MAX_QPS + 1) * kMsgSz;
|
||
constexpr int kBaseTag = 1300;
|
||
|
||
std::vector<char> sendBuf(kBufSz, 0);
|
||
std::vector<char> recvBuf(kBufSz, 0);
|
||
for (size_t i = 0; i < kBufSz; i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kBufSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps (driver caps 128 at hardware limit).
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, kMsgSz, 999, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Set 1 token per QP — each QP selected exactly once across kNMsgs sends.
|
||
if (rank == 1) {
|
||
std::vector<int> tokens(actualNqps, 1);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// Post all actualNqps sends/recvs concurrently (skip slot 0 used by warmup).
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm,
|
||
sendBuf.data() + kMsgSz, recvBuf.data() + kMsgSz,
|
||
kMsgSz, actualNqps, kBaseTag, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
|
||
ASSERT_TRUE(state.schedInit);
|
||
for (int i = 0; i < state.nqps; i++)
|
||
EXPECT_EQ(state.initQpTokens[i], 1) << "QP " << i << " initToken must be 1";
|
||
EXPECT_EQ(state.activeTotTokens, 0);
|
||
for (int i = 0; i < state.nqps; i++)
|
||
EXPECT_EQ(state.activeQpTokens[i], 0) << "QP " << i << " activeToken must be 0";
|
||
EXPECT_EQ(state.qpIndex, 0);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastFourQPsMonotonicOrder
|
||
//
|
||
// White-box: request 4 QPs; actual nqps determined at runtime.
|
||
// Assign tokens as a strictly-decreasing sequence summing to 100 using
|
||
// triangular proportions: token[i] = round(100 * (nqps-i) / weightSum),
|
||
// where weightSum = nqps*(nqps+1)/2; the last token absorbs rounding residual.
|
||
// For nqps=1: sequence is trivially {100}; monotonic check is skipped.
|
||
// Run exactly 100 sends (= totTokens). Verify:
|
||
// - initQpTokens are strictly decreasing (nqps>1)
|
||
// - activeTotTokens == 0 after exactly totTokens sends
|
||
// - activeQpTokens[i] == 0 for all i
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastFourQPsMonotonicOrder) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr int kNMsgs = 100;
|
||
constexpr size_t kMsgSz = 32;
|
||
constexpr size_t kBufSz = (kNMsgs + 1) * kMsgSz;
|
||
constexpr int kBaseTag = 1500;
|
||
|
||
std::vector<char> sendBuf(kBufSz, 0);
|
||
std::vector<char> recvBuf(kBufSz, 0);
|
||
for (size_t i = 0; i < kBufSz; i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kBufSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, kMsgSz, 998, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Build strictly-decreasing token sequence summing to kNMsgs.
|
||
if (rank == 1) {
|
||
const int weightSum = actualNqps * (actualNqps + 1) / 2;
|
||
std::vector<int> tokens(actualNqps);
|
||
int allocated = 0;
|
||
for (int i = 0; i < actualNqps - 1; i++) {
|
||
tokens[i] = kNMsgs * (actualNqps - i) / weightSum;
|
||
allocated += tokens[i];
|
||
}
|
||
tokens[actualNqps - 1] = kNMsgs - allocated;
|
||
|
||
if (actualNqps > 1) {
|
||
bool isDecreasing = true;
|
||
for (int i = 0; i < actualNqps - 1; i++)
|
||
if (tokens[i] <= tokens[i + 1]) { isDecreasing = false; break; }
|
||
ASSERT_TRUE(isDecreasing)
|
||
<< "Generated token sequence is not strictly decreasing for nqps=" << actualNqps;
|
||
}
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// Verify initTokens immediately after SetTokens (before timer can overwrite).
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState probe = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &probe), ncclSuccess);
|
||
ASSERT_TRUE(probe.schedInit);
|
||
if (actualNqps > 1) {
|
||
for (int i = 0; i < probe.nqps - 1; i++)
|
||
EXPECT_GT(probe.initQpTokens[i], probe.initQpTokens[i + 1])
|
||
<< "initTokens not strictly decreasing at index " << i;
|
||
}
|
||
}
|
||
|
||
// Phase 1: exactly kNMsgs sends = one full WRR round.
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm,
|
||
sendBuf.data() + kMsgSz, recvBuf.data() + kMsgSz,
|
||
kMsgSz, kNMsgs, kBaseTag, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState state = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &state), ncclSuccess);
|
||
ASSERT_TRUE(state.schedInit);
|
||
EXPECT_EQ(state.activeTotTokens, 0);
|
||
for (int i = 0; i < state.nqps; i++)
|
||
EXPECT_EQ(state.activeQpTokens[i], 0)
|
||
<< "QP " << i << " activeToken must be 0 after full round";
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastSplitDataThresholdBoundary
|
||
//
|
||
// White-box: splitDataMin from env. Per-QP threshold = splitDataMin * nqps.
|
||
// size >= threshold → split path → 0 WRR tokens consumed
|
||
// size < threshold → WRR path → 1 WRR token consumed
|
||
// Actual nqps determines threshold dynamically.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastSplitDataThresholdBoundary) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
{
|
||
const char* v = getenv("NCCL_IB_SPLIT_DATA_ON_QPS");
|
||
if (!v || strcmp(v, "1") != 0)
|
||
GTEST_SKIP() << "CastSplitDataThresholdBoundary requires NCCL_IB_SPLIT_DATA_ON_QPS=1";
|
||
}
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
|
||
// Phase 0: small warmup buffer to learn actualNqps before allocating the real buffer.
|
||
constexpr size_t kWarmupSz = 128;
|
||
std::vector<char> warmupBuf(kWarmupSz, 0);
|
||
void* warmupBase = warmupBuf.data();
|
||
void* warmupHandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, warmupBase, kWarmupSz, NCCL_PTR_HOST, &warmupHandle), ncclSuccess);
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, warmupBase, 64, 1599, warmupHandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
ASSERT_EQ(DeregisterMemory(comm, warmupHandle), ncclSuccess);
|
||
|
||
// threshold = splitDataMin * nqps (from net_ib_cast.cc: dataPerQp = size*nreqs/nqps)
|
||
const size_t kSplitDataMin = GetSplitDataMin();
|
||
if (kSplitDataMin == 0) {
|
||
// With splitDataMin=0 every message goes through the split path regardless of size.
|
||
// There is no WRR boundary to test, so skip rather than crash on zero-size buffer.
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: no WRR/split boundary exists";
|
||
}
|
||
const size_t kThreshold = kSplitDataMin * static_cast<size_t>(actualNqps);
|
||
const size_t kSplitSz = kThreshold; // dataPerQp == splitDataMin → split
|
||
const size_t kWrrSz = kThreshold - 1; // dataPerQp < splitDataMin → WRR
|
||
|
||
// Allocate and register the real buffer sized for the actual threshold.
|
||
const size_t kBufSz = kThreshold;
|
||
std::vector<char> sendBuf(kBufSz, 0x5A);
|
||
std::vector<char> recvBuf(kBufSz, 0x00);
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kBufSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// Large send: split path — activeTotTokens must NOT change.
|
||
int activeTotBeforeSplit = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBeforeSplit = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kSplitSz, 1600, mhandle);
|
||
|
||
int activeTotAfterSplit = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(st.activeTotTokens, activeTotBeforeSplit)
|
||
<< "split path must not consume WRR tokens";
|
||
activeTotAfterSplit = st.activeTotTokens;
|
||
}
|
||
|
||
// Small send: WRR path — activeTotTokens must decrease by exactly 1.
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kWrrSz, 1601, mhandle);
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
int delta = activeTotAfterSplit - st.activeTotTokens;
|
||
EXPECT_EQ(delta, 1) << "WRR path must consume exactly one token";
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastAlternatingWrrNonWrr
|
||
//
|
||
// White-box: Toggle doWrr mid-test on an established connection.
|
||
// Phase 1 (doWrr=true): 10 sends → 10 WRR tokens consumed
|
||
// Phase 2 (doWrr=false): 10 sends → 0 WRR tokens consumed
|
||
// Phase 3 (doWrr=true): 10 sends → 10 WRR tokens consumed
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastAlternatingWrrNonWrr) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
CAST_REQUIRE_UPDATE_INTERVAL_OR_SKIP(10000000);
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr int kPhase = 10;
|
||
constexpr size_t kMsgSz = 64;
|
||
constexpr int kBaseTag = 1700;
|
||
|
||
char sendBuf[kPhase * 3 * kMsgSz], recvBuf[kPhase * 3 * kMsgSz];
|
||
memset(sendBuf, 0xAA, sizeof(sendBuf));
|
||
memset(recvBuf, 0, sizeof(recvBuf));
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf) : static_cast<void*>(sendBuf);
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, sizeof(sendBuf), NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, kMsgSz, 1699, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
auto doPhase = [&](int phaseBase) {
|
||
CastDoBatchSendRecv(rank, sendComm, recvComm,
|
||
sendBuf + phaseBase * kMsgSz, recvBuf + phaseBase * kMsgSz,
|
||
kMsgSz, kPhase, kBaseTag + phaseBase, mhandle);
|
||
};
|
||
|
||
// Phase 1: WRR on
|
||
int activeTotBefore1 = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore1 = st.activeTotTokens;
|
||
}
|
||
doPhase(0);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
int consumed = activeTotBefore1 - st.activeTotTokens;
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(consumed, kPhase) << "phase 1: expected " << kPhase << " WRR selections";
|
||
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, false, false, GetSplitDataMin()), ncclSuccess);
|
||
}
|
||
|
||
int activeTotBefore2 = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore2 = st.activeTotTokens;
|
||
}
|
||
doPhase(kPhase);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
int consumed = activeTotBefore2 - st.activeTotTokens;
|
||
EXPECT_EQ(consumed, 0) << "phase 2: doWrr=false must not consume WRR tokens";
|
||
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, true, false, GetSplitDataMin()), ncclSuccess);
|
||
}
|
||
|
||
int activeTotBefore3 = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore3 = st.activeTotTokens;
|
||
}
|
||
doPhase(kPhase * 2);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
int consumed = activeTotBefore3 - st.activeTotTokens;
|
||
EXPECT_EQ(consumed, kPhase) << "phase 3: expected " << kPhase << " WRR selections";
|
||
}
|
||
|
||
// Forward rank-1 gtest failures to rank 0 (default listeners detached on non-zero ranks).
|
||
{
|
||
int total_failed = 0; std::string msg;
|
||
if (rank == 1) {
|
||
auto* r = ::testing::UnitTest::GetInstance()->current_test_info()->result();
|
||
for (int i = 0; i < r->total_part_count(); i++) {
|
||
const auto& p = r->GetTestPartResult(i);
|
||
if (p.failed()) { total_failed++; msg += std::string(" - ") + (p.summary() ? p.summary() : "") + "\n"; }
|
||
}
|
||
int len = (int)msg.size();
|
||
MPI_Send(&total_failed, 1, MPI_INT, 0, 9630, MPI_COMM_WORLD);
|
||
MPI_Send(&len, 1, MPI_INT, 0, 9631, MPI_COMM_WORLD);
|
||
if (len > 0) MPI_Send(msg.data(), len, MPI_CHAR, 0, 9632, MPI_COMM_WORLD);
|
||
} else if (rank == 0) {
|
||
int len = 0;
|
||
MPI_Recv(&total_failed, 1, MPI_INT, 1, 9630, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
MPI_Recv(&len, 1, MPI_INT, 1, 9631, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
if (len > 0) {
|
||
msg.assign(len, '\0');
|
||
MPI_Recv(msg.data(), len, MPI_CHAR, 1, 9632, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
ADD_FAILURE() << "rank 1 reported " << total_failed << " failure(s):\n" << msg;
|
||
}
|
||
}
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastEnableDisableSplitData
|
||
//
|
||
// White-box: Toggle splitData on an established connection.
|
||
// Message size = splitDataMin * nqps (at the per-QP threshold boundary).
|
||
// Phase 1 (splitData=false): oneQp WRR → 1 token consumed
|
||
// Phase 2 (splitData=true): split path → 0 tokens consumed
|
||
// Phase 3 (splitData=false): oneQp WRR → 1 token consumed
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastEnableDisableSplitData) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
CAST_REQUIRE_UPDATE_INTERVAL_OR_SKIP(10000000);
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
|
||
// Phase 0: small warmup buffer to learn actualNqps before allocating the real buffer.
|
||
constexpr size_t kWarmupSz = 128;
|
||
std::vector<char> warmupBuf(kWarmupSz, 0);
|
||
void* warmupBase = warmupBuf.data();
|
||
void* warmupHandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, warmupBase, kWarmupSz, NCCL_PTR_HOST, &warmupHandle), ncclSuccess);
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, warmupBase, 64, 1799, warmupHandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
ASSERT_EQ(DeregisterMemory(comm, warmupHandle), ncclSuccess);
|
||
|
||
// Message at the split threshold: dataPerQp = splitDataMin → split when splitData=true.
|
||
const size_t kSplitDataMin = GetSplitDataMin();
|
||
if (kSplitDataMin == 0) {
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: no WRR/split boundary exists";
|
||
}
|
||
const size_t kMsgSz = kSplitDataMin * static_cast<size_t>(actualNqps);
|
||
|
||
// Allocate and register the real buffer sized for the actual threshold.
|
||
std::vector<char> sendBuf(kMsgSz, 0xBC);
|
||
std::vector<char> recvBuf(kMsgSz, 0x00);
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kMsgSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, true, false, kSplitDataMin), ncclSuccess);
|
||
}
|
||
|
||
// Phase 1: splitData=false → oneQp WRR → 1 token consumed
|
||
int activeTotBefore1ed = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore1ed = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kMsgSz, 1800, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(activeTotBefore1ed - st.activeTotTokens, 1);
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, true, true, kSplitDataMin), ncclSuccess);
|
||
}
|
||
|
||
// Phase 2: splitData=true → split path → 0 tokens consumed
|
||
int activeTotBefore2ed = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore2ed = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kMsgSz, 1801, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(activeTotBefore2ed - st.activeTotTokens, 0)
|
||
<< "split path must not consume tokens";
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, true, false, kSplitDataMin), ncclSuccess);
|
||
}
|
||
|
||
// Phase 3: splitData=false → oneQp WRR → 1 token consumed
|
||
int activeTotBefore3ed = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore3ed = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kMsgSz, 1802, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(activeTotBefore3ed - st.activeTotTokens, 1);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastEnableDisableSched
|
||
//
|
||
// White-box: Toggle schedEnable on an established connection.
|
||
// Message is small (dataPerQp < splitDataMin) with splitData=true.
|
||
// enable=true → WRR oneQp path → 1 token consumed
|
||
// enable=false → scheduler disabled, all-QP path → 0 tokens consumed
|
||
// enable=true → WRR resumes → 1 token consumed
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastEnableDisableSched) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
CAST_REQUIRE_UPDATE_INTERVAL_OR_SKIP(10000000);
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
// Must be strictly below splitDataMin so dataPerQp < splitDataMin for any nqps,
|
||
// ensuring messages take the WRR path (not the split path).
|
||
if (GetSplitDataMin() == 0) {
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: no WRR/split boundary exists";
|
||
}
|
||
const size_t kMsgSz = std::max<size_t>(64, GetSplitDataMin() - 1);
|
||
std::vector<char> sendBuf(kMsgSz, 0xCD);
|
||
std::vector<char> recvBuf(kMsgSz, 0);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* buf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kMsgSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, buf, kMsgSz, 1899, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// Phase 1: enable=true → WRR path
|
||
int activeTotBefore1es = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore1es = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSz, 1900, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(activeTotBefore1es - st.activeTotTokens, 1);
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, false, true, true, GetSplitDataMin()), ncclSuccess);
|
||
}
|
||
|
||
// Phase 2: enable=false → bypass (all-QP path)
|
||
int activeTotBefore2es = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore2es = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSz, 1901, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(activeTotBefore2es - st.activeTotTokens, 0)
|
||
<< "enable=false must not consume WRR tokens";
|
||
ASSERT_EQ(ncclIbCastSetSchedParms(sendComm, true, true, true, GetSplitDataMin()), ncclSuccess);
|
||
}
|
||
|
||
// Phase 3: enable=true → WRR resumes
|
||
int activeTotBefore3es = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBefore3es = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, buf, kMsgSz, 1902, mhandle);
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(activeTotBefore3es - st.activeTotTokens, 1);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastSendRecvMultipleSizes
|
||
//
|
||
// splitData=true, splitDataMin from env. Threshold = splitDataMin * nqps.
|
||
// WRR sizes (< threshold): each consumes 1 WRR token.
|
||
// Split sizes (>= threshold): 0 WRR tokens consumed.
|
||
// Data integrity verified for all sizes.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastSendRecvMultipleSizes) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
{
|
||
const char* v = getenv("NCCL_IB_SPLIT_DATA_ON_QPS");
|
||
if (!v || strcmp(v, "1") != 0)
|
||
GTEST_SKIP() << "CastSendRecvMultipleSizes requires NCCL_IB_SPLIT_DATA_ON_QPS=1";
|
||
}
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kBufSz = 524288; // 512 KB — fits all test sizes
|
||
std::vector<char> sendBuf(kBufSz);
|
||
std::vector<char> recvBuf(kBufSz);
|
||
for (size_t i = 0; i < kBufSz; i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kBufSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps to compute split threshold.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, 64, 1999, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
const size_t kSplitDataMin = GetSplitDataMin();
|
||
if (kSplitDataMin == 0) {
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: no WRR/split boundary exists";
|
||
}
|
||
const size_t kThreshold = kSplitDataMin * static_cast<size_t>(actualNqps);
|
||
ASSERT_LE(kThreshold, kBufSz) << "buffer too small for threshold=" << kThreshold;
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
// Sizes below threshold → WRR path (1 token consumed each).
|
||
// All entries must be strictly below kThreshold; filter out any that aren't
|
||
// (e.g. 512 and 4096 may exceed kThreshold when splitDataMin is small).
|
||
std::vector<size_t> kWrrSizes;
|
||
for (size_t sz : std::initializer_list<size_t>{512, 4096, kSplitDataMin, kThreshold - 1}) {
|
||
if (sz > 0 && sz < kThreshold && sz <= kBufSz)
|
||
kWrrSizes.push_back(sz);
|
||
}
|
||
// Deduplicate (kSplitDataMin or 512 might equal kThreshold-1).
|
||
kWrrSizes.erase(std::unique(kWrrSizes.begin(), kWrrSizes.end()), kWrrSizes.end());
|
||
// Sizes at/above threshold → split path (0 tokens consumed).
|
||
const std::vector<size_t> kSplitSizes = {kThreshold, kThreshold * 2};
|
||
int baseTag = 2000;
|
||
|
||
for (size_t sz : kWrrSizes) {
|
||
if (sz > kBufSz) continue; // skip if buffer too small (shouldn't happen)
|
||
int prevActiveTot = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
prevActiveTot = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, sz, baseTag++, mhandle);
|
||
if (rank == 0)
|
||
EXPECT_EQ(memcmp(sendBuf.data(), recvBuf.data(), sz), 0) << "data mismatch at size " << sz;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(prevActiveTot - st.activeTotTokens, 1)
|
||
<< "WRR path must consume exactly 1 token for size=" << sz;
|
||
}
|
||
}
|
||
|
||
for (size_t sz : kSplitSizes) {
|
||
if (sz > kBufSz) continue;
|
||
int prevActiveTot = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
prevActiveTot = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, sz, baseTag++, mhandle);
|
||
if (rank == 0)
|
||
EXPECT_EQ(memcmp(sendBuf.data(), recvBuf.data(), sz), 0) << "data mismatch at size " << sz;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
EXPECT_EQ(prevActiveTot - st.activeTotTokens, 0)
|
||
<< "split path must not consume WRR tokens for size=" << sz;
|
||
}
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastLargeTransfer
|
||
//
|
||
// Black/white-box: 16 MB transfer with splitData=true.
|
||
// dataPerQp = 16MB / nqps >> splitDataMin → split path taken.
|
||
// Verify: data integrity + 0 WRR tokens consumed.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastLargeTransfer) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
{
|
||
const char* v = getenv("NCCL_IB_SPLIT_DATA_ON_QPS");
|
||
if (!v || strcmp(v, "1") != 0)
|
||
GTEST_SKIP() << "CastLargeTransfer requires NCCL_IB_SPLIT_DATA_ON_QPS=1";
|
||
}
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kMsgSz = kLargeBufferSize; // 16 MB
|
||
std::vector<char> sendBuf(kMsgSz);
|
||
std::vector<char> recvBuf(kMsgSz, 0);
|
||
for (size_t i = 0; i < kMsgSz; i++) sendBuf[i] = static_cast<char>(i & 0xFF);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* baseBuf = (rank == 0) ? static_cast<void*>(recvBuf.data())
|
||
: static_cast<void*>(sendBuf.data());
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, baseBuf, kMsgSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, baseBuf, 64, 2099, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
int activeTotBeforeLarge = 0;
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
activeTotBeforeLarge = st.activeTotTokens;
|
||
}
|
||
CastDoSendRecv(rank, sendComm, recvComm, baseBuf, kMsgSz, 2100, mhandle);
|
||
|
||
if (rank == 0)
|
||
EXPECT_EQ(memcmp(sendBuf.data(), recvBuf.data(), kMsgSz), 0) << "data mismatch";
|
||
|
||
if (rank == 1) {
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(activeTotBeforeLarge - st.activeTotTokens, 0)
|
||
<< "16 MB transfer must use split path, not WRR";
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastSendRecvZeroSize
|
||
//
|
||
// White-box: Zero-byte send with splitData=true.
|
||
// dataPerQp = 0 < splitDataMin → oneQp WRR path (regardless of nqps).
|
||
// Verify: send/recv complete with ncclSuccess, received size=0, 1 WRR token consumed.
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastSendRecvZeroSize) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
void* listenComm = nullptr;
|
||
void* sendComm = nullptr;
|
||
void* recvComm = nullptr;
|
||
SetupCastConnection(0, &listenComm, &sendComm, &recvComm);
|
||
|
||
constexpr size_t kRegSz = 64;
|
||
char buf[kRegSz];
|
||
memset(buf, 0, kRegSz);
|
||
|
||
void* comm = (rank == 0) ? recvComm : sendComm;
|
||
void* mhandle = nullptr;
|
||
ASSERT_EQ(RegisterMemory(comm, buf, kRegSz, NCCL_PTR_HOST, &mhandle), ncclSuccess);
|
||
|
||
// Warmup: learn real nqps.
|
||
const int actualNqps = GetActualNqps(sendComm, recvComm, buf, kRegSz, 2199, mhandle);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
if (rank == 1) {
|
||
const std::vector<int> tokens = EqualTokens(actualNqps);
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComm, tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
void* req = nullptr;
|
||
int sz = -1;
|
||
if (rank == 0) {
|
||
void* bufs[1] = {buf};
|
||
size_t sizes[1] = {0};
|
||
int tags[1] = {2200};
|
||
void* handles[1] = {mhandle};
|
||
ASSERT_EQ(PostRecv(recvComm, 1, bufs, sizes, tags, handles, &req), ncclSuccess);
|
||
ASSERT_NE(req, nullptr);
|
||
ASSERT_EQ(WaitForCompletion(req, &sz, 10000), ncclSuccess);
|
||
EXPECT_EQ(sz, 0) << "received size must be 0";
|
||
} else {
|
||
PostSendWithRetry(sendComm, buf, 0, 2200, mhandle, &req);
|
||
ASSERT_EQ(WaitForCompletion(req, &sz, 10000), ncclSuccess);
|
||
|
||
struct ncclIbCastSchedState st = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComm, &st), ncclSuccess);
|
||
ASSERT_TRUE(st.schedInit);
|
||
EXPECT_EQ(st.initTotTokens - st.activeTotTokens, 1)
|
||
<< "zero-size send must still trigger WRR (dataPerQp=0 < splitDataMin)";
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
TeardownConnection(recvComm, listenComm, sendComm, mhandle);
|
||
}
|
||
|
||
// =============================================================================
|
||
// Test: CastStressMultiRoundTwoConns
|
||
//
|
||
// Stress test: 100 independent connections running concurrently, 5000 sends each
|
||
// (WRR small-message phase), followed by a ramping-size phase that exercises the
|
||
// WRR/split boundary across {splitDataMin/4, /2, *1, *4, *16} with 20 rounds each.
|
||
// Actual nqps determined at runtime. Works with nqps=1 (no WRR timer check) or
|
||
// nqps>1 (RTT timer fires and rewrites initTokens from asymmetric initial values).
|
||
//
|
||
// Invariants asserted:
|
||
// - schedInit == true on conn[0].
|
||
// - sum(activeQpTokens[i]) == activeTotTokens for conn[0].
|
||
// - 0 <= activeTotTokens <= initTotTokens for conn[0].
|
||
// - If nqps > 1: initQpTokens changed from asymmetric initial values (timer fired).
|
||
// =============================================================================
|
||
TEST_F(NetIbMPITest, CastStressMultiRoundTwoConns) {
|
||
ASSERT_TRUE(validateTestPrerequisites(kExactTwoProcesses, kExactTwoProcesses,
|
||
false, kMinGpusPerNode, kNoNodeLimit))
|
||
<< "Test requires exactly " << kExactTwoProcesses << " processes";
|
||
|
||
const int rank = MPIEnvironment::world_rank;
|
||
|
||
CAST_ENV_CHECK_OR_SKIP();
|
||
net_ = &netIbCast;
|
||
AssertInitAndGetDevices(nullptr);
|
||
|
||
constexpr int kNConns = 100;
|
||
std::vector<void*> listenComms(kNConns, nullptr);
|
||
std::vector<void*> sendComms(kNConns, nullptr);
|
||
std::vector<void*> recvComms(kNConns, nullptr);
|
||
for (int c = 0; c < kNConns; c++)
|
||
SetupCastConnection(/*dev=*/0, &listenComms[c], &sendComms[c], &recvComms[c]);
|
||
|
||
// Scale msgs per connection inversely with connection count so total work stays constant.
|
||
constexpr int kNMsgsTotal = 10000;
|
||
constexpr int kNMsgs = kNMsgsTotal / kNConns; // 100 msgs × 100 conns = 10000 total
|
||
if (GetSplitDataMin() == 0) {
|
||
GTEST_SKIP() << "RCCL_IB_QP_SCHED_SPLIT_DATA_MIN=0: no WRR/split boundary exists";
|
||
}
|
||
const size_t kMsgSz = std::max<size_t>(64, GetSplitDataMin() - 1);
|
||
const size_t kBufSz = static_cast<size_t>(kNMsgs) * kMsgSz;
|
||
|
||
// One send/recv buffer pair per connection.
|
||
std::vector<std::vector<char>> sendBufs(kNConns, std::vector<char>(kBufSz));
|
||
std::vector<std::vector<char>> recvBufs(kNConns, std::vector<char>(kBufSz));
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (size_t i = 0; i < kBufSz; i++) {
|
||
sendBufs[c][i] = static_cast<char>(((i + c) * 3 + 7) & 0xFF);
|
||
}
|
||
memset(recvBufs[c].data(), 0, kBufSz);
|
||
}
|
||
|
||
std::vector<void*> mhandles(kNConns, nullptr);
|
||
for (int c = 0; c < kNConns; c++) {
|
||
void* comm = (rank == 0) ? recvComms[c] : sendComms[c];
|
||
char* regBuf = (rank == 0) ? recvBufs[c].data() : sendBufs[c].data();
|
||
ASSERT_EQ(RegisterMemory(comm, regBuf, kBufSz, NCCL_PTR_HOST, &mhandles[c]), ncclSuccess);
|
||
}
|
||
|
||
// Warmup on conn[0]: learn real nqps. All conns share same NCCL_PARAM → same nqps.
|
||
const int actualNqps = GetActualNqps(sendComms[0], recvComms[0],
|
||
(rank == 0) ? recvBufs[0].data() : sendBufs[0].data(),
|
||
kMsgSz, 3999, mhandles[0]);
|
||
ASSERT_GT(actualNqps, 0);
|
||
|
||
// Arm tokens asymmetrically on conn[0] only (timer-fired check uses conn[0]).
|
||
if (rank == 1) {
|
||
std::vector<int> tokens(actualNqps);
|
||
if (actualNqps == 1) {
|
||
tokens[0] = 100;
|
||
} else {
|
||
tokens[0] = 70;
|
||
tokens[1] = 30;
|
||
for (int i = 2; i < actualNqps; i++) tokens[i] = 50;
|
||
}
|
||
ASSERT_EQ(ncclIbCastSetTokens(sendComms[0], tokens.data(), actualNqps), ncclSuccess);
|
||
}
|
||
|
||
constexpr int kBatch = 16;
|
||
constexpr int kTagBase = 10000; // tag = kTagBase + c * kNMsgs + i
|
||
|
||
// ── Phase 1: small-message WRR stress ────────────────────────────────────
|
||
if (rank == 0) {
|
||
for (int base = 0; base < kNMsgs; base += kBatch) {
|
||
const int end = std::min(base + kBatch, kNMsgs);
|
||
const int bsz = end - base;
|
||
std::vector<std::vector<void*>> reqs(kNConns, std::vector<void*>(bsz, nullptr));
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (int i = base; i < end; i++) {
|
||
void* bufs[1] = {recvBufs[c].data() + i * kMsgSz};
|
||
size_t sizes[1] = {kMsgSz};
|
||
int tags[1] = {kTagBase + c * kNMsgs + i};
|
||
void* handles[1] = {mhandles[c]};
|
||
ASSERT_EQ(PostRecv(recvComms[c], 1, bufs, sizes, tags, handles,
|
||
&reqs[c][i - base]), ncclSuccess);
|
||
ASSERT_NE(reqs[c][i - base], nullptr);
|
||
}
|
||
}
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (int i = 0; i < bsz; i++) {
|
||
int sz = 0;
|
||
ASSERT_EQ(WaitForCompletion(reqs[c][i], &sz, 10000), ncclSuccess);
|
||
}
|
||
}
|
||
}
|
||
for (int c = 0; c < kNConns; c++) {
|
||
EXPECT_EQ(memcmp(recvBufs[c].data(), sendBufs[c].data(), kBufSz), 0)
|
||
<< "phase1 data corruption on conn " << c;
|
||
}
|
||
} else {
|
||
for (int base = 0; base < kNMsgs; base += kBatch) {
|
||
const int end = std::min(base + kBatch, kNMsgs);
|
||
const int bsz = end - base;
|
||
std::vector<std::vector<void*>> reqs(kNConns, std::vector<void*>(bsz, nullptr));
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (int i = base; i < end; i++)
|
||
PostSendWithRetry(sendComms[c], sendBufs[c].data() + i * kMsgSz, kMsgSz,
|
||
kTagBase + c * kNMsgs + i, mhandles[c], &reqs[c][i - base]);
|
||
}
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (int i = 0; i < bsz; i++) {
|
||
int sz = 0;
|
||
ASSERT_EQ(WaitForCompletion(reqs[c][i], &sz, 10000), ncclSuccess);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Collect sched state from conn[0] after phase 1.
|
||
struct ncclIbCastSchedState stAfterPhase1 = {};
|
||
if (rank == 0) {
|
||
MPI_Recv(&stAfterPhase1, sizeof(stAfterPhase1), MPI_BYTE, 1, 9870, MPI_COMM_WORLD,
|
||
MPI_STATUS_IGNORE);
|
||
int nqps = 0;
|
||
MPI_Recv(&nqps, 1, MPI_INT, 1, 9872, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
||
|
||
EXPECT_TRUE(stAfterPhase1.schedInit) << "conn[0]: schedInit must be true after phase 1";
|
||
{
|
||
int s = 0;
|
||
for (int q = 0; q < stAfterPhase1.nqps; q++) s += stAfterPhase1.activeQpTokens[q];
|
||
EXPECT_EQ(s, stAfterPhase1.activeTotTokens) << "conn[0]: token sum mismatch";
|
||
}
|
||
EXPECT_GE(stAfterPhase1.activeTotTokens, 0);
|
||
EXPECT_LE(stAfterPhase1.activeTotTokens, stAfterPhase1.initTotTokens);
|
||
if (nqps > 1) {
|
||
EXPECT_NE(stAfterPhase1.initQpTokens[0], 70) << "conn[0]: RTT timer did not fire";
|
||
EXPECT_NE(stAfterPhase1.initQpTokens[1], 30) << "conn[0]: RTT timer did not fire";
|
||
}
|
||
} else {
|
||
struct ncclIbCastSchedState st0 = {};
|
||
ASSERT_EQ(ncclIbCastGetSchedState(sendComms[0], &st0), ncclSuccess);
|
||
MPI_Send(&st0, sizeof(st0), MPI_BYTE, 0, 9870, MPI_COMM_WORLD);
|
||
MPI_Send(&actualNqps, 1, MPI_INT, 0, 9872, MPI_COMM_WORLD);
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
|
||
// ── Phase 2: ramping-size stress ─────────────────────────────────────────
|
||
// Sizes span WRR and split paths: /4, /2, *1, *4, *16 relative to splitDataMin.
|
||
const uint32_t sdm = GetSplitDataMin();
|
||
const std::vector<size_t> kRampSizes = {
|
||
std::max<size_t>(64, sdm / 4),
|
||
std::max<size_t>(64, sdm / 2),
|
||
static_cast<size_t>(sdm),
|
||
static_cast<size_t>(sdm) * 4,
|
||
static_cast<size_t>(sdm) * 16,
|
||
};
|
||
constexpr int kLargeRounds = 20;
|
||
const size_t kRampBufSz = kRampSizes.back(); // one MR per conn, sized for largest
|
||
|
||
std::vector<std::vector<char>> rampSend(kNConns, std::vector<char>(kRampBufSz));
|
||
std::vector<std::vector<char>> rampRecv(kNConns, std::vector<char>(kRampBufSz));
|
||
for (int c = 0; c < kNConns; c++) {
|
||
for (size_t i = 0; i < kRampBufSz; i++)
|
||
rampSend[c][i] = static_cast<char>(((i + c) * 7 + 3) & 0xFF);
|
||
memset(rampRecv[c].data(), 0, kRampBufSz);
|
||
}
|
||
|
||
std::vector<void*> rampHandles(kNConns, nullptr);
|
||
for (int c = 0; c < kNConns; c++) {
|
||
void* comm = (rank == 0) ? recvComms[c] : sendComms[c];
|
||
char* regBuf = (rank == 0) ? rampRecv[c].data() : rampSend[c].data();
|
||
ASSERT_EQ(RegisterMemory(comm, regBuf, kRampBufSz, NCCL_PTR_HOST, &rampHandles[c]),
|
||
ncclSuccess);
|
||
}
|
||
|
||
// Tag space: 20000 + sizeIdx*kNConns*kLargeRounds + c*kLargeRounds + round
|
||
constexpr int kRampTagBase = 20000;
|
||
for (int si = 0; si < static_cast<int>(kRampSizes.size()); si++) {
|
||
const size_t sz = kRampSizes[si];
|
||
for (int round = 0; round < kLargeRounds; round++) {
|
||
for (int c = 0; c < kNConns; c++) {
|
||
const int tag = kRampTagBase + si * kNConns * kLargeRounds
|
||
+ c * kLargeRounds + round;
|
||
CastDoSendRecv(rank, sendComms[c], recvComms[c],
|
||
(rank == 0) ? rampRecv[c].data() : rampSend[c].data(),
|
||
sz, tag, rampHandles[c]);
|
||
}
|
||
}
|
||
}
|
||
|
||
if (rank == 0) {
|
||
for (int c = 0; c < kNConns; c++) {
|
||
EXPECT_EQ(memcmp(rampRecv[c].data(), rampSend[c].data(), kRampBufSz), 0)
|
||
<< "phase2 data corruption on conn " << c;
|
||
}
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
|
||
// ── Teardown ─────────────────────────────────────────────────────────────
|
||
for (int c = 0; c < kNConns; c++) {
|
||
void* comm = (rank == 0) ? recvComms[c] : sendComms[c];
|
||
ASSERT_EQ(DeregisterMemory(comm, rampHandles[c]), ncclSuccess);
|
||
ASSERT_EQ(DeregisterMemory(comm, mhandles[c]), ncclSuccess);
|
||
if (rank == 0) {
|
||
ASSERT_EQ(CloseRecvComm(recvComms[c]), ncclSuccess);
|
||
ASSERT_EQ(CloseListenComm(listenComms[c]), ncclSuccess);
|
||
} else {
|
||
ASSERT_EQ(CloseSendComm(sendComms[c]), ncclSuccess);
|
||
}
|
||
}
|
||
|
||
MPI_Barrier(MPI_COMM_WORLD);
|
||
}
|
||
|
||
#endif // MPI_TESTS_ENABLED
|