Skip to content

Declare and instrument PCIe H2D and D2H byte counter and per-device label telemetry metrics: #950

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Declare and instrument PCIe H2D and D2H byte counter and per-device label telemetry metrics: #950

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 61 additions & 17 deletions tpu_sync/core/kv_cache_manager_with_transfer_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,20 @@ namespace tpu_raiden {
namespace {

using ::absl_testing::IsOk;
using ::testing::AllOf;
using ::testing::ElementsAre;
using ::testing::Field;
using ::testing::Gt;
using ::testing::IsEmpty;
using ::testing::Ne;
using ::testing::Not;

auto HasResolvedIpLabel(absl::string_view key) {
return AllOf(
Field(&telemetry::MetricLabel::key, key),
Field(&telemetry::MetricLabel::value,
AllOf(Not(IsEmpty()), Ne(telemetry::metric_labels::kUnknownIp))));
}

class TestBase : public kv_cache::KVCacheManagerBase {
public:
Expand Down Expand Up @@ -109,41 +121,73 @@ TEST(KVCacheManagerWithTransferTest, LocalOrchestratedTransfer) {

ASSERT_THAT(buffer->GetReadyFuture().Await(), IsOk());

auto handle_or = raiden::RaidenBufferHandle::Acquire(buffer.get());
ASSERT_THAT(handle_or, IsOk());
std::vector<std::vector<raiden::RaidenBufferHandle>> layer_buffers = {
{handle_or.value()}};
auto engine = std::make_unique<KVCacheManagerWithTransfer>(
layer_buffers,
/*local_port=*/std::nullopt,
/*host_blocks_to_allocate=*/std::nullopt,
/*unsafe_skip_buffer_lock=*/true,
/*parallelism=*/1,
/*host_allocator=*/nullptr,
/*node_id=*/0,
/*local_control_port=*/0,
/*max_blocks=*/2,
/*num_slots=*/2,
/*timeout_s=*/10.0);
const uint64_t expected_slice_bytes = engine->base()->slice_byte_size();

const auto pcie_labels_matcher = ElementsAre(
HasResolvedIpLabel(telemetry::metric_labels::kHostIp),
telemetry::MetricLabel{.key = telemetry::metric_labels::kLocalRank,
.value = "0"});

auto mock_backend = std::make_unique<telemetry::MockMetricsBackend>();
telemetry::MockMetricsBackend* raw_mock = mock_backend.get();
EXPECT_CALL(*raw_mock,
IncrementCounter(telemetry::metric_names::kSentBytesTotal,
testing::_, Gt(0)))
.Times(1);
EXPECT_CALL(*raw_mock,
IncrementCounter(telemetry::metric_names::kReceivedBytesTotal,
testing::_, Gt(0)))
.Times(1);
EXPECT_CALL(*raw_mock,
ObserveHistogram(telemetry::metric_names::kTransferDurationMs,
IsEmpty(), Gt(0.0)))
.Times(1);
EXPECT_CALL(
*raw_mock,
ObserveHistogram(
telemetry::metric_names::kP2pTransferTimeMs,
ElementsAre(
telemetry::MetricLabel{.key = telemetry::metric_labels::kSrcIp,
.value = "127.0.0.1"},
HasResolvedIpLabel(telemetry::metric_labels::kDstIp)),
Gt(0.0)))
.Times(1);
EXPECT_CALL(*raw_mock,
ObserveHistogram(telemetry::metric_names::kD2hTransferTimeMs,
IsEmpty(), Gt(0.0)))
.Times(1);
EXPECT_CALL(*raw_mock,
IncrementCounter(telemetry::metric_names::kD2hBytesTotal,
pcie_labels_matcher, expected_slice_bytes))
.Times(1);
EXPECT_CALL(*raw_mock,
ObserveHistogram(telemetry::metric_names::kH2dTransferTimeMs,
IsEmpty(), Gt(0.0)))
.Times(1);
EXPECT_CALL(*raw_mock,
IncrementCounter(telemetry::metric_names::kH2dBytesTotal,
pcie_labels_matcher, expected_slice_bytes))
.Times(1);
// Register mock backend
telemetry::ScopedMetricsBackendReset scoped_metrics_reset(
std::move(mock_backend));

// Create KVCacheManagerWithTransfer
auto handle_or = raiden::RaidenBufferHandle::Acquire(buffer.get());
std::vector<std::vector<raiden::RaidenBufferHandle>> layer_buffers = {
{handle_or.value()}};
auto engine = std::make_unique<KVCacheManagerWithTransfer>(
layer_buffers,
/*local_port=*/std::nullopt,
/*host_blocks_to_allocate=*/std::nullopt,
/*unsafe_skip_buffer_lock=*/true,
/*parallelism=*/1,
/*host_allocator=*/nullptr,
/*node_id=*/0,
/*local_control_port=*/0,
/*max_blocks=*/2,
/*num_slots=*/2,
/*timeout_s=*/10.0);

// Configure staging slots: 2 slots, max 2 blocks per slot
ASSERT_THAT(engine->base()->ConfigureHostStagingSlots(2, 2), IsOk());

Expand Down
Loading
Loading