Add ncclCommDump API (#2068)

* Add ncclCommDump API

* remove trailing whitespace changes

* Add more proxy trace timestamps

* Add facebook_rccl namespace before proxyTrace timestamp call

* Clean up ProxyTrae construction

* Move updateProxyOpCounter to member function

* Move setProxyOpTimestamp to member function

* Move addNewProxyOp to member function

* Make internal methods private

* Make ProxyTrace thread safe

* Fix unit tests

* Fix overwritten ProxyTrace DONE setting in net.cc
Tento commit je obsažen v:
Ahmed Khan
2025-12-11 17:02:35 -05:00
odevzdal GitHub
rodič 1f2f9f33ba
revize 08dd75712f
9 změnil soubory, kde provedl 230 přidání a 138 odebrání
+24 -37
Zobrazit soubor
@@ -9,6 +9,7 @@
#include "proxy_trace/proxy_trace.h"
#include <cstdint>
#include <gtest/gtest.h>
#include <memory>
#include <unistd.h>
namespace RcclUnitTesting {
@@ -21,7 +22,7 @@ public:
int nSteps = 10;
void SetUp() override {
proxyState = new ncclProxyState();
facebook_rccl::proxyTraceInit(proxyState->proxyTrace, 0, commHash);
proxyState->proxyTrace = std::make_unique<facebook_rccl::ProxyTrace>(0);
EXPECT_NE(proxyState->proxyTrace, nullptr);
sub1 = new ncclProxySubArgs();
sub2 = new ncclProxySubArgs();
@@ -36,9 +37,12 @@ public:
delete proxyState;
}
void AddTraceOp(ncclProxySubArgs *sub, facebook_rccl::ProxyOpType opType) {
facebook_rccl::addNewProxyOp(proxyState->proxyTrace, sub->traceKey,
sub->traceInfo, opType, sub->channelId,
sub->nsteps, sub->nbytes, sub->peer);
proxyState->proxyTrace->addNewProxyOp(
sub->traceKey,
sub->traceInfo,
opType,
sub->channelId,
sub->nsteps, sub->nbytes, sub->peer);
}
};
@@ -49,16 +53,10 @@ TEST_F(ProxyTraceTestFixture, nonEmptySingleton) {
TEST_F(ProxyTraceTestFixture, addTraceOp) {
auto &tracer = proxyState->proxyTrace;
EXPECT_EQ(tracer->getOrCreateProxyOpId(sub1->traceKey.commHash,
sub1->traceKey.opCount),
0);
AddTraceOp(sub1, facebook_rccl::ProxyOpType::SEND);
EXPECT_EQ(sub1->traceKey.proxyOpId, 0);
AddTraceOp(sub2, facebook_rccl::ProxyOpType::RECV);
EXPECT_EQ(sub2->traceKey.proxyOpId, 1);
EXPECT_EQ(tracer->getOrCreateProxyOpId(sub1->traceKey.commHash,
sub1->traceKey.opCount),
2);
auto traceRecordPtr = tracer->getProxyTraceOpPtr(sub1->traceKey);
EXPECT_EQ(traceRecordPtr->opType, facebook_rccl::ProxyOpType::SEND);
}
@@ -73,9 +71,10 @@ TEST_F(ProxyTraceTestFixture, getMapSizeMB) {
EXPECT_GT(size2, size1);
// finish sub1
sub1->done = nSteps;
facebook_rccl::updateProxyOpCounter(tracer, sub1->traceKey,
facebook_rccl::ProxyCounterTypes::DONE,
sub1->done);
tracer->updateProxyOpCounter(
sub1->traceKey,
facebook_rccl::ProxyCounterTypes::DONE,
sub1->done);
// sub1 is now serialized and should be moved from activeOps to finishedOps
auto size3 = tracer->getMapSizeMB();
EXPECT_GT(size3, size1);
@@ -84,13 +83,14 @@ TEST_F(ProxyTraceTestFixture, getMapSizeMB) {
TEST_F(ProxyTraceTestFixture, updateTraceOp) {
auto &tracer = proxyState->proxyTrace;
AddTraceOp(sub1, facebook_rccl::ProxyOpType::SEND);
facebook_rccl::updateProxyOpCounter(
tracer, sub1->traceKey,
facebook_rccl::ProxyCounterTypes::KERNEL_COPY_READY, 1);
facebook_rccl::updateProxyOpCounter(
tracer, sub1->traceKey, facebook_rccl::ProxyCounterTypes::POSTED, 3);
facebook_rccl::updateProxyOpCounter(
tracer, sub1->traceKey, facebook_rccl::ProxyCounterTypes::TRANSMITTED, 2);
tracer->updateProxyOpCounter(
sub1->traceKey,
facebook_rccl::ProxyCounterTypes::KERNEL_COPY_READY,
1);
tracer->updateProxyOpCounter(
sub1->traceKey, facebook_rccl::ProxyCounterTypes::POSTED, 3);
tracer->updateProxyOpCounter(
sub1->traceKey, facebook_rccl::ProxyCounterTypes::TRANSMITTED, 2);
auto traceRecordPtr = tracer->getProxyTraceOpPtr(sub1->traceKey);
EXPECT_NE(traceRecordPtr, nullptr);
@@ -110,25 +110,12 @@ TEST_F(ProxyTraceTestFixture, updateTraceOp2) {
AddTraceOp(sub1, facebook_rccl::ProxyOpType::SEND);
int64_t rand = 123456789;
sub1->posted = rand;
facebook_rccl::updateProxyOpCounter(tracer, sub1->traceKey,
facebook_rccl::ProxyCounterTypes::POSTED,
sub1->posted);
tracer->updateProxyOpCounter(sub1->traceKey,
facebook_rccl::ProxyCounterTypes::POSTED,
sub1->posted);
auto traceRecordPtr = tracer->getProxyTraceOpPtr(sub1->traceKey);
EXPECT_EQ(traceRecordPtr->counters[facebook_rccl::ProxyCounterTypes::POSTED],
rand);
}
TEST_F(ProxyTraceTestFixture, memoryReclaim) {
auto &tracer = proxyState->proxyTrace;
tracer->resetAll();
AddTraceOp(sub1, facebook_rccl::ProxyOpType::SEND);
sub1->done = nSteps;
facebook_rccl::updateProxyOpCounter(tracer, sub1->traceKey,
facebook_rccl::ProxyCounterTypes::DONE,
sub1->done);
auto traceRecordPtr = tracer->getProxyTraceOpPtr(sub1->traceKey);
EXPECT_EQ(traceRecordPtr, nullptr);
EXPECT_GT(tracer->getMapSizeMB(), 0);
}
} // namespace RcclUnitTesting
} // namespace RcclUnitTesting