[tpu_raiden] Add Python/C++ bridge module for Prometheus telemetry

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",
+    ],
+)

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 <atomic>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#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<RaidenMetricStore> global_store;
+  return *global_store;
+}
+
+void RaidenMetricStore::AddBackend(std::unique_ptr<MetricsBackend> 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

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 <atomic>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#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<const MetricLabel>;
+
+// 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<MetricsBackend> 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<std::unique_ptr<MetricsBackend>> backends_
+      ABSL_GUARDED_BY(mutex_);
+  std::atomic<bool> has_backends_{false};
+};
+
+}  // namespace tpu_raiden::telemetry
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_