diff --git a/tpu_raiden/telemetry/BUILD b/tpu_raiden/telemetry/BUILD new file mode 100644 index 00000000..6e83f3fb --- /dev/null +++ b/tpu_raiden/telemetry/BUILD @@ -0,0 +1,100 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 Google LLC + +load("@nanobind_bazel//:build_defs.bzl", "nanobind_extension") +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") +load("@rules_python//python:defs.bzl", "py_test") + +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) + +cc_library( + name = "metrics_api", + srcs = ["metrics_api.cc"], + hdrs = ["metrics_api.h"], + deps = [ + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "metrics_api_test", + srcs = ["metrics_api_test.cc"], + deps = [ + ":metrics_api", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "metrics_3p_prometheus_exporter", + srcs = ["prometheus_exporter.cc"], + hdrs = ["prometheus_exporter.h"], + deps = [ + ":metrics_api", + "//third_party/prometheus_cpp_client:prometheus_client_core", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "prometheus_exporter_test", + srcs = ["prometheus_exporter_test.cc"], + deps = [ + ":metrics_3p_prometheus_exporter", + ":metrics_api", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +nanobind_extension( + name = "_raiden_telemetry", + srcs = ["raiden_telemetry_module.cc"], + copts = [ + "-fexceptions", + "-frtti", + "-fvisibility=hidden", + ], + features = [ + "-use_header_modules", + ], + linkopts = ["-Wl,--exclude-libs,ALL"], + visibility = ["//visibility:public"], + deps = [ + ":metrics_3p_prometheus_exporter", + ":metrics_api", + "@com_google_absl//absl/base", + "@nanobind", + ], +) + +py_test( + name = "raiden_telemetry_test", + srcs = ["raiden_telemetry_test.py"], + deps = [ + ":_raiden_telemetry", + "@com_google_absl_py//absl/testing:absltest", + ], +) diff --git a/tpu_raiden/telemetry/metrics_api.cc b/tpu_raiden/telemetry/metrics_api.cc new file mode 100644 index 00000000..e0b1816b --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api.cc @@ -0,0 +1,100 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/telemetry/metrics_api.h" + +#include +#include +#include +#include +#include +#include + +#include "absl/base/no_destructor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" + +namespace tpu_raiden::telemetry { + +RaidenMetricStore& RaidenMetricStore::GetGlobalMetricStore() { + static absl::NoDestructor global_store; + return *global_store; +} + +void RaidenMetricStore::AddBackend(std::unique_ptr backend) { + if (!backend) return; + absl::MutexLock lock(mutex_); + backends_.push_back(std::move(backend)); + has_backends_.store(true, std::memory_order_release); +} + + +bool RaidenMetricStore::HasBackends() const { + return has_backends_.load(std::memory_order_acquire); +} + +void RaidenMetricStore::IncrementCounter(absl::string_view name, + LabelSpan labels, uint64_t val) const { + if (!HasBackends()) return; + // TODO: Explore RCU optimization for lock-free reads. + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->IncrementCounter(name, labels, val); + } +} + +void RaidenMetricStore::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + if (!HasBackends()) return; + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->SetGauge(name, labels, val); + } +} + +void RaidenMetricStore::ObserveHistogram(absl::string_view name, + LabelSpan labels, double val) const { + if (!HasBackends()) return; + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->ObserveHistogram(name, labels, val); + } +} + +std::string RaidenMetricStore::GetTextSnapshot() const { + if (!HasBackends()) return ""; + absl::ReaderMutexLock lock(mutex_); + std::string result; + for (const auto& backend : backends_) { + absl::StrAppend(&result, backend->GetTextSnapshot()); + } + return result; +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/metrics_api.h b/tpu_raiden/telemetry/metrics_api.h new file mode 100644 index 00000000..51ff50e4 --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api.h @@ -0,0 +1,150 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_ + +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" + +namespace tpu_raiden::telemetry { + +enum class MetricType { + kCounter, + kGauge, + kHistogram, +}; + +// Structure defining centralized metadata for a Raiden metric across all +// exporter backends (Prometheus, Streamz, etc.). +struct MetricMetadata { + absl::string_view name; + absl::string_view description; + absl::string_view prometheus_name; + absl::string_view streamz_name; + MetricType type; +}; + +namespace metric_names { + +inline constexpr absl::string_view kSentBytesTotal = "sent_bytes_total"; + +} // namespace metric_names + +namespace metric_metadata { + +inline constexpr MetricMetadata kSentBytesTotal{ + .name = metric_names::kSentBytesTotal, + .description = "Total count of bytes sent over TPU Raiden interfaces.", + .prometheus_name = "tpu_raiden_sent_bytes_total", + .streamz_name = "/tpu_raiden/sent_bytes_total", + .type = MetricType::kCounter}; + +inline constexpr MetricMetadata kAllMetrics[] = { + kSentBytesTotal, +}; + +} // namespace metric_metadata + +// Structure defining a metric key-value label pair. +struct MetricLabel { + absl::string_view key; + absl::string_view value; +}; + +// Allocation-free label view span type definition +using LabelSpan = absl::Span; + +// Abstract Dual-Backend Interface +class MetricsBackend { + public: + MetricsBackend() = default; + MetricsBackend(const MetricsBackend&) = delete; + MetricsBackend& operator=(const MetricsBackend&) = delete; + MetricsBackend(MetricsBackend&&) = delete; + MetricsBackend& operator=(MetricsBackend&&) = delete; + + virtual ~MetricsBackend() = default; + + virtual void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const = 0; + + virtual void SetGauge(absl::string_view name, LabelSpan labels, + double val) const = 0; + + virtual void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const = 0; + + virtual std::string GetTextSnapshot() const = 0; +}; + +// Central Telemetry Facade for managing metrics across registered backends. +// This class is thread-safe for all concurrent operations. +class RaidenMetricStore { + public: + static RaidenMetricStore& GetGlobalMetricStore(); + + RaidenMetricStore() = default; + ~RaidenMetricStore() = default; + + RaidenMetricStore(const RaidenMetricStore&) = delete; + RaidenMetricStore& operator=(const RaidenMetricStore&) = delete; + RaidenMetricStore(RaidenMetricStore&&) = delete; + RaidenMetricStore& operator=(RaidenMetricStore&&) = delete; + + void AddBackend(std::unique_ptr backend); + bool HasBackends() const; + + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val = 1) const; + + void SetGauge(absl::string_view name, LabelSpan labels, double val) const; + + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const; + + std::string GetTextSnapshot() const; + + private: + mutable absl::Mutex mutex_; + std::vector> backends_ + ABSL_GUARDED_BY(mutex_); + std::atomic has_backends_{false}; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_ diff --git a/tpu_raiden/telemetry/metrics_api_test.cc b/tpu_raiden/telemetry/metrics_api_test.cc new file mode 100644 index 00000000..86cf5751 --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api_test.cc @@ -0,0 +1,167 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/telemetry/metrics_api.h" + +#include +#include +#include +#include // NOLINT(build/c++11) +#include +#include +#include + +#include +#include +#include "absl/strings/string_view.h" + +namespace tpu_raiden::telemetry { +namespace { + +static_assert(!std::is_copy_constructible_v, + "MetricsBackend must not be copy constructible"); +static_assert(!std::is_copy_assignable_v, + "MetricsBackend must not be copy assignable"); +static_assert(!std::is_move_constructible_v, + "MetricsBackend must not be move constructible"); +static_assert(!std::is_move_assignable_v, + "MetricsBackend must not be move assignable"); + +using testing::_; +using testing::Eq; +using testing::Return; + +class MockMetricsBackend : public MetricsBackend { + public: + MOCK_METHOD(void, IncrementCounter, + (absl::string_view name, LabelSpan labels, uint64_t val), + (override, const)); + MOCK_METHOD(void, SetGauge, + (absl::string_view name, LabelSpan labels, double val), + (override, const)); + MOCK_METHOD(void, ObserveHistogram, + (absl::string_view name, LabelSpan labels, double val), + (override, const)); + MOCK_METHOD(std::string, GetTextSnapshot, (), (override, const)); +}; + +class MetricsApiTest : public testing::Test { + protected: + RaidenMetricStore store_; +}; + +TEST_F(MetricsApiTest, GlobalMetricStoreSingleton) { + RaidenMetricStore& global1 = RaidenMetricStore::GetGlobalMetricStore(); + RaidenMetricStore& global2 = RaidenMetricStore::GetGlobalMetricStore(); + EXPECT_EQ(&global1, &global2); +} + +TEST_F(MetricsApiTest, MetricMetadataConstants) { + EXPECT_EQ(metric_metadata::kSentBytesTotal.description, + "Total count of bytes sent over TPU Raiden interfaces."); +} + +TEST_F(MetricsApiTest, FastPathExitWhenNoBackends) { + EXPECT_FALSE(store_.HasBackends()); + + store_.IncrementCounter(metric_names::kSentBytesTotal, {}, 1024); +} + +TEST_F(MetricsApiTest, DispatchesToRegisteredBackend) { + auto mock_backend = std::make_unique(); + MockMetricsBackend* raw_mock = mock_backend.get(); + + EXPECT_CALL(*raw_mock, + IncrementCounter(Eq(metric_names::kSentBytesTotal), _, 2048)) + .Times(1); + EXPECT_CALL(*raw_mock, GetTextSnapshot()).WillOnce(Return("# HELP mock\n")); + + store_.AddBackend(std::move(mock_backend)); + EXPECT_TRUE(store_.HasBackends()); + + store_.IncrementCounter(metric_names::kSentBytesTotal, {}, 2048); + EXPECT_EQ(store_.GetTextSnapshot(), "# HELP mock\n"); +} + + +TEST_F(MetricsApiTest, ConcurrentTelemetryEmissions) { + auto backend = std::make_unique(); + MockMetricsBackend* raw_backend = backend.get(); + + constexpr int kNumThreads = 8; + constexpr int kIterations = 100; + constexpr int kTotalCalls = kNumThreads * kIterations; + + EXPECT_CALL(*raw_backend, IncrementCounter(Eq("counter"), _, 1)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, SetGauge(Eq("gauge"), _, 42)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, ObserveHistogram(Eq("histogram"), _, 3.14)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, GetTextSnapshot()) + .Times(kTotalCalls) + .WillRepeatedly(Return("snapshot\n")); + + store_.AddBackend(std::move(backend)); + + std::vector threads; + threads.reserve(kNumThreads); + + for (int i = 0; i < kNumThreads; ++i) { + threads.emplace_back([this] { + for (int j = 0; j < kIterations; ++j) { + store_.IncrementCounter("counter", {}, 1); + store_.SetGauge("gauge", {}, 42); + store_.ObserveHistogram("histogram", {}, 3.14); + (void)store_.GetTextSnapshot(); + } + }); + } + + for (auto& t : threads) { + t.join(); + } +} + +TEST_F(MetricsApiTest, ConstMetricsBackendReference) { + MockMetricsBackend backend; + const MetricsBackend& const_ref = backend; + + EXPECT_CALL(backend, IncrementCounter(Eq("counter"), _, 5)).Times(1); + EXPECT_CALL(backend, SetGauge(Eq("gauge"), _, 10)).Times(1); + EXPECT_CALL(backend, ObserveHistogram(Eq("histogram"), _, 1.23)).Times(1); + EXPECT_CALL(backend, GetTextSnapshot()).WillOnce(Return("snapshot\n")); + + const_ref.IncrementCounter("counter", {}, 5); + const_ref.SetGauge("gauge", {}, 10); + const_ref.ObserveHistogram("histogram", {}, 1.23); + EXPECT_EQ(const_ref.GetTextSnapshot(), "snapshot\n"); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/prometheus_exporter.cc b/tpu_raiden/telemetry/prometheus_exporter.cc new file mode 100644 index 00000000..4fd844f6 --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter.cc @@ -0,0 +1,175 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/telemetry/prometheus_exporter.h" + +#include +#include // NOLINT: Required by prometheus-cpp client API. +#include +#include +#include + +#include "prometheus/counter.h" +#include "prometheus/family.h" +#include "prometheus/gauge.h" +#include "prometheus/histogram.h" +#include "prometheus/registry.h" +#include "prometheus/text_serializer.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { + +namespace { + +std::map ConvertLabels(LabelSpan labels) { + if (labels.empty()) { + return {}; + } + std::map result; + for (const auto& [key, value] : labels) { + result.emplace(key, value); + } + return result; +} + +} // namespace + +void PrometheusExporter::RegisterKnownFamilies() { + for (const auto& meta : metric_metadata::kAllMetrics) { + switch (meta.type) { + case MetricType::kCounter: { + auto* family = &prometheus::BuildCounter() + .Name(std::string(meta.prometheus_name)) + .Help(std::string(meta.description)) + .Register(*registry_); + counter_families_.emplace(std::string(meta.name), family); + counter_families_.emplace(std::string(meta.prometheus_name), family); + counter_families_.emplace(std::string(meta.streamz_name), family); + break; + } + case MetricType::kGauge: { + auto* family = &prometheus::BuildGauge() + .Name(std::string(meta.prometheus_name)) + .Help(std::string(meta.description)) + .Register(*registry_); + gauge_families_.emplace(std::string(meta.name), family); + gauge_families_.emplace(std::string(meta.prometheus_name), family); + gauge_families_.emplace(std::string(meta.streamz_name), family); + break; + } + case MetricType::kHistogram: { + auto* family = &prometheus::BuildHistogram() + .Name(std::string(meta.prometheus_name)) + .Help(std::string(meta.description)) + .Register(*registry_); + histogram_families_.emplace(std::string(meta.name), family); + histogram_families_.emplace(std::string(meta.prometheus_name), family); + histogram_families_.emplace(std::string(meta.streamz_name), family); + break; + } + } + } +} + +PrometheusExporter::PrometheusExporter( + prometheus::Histogram::BucketBoundaries custom_buckets) + : registry_(std::make_shared()), + default_buckets_(std::move(custom_buckets)) { + RegisterKnownFamilies(); +} + +prometheus::Family* PrometheusExporter::GetCounterFamily( + absl::string_view name) const { + auto it = counter_families_.find(name); + if (it == counter_families_.end()) { + return nullptr; + } + return it->second; +} + +prometheus::Family* PrometheusExporter::GetGaugeFamily( + absl::string_view name) const { + auto it = gauge_families_.find(name); + if (it == gauge_families_.end()) { + return nullptr; + } + return it->second; +} + +prometheus::Family* +PrometheusExporter::GetHistogramFamily(absl::string_view name) const { + auto it = histogram_families_.find(name); + if (it == histogram_families_.end()) { + return nullptr; + } + return it->second; +} + +void PrometheusExporter::IncrementCounter(absl::string_view name, + LabelSpan labels, + uint64_t val) const { + prometheus::Family* family = nullptr; + family = GetCounterFamily(name); + if (family == nullptr) { + return; + } + prometheus::Counter& counter = family->Add(ConvertLabels(labels)); + counter.Increment(static_cast(val)); +} + +void PrometheusExporter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + prometheus::Family* family = nullptr; + family = GetGaugeFamily(name); + if (family == nullptr) { + return; + } + prometheus::Gauge& gauge = family->Add(ConvertLabels(labels)); + gauge.Set(val); +} + +void PrometheusExporter::ObserveHistogram(absl::string_view name, + LabelSpan labels, double val) const { + prometheus::Family* family = nullptr; + family = GetHistogramFamily(name); + if (family == nullptr) { + return; + } + prometheus::Histogram& histogram = + family->Add(ConvertLabels(labels), default_buckets_); + histogram.Observe(val); +} + +std::string PrometheusExporter::GetTextSnapshot() const { + prometheus::TextSerializer serializer; + return serializer.Serialize(registry_->Collect()); +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/prometheus_exporter.h b/tpu_raiden/telemetry/prometheus_exporter.h new file mode 100644 index 00000000..6d58c0dd --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter.h @@ -0,0 +1,108 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ + +#include +#include +#include + +#include "prometheus/counter.h" +#include "prometheus/family.h" +#include "prometheus/gauge.h" +#include "prometheus/histogram.h" +#include "prometheus/registry.h" +#include "absl/base/no_destructor.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { + +inline const prometheus::Histogram::BucketBoundaries& +DefaultHistogramBuckets() { + static const absl::NoDestructor + kBuckets(prometheus::Histogram::BucketBoundaries{ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0}); + return *kBuckets; +} + +class PrometheusExporter : public MetricsBackend { + public: + explicit PrometheusExporter(prometheus::Histogram::BucketBoundaries + custom_buckets = DefaultHistogramBuckets()); + + ~PrometheusExporter() override = default; + + PrometheusExporter(const PrometheusExporter&) = delete; + PrometheusExporter& operator=(const PrometheusExporter&) = delete; + PrometheusExporter(PrometheusExporter&&) = delete; + PrometheusExporter& operator=(PrometheusExporter&&) = delete; + + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const override; + + void SetGauge(absl::string_view name, LabelSpan labels, + double val) const override; + + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const override; + + std::string GetTextSnapshot() const override; + + std::shared_ptr GetRegistry() const { + return registry_; + } + + private: + void RegisterKnownFamilies(); + + prometheus::Family* GetCounterFamily( + absl::string_view name) const; + prometheus::Family* GetGaugeFamily(absl::string_view name) + const; + prometheus::Family* GetHistogramFamily( + absl::string_view name) const; + std::shared_ptr registry_; + prometheus::Histogram::BucketBoundaries default_buckets_; + + mutable absl::flat_hash_map*> + counter_families_; + mutable absl::flat_hash_map*> + gauge_families_; + mutable absl::flat_hash_map*> + histogram_families_; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ diff --git a/tpu_raiden/telemetry/prometheus_exporter_test.cc b/tpu_raiden/telemetry/prometheus_exporter_test.cc new file mode 100644 index 00000000..5dcc76a7 --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter_test.cc @@ -0,0 +1,123 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_raiden/telemetry/prometheus_exporter.h" + +#include +#include +#include // NOLINT(build/c++11) +#include +#include + +#include +#include +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { + +using ::testing::HasSubstr; + +TEST(PrometheusExporterTest, RecordAndExportFormat) { + auto exporter = std::make_unique(); + RaidenMetricStore store; + store.AddBackend(std::move(exporter)); + + MetricLabel label1{"interface", "ICI"}; + + store.IncrementCounter(metric_names::kSentBytesTotal, {label1}, 1024); + + std::string output = store.GetTextSnapshot(); + + EXPECT_TRUE( + absl::StrContains(output, "# TYPE tpu_raiden_sent_bytes_total counter")); + EXPECT_TRUE(absl::StrContains( + output, "tpu_raiden_sent_bytes_total{interface=\"ICI\"} 1024")); +} + +TEST(PrometheusExporterTest, MetricMetadataConstantsMapped) { + EXPECT_EQ(metric_metadata::kSentBytesTotal.prometheus_name, + "tpu_raiden_sent_bytes_total"); + EXPECT_EQ(metric_metadata::kSentBytesTotal.streamz_name, + "/tpu_raiden/sent_bytes_total"); + EXPECT_THAT( + metric_metadata::kSentBytesTotal.description, + HasSubstr("Total count of bytes sent over TPU Raiden interfaces.")); +} + +TEST(PrometheusExporterTest, UnmappedMetricIgnored) { + PrometheusExporter exporter; + exporter.IncrementCounter("custom_unmapped_counter", {}, 42); + exporter.SetGauge("custom_unmapped_gauge", {}, 99); + exporter.ObserveHistogram("custom_unmapped_histogram", {}, 1.23); + + std::string output = exporter.GetTextSnapshot(); + + EXPECT_FALSE(absl::StrContains(output, "custom_unmapped_counter")); + EXPECT_FALSE(absl::StrContains(output, "custom_unmapped_gauge")); + EXPECT_FALSE(absl::StrContains(output, "custom_unmapped_histogram")); +} + +TEST(PrometheusExporterTest, ConstReferenceAccess) { + PrometheusExporter exporter; + const PrometheusExporter& const_exporter = exporter; + const MetricsBackend& const_backend = exporter; + + const_exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 500); + + std::string snapshot = const_backend.GetTextSnapshot(); + EXPECT_TRUE(absl::StrContains(snapshot, "tpu_raiden_sent_bytes_total 500")); +} + +TEST(PrometheusExporterTest, ConcurrentMetricUpdates) { + PrometheusExporter exporter; + constexpr int kNumThreads = 8; + constexpr int kIterations = 1000; + std::vector threads; + threads.reserve(kNumThreads); + + for (int i = 0; i < kNumThreads; ++i) { + threads.emplace_back([&exporter]() { + for (int j = 0; j < kIterations; ++j) { + exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 1); + } + }); + } + + for (auto& t : threads) { + t.join(); + } + + std::string snapshot = exporter.GetTextSnapshot(); + EXPECT_TRUE(absl::StrContains( + snapshot, + absl::StrCat("tpu_raiden_sent_bytes_total ", kNumThreads * kIterations))); +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/raiden_telemetry_module.cc b/tpu_raiden/telemetry/raiden_telemetry_module.cc new file mode 100644 index 00000000..1cb1c87a --- /dev/null +++ b/tpu_raiden/telemetry/raiden_telemetry_module.cc @@ -0,0 +1,65 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include "absl/base/call_once.h" +#include +#include // IWYU pragma: keep +#include "tpu_raiden/telemetry/metrics_api.h" +#include "tpu_raiden/telemetry/prometheus_exporter.h" + +namespace nb = nanobind; + +NB_MODULE(_raiden_telemetry, m) { + m.doc() = "Python C++ bridge for TPU Raiden telemetry text snapshots."; + + m.def( + "init_prometheus_backend", + []() { + static absl::once_flag prometheus_init_flag; + absl::call_once(prometheus_init_flag, []() { + tpu_raiden::telemetry::RaidenMetricStore::GetGlobalMetricStore() + .AddBackend(std::make_unique< + tpu_raiden::telemetry::PrometheusExporter>()); + }); + }, + "Initializes and registers the C++ 3P Prometheus backend with the " + "global RaidenMetricStore if not already registered."); + + m.def( + "get_raiden_metrics_prometheus_text", + []() -> std::string { + nb::gil_scoped_release release; + return tpu_raiden::telemetry::RaidenMetricStore::GetGlobalMetricStore() + .GetTextSnapshot(); + }, + "Exports the Prometheus text snapshot of TPU Raiden metrics without " + "holding the Python GIL."); +} diff --git a/tpu_raiden/telemetry/raiden_telemetry_test.py b/tpu_raiden/telemetry/raiden_telemetry_test.py new file mode 100644 index 00000000..7b5c67a9 --- /dev/null +++ b/tpu_raiden/telemetry/raiden_telemetry_test.py @@ -0,0 +1,50 @@ +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _raiden_telemetry C++/Python FFI bridge module.""" + +from absl.testing import absltest +from tpu_raiden.telemetry import _raiden_telemetry + + +class RaidenTelemetryTest(absltest.TestCase): + + def test_get_prometheus_text_snapshot(self): + # Call before backend initialization (should return empty string or valid text) + snapshot = _raiden_telemetry.get_raiden_metrics_prometheus_text() + self.assertIsInstance(snapshot, str) + + def test_init_prometheus_backend(self): + # Initialize the C++ Prometheus backend + _raiden_telemetry.init_prometheus_backend() + snapshot = _raiden_telemetry.get_raiden_metrics_prometheus_text() + self.assertIsInstance(snapshot, str) + + +if __name__ == "__main__": + absltest.main()