Define BaseShmExporter for multi-process shared-memory telemetry exporters.

tpu_sync/telemetry/BUILD

@@ -161,3 +161,32 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_library(
+    name = "base_shm_exporter",
+    srcs = ["base_shm_exporter.cc"],
+    hdrs = ["base_shm_exporter.h"],
+    deps = [
+        ":metrics_backend",
+        "//tpu_sync/telemetry/shm:shm_collector",
+        "//tpu_sync/telemetry/shm:shm_writer",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "base_shm_exporter_test",
+    srcs = ["base_shm_exporter_test.cc"],
+    deps = [
+        ":base_shm_exporter",
+        ":metrics_backend",
+        "//tpu_sync/telemetry/shm:shm_collector",
+        "//tpu_sync/telemetry/shm:shm_layout",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tpu_sync/telemetry/base_shm_exporter.cc

@@ -0,0 +1,118 @@
+// 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_sync/telemetry/base_shm_exporter.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/string_view.h"
+#include "tpu_sync/telemetry/metrics_backend.h"
+#include "tpu_sync/telemetry/shm/shm_collector.h"
+#include "tpu_sync/telemetry/shm/shm_writer.h"
+
+namespace tpu_raiden::telemetry {
+namespace {
+
+// Validates that `options.local_rank` is a valid non-negative integer
+// (canonicalizing any leading zeros) and that `options.shm_dir` is an absolute,
+// non-empty directory path with trailing slashes stripped, rejecting the root
+// path "/".
+ExporterOptions NormalizeAndValidateOptions(ExporterOptions options) {
+  CHECK(options.local_rank.has_value() && !options.local_rank->empty())
+      << "options.local_rank must be specified and non-empty for "
+         "BaseShmExporter";
+
+  int rank_val = -1;
+  CHECK(absl::SimpleAtoi(*options.local_rank, &rank_val) && rank_val >= 0)
+      << "options.local_rank must be a valid non-negative integer for "
+         "BaseShmExporter, got: '"
+      << *options.local_rank << "'";
+  options.local_rank = std::to_string(rank_val);
+
+  while (options.shm_dir.has_value() && options.shm_dir->ends_with("/")) {
+    options.shm_dir->pop_back();
+  }
+  CHECK(options.shm_dir.has_value() && !options.shm_dir->empty())
+      << "options.shm_dir must be specified and non-empty and a valid path";
+
+  return options;
+}
+
+}  // namespace
+
+BaseShmExporter::BaseShmExporter(ExporterOptions options)
+    : options_(NormalizeAndValidateOptions(std::move(options))),
+      shm_writer_(ShmWriterOptions{
+          .shm_dir = *options_.shm_dir,
+          .local_rank = *options_.local_rank,
+      }),
+      shm_collector_(ShmCollectorOptions{
+          .shm_dir = *options_.shm_dir,
+      }) {}
+
+BaseShmExporter::~BaseShmExporter() = default;
+
+void BaseShmExporter::IncrementCounter(absl::string_view name, LabelSpan labels,
+                                       uint64_t val) const {
+  shm_writer_.IncrementCounter(name, labels, val);
+}
+
+void BaseShmExporter::SetGauge(absl::string_view name, LabelSpan labels,
+                               double val) const {
+  if (!std::isfinite(val)) return;
+
+  const MetricLabel rank_label = local_rank_label();
+  if (labels.empty()) {
+    shm_writer_.SetGauge(name, LabelSpan(&rank_label, 1), val);
+    return;
+  }
+  if (absl::c_any_of(labels, [](const MetricLabel& label) {
+        return label.key == metric_labels::kLocalRank;
+      })) {
+    shm_writer_.SetGauge(name, labels, val);
+    return;
+  }
+  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity + 1>
+      augmented_labels(labels.begin(), labels.end());
+  if (absl::c_is_sorted(labels)) {
+    auto it = std::lower_bound(augmented_labels.begin(), augmented_labels.end(),
+                               rank_label);
+    augmented_labels.insert(it, rank_label);
+  } else {
+    augmented_labels.push_back(rank_label);
+    absl::c_sort(augmented_labels);
+  }
+  shm_writer_.SetGauge(name, augmented_labels, val);
+}
+
+void BaseShmExporter::ObserveHistogram(absl::string_view name, LabelSpan labels,
+                                       double val) const {
+  shm_writer_.ObserveHistogram(name, labels, val);
+}
+
+std::string BaseShmExporter::GetTextSnapshot() const { return std::string(); }
+
+AggregatedMetrics BaseShmExporter::CollectMetrics() const {
+  return shm_collector_.CollectMetrics();
+}
+
+}  // namespace tpu_raiden::telemetry

tpu_sync/telemetry/base_shm_exporter.h

@@ -0,0 +1,89 @@
+// 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_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_
+#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_
+
+#include <cstdint>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "tpu_sync/telemetry/metrics_backend.h"
+#include "tpu_sync/telemetry/shm/shm_collector.h"
+#include "tpu_sync/telemetry/shm/shm_writer.h"
+
+namespace tpu_raiden::telemetry {
+
+// Base class for multi-process shared-memory telemetry exporters.
+// Manages an underlying ShmWriter for low-overhead metric publishing and an
+// ShmCollector for multi-worker aggregation.
+//
+// Preconditions:
+// Requires valid local_rank (via options.local_rank) and shm_dir (via
+// options.shm_dir). Fails fast with CHECK if either is missing or invalid.
+// Trailing slashes in shm_dir are stripped; the root path "/" is prohibited.
+// Note: custom_buckets is ignored by shared-memory backends in favor of fixed
+// compile-time bucket layouts (kDefaultHistogramBuckets).
+//
+// Thread-safety & Destruction contract:
+// Callers must ensure all concurrent metric recording and CollectMetrics
+// calls have finished prior to exporter destruction.
+class BaseShmExporter : public MetricsBackend {
+ public:
+  explicit BaseShmExporter(ExporterOptions options);
+  ~BaseShmExporter() override;
+
+  BaseShmExporter(const BaseShmExporter&) = delete;
+  BaseShmExporter& operator=(const BaseShmExporter&) = delete;
+  BaseShmExporter(BaseShmExporter&&) = delete;
+  BaseShmExporter& operator=(BaseShmExporter&&) = delete;
+
+  // Metric recording methods. Safe for concurrent execution across worker
+  // threads. Writes directly to the memory-mapped shared segment.
+  void IncrementCounter(absl::string_view name, LabelSpan labels,
+                        uint64_t val) const override;
+  // Records a gauge value. Automatically attaches the "local_rank" label if not
+  // already present in caller labels.
+  void SetGauge(absl::string_view name, LabelSpan labels,
+                double val) const override;
+  // Records a histogram observation. Safe for concurrent execution across
+  // worker threads. Writes directly to the memory-mapped shared segment.
+  void ObserveHistogram(absl::string_view name, LabelSpan labels,
+                        double val) const override;
+
+  // Returns an empty string. BaseShmExporter is un-opinionated and does not
+  // prescribe an exposition format (Prometheus vs OTel); this method is
+  // intended to be overridden by format-specific subclasses.
+  std::string GetTextSnapshot() const override;
+
+  // Scans the shared-memory directory and returns aggregated metrics across
+  // all local worker processes.
+  AggregatedMetrics CollectMetrics() const;
+
+  // Returns the validated exporter configuration options.
+  const ExporterOptions& options() const { return options_; }
+
+ private:
+  MetricLabel local_rank_label() const {
+    return MetricLabel{metric_labels::kLocalRank, *options_.local_rank};
+  }
+
+  ExporterOptions options_;
+  ShmWriter shm_writer_;
+  ShmCollector shm_collector_;
+};
+
+}  // namespace tpu_raiden::telemetry
+
+#endif  // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_