diff --git a/tpu_sync/telemetry/BUILD b/tpu_sync/telemetry/BUILD index ef6825aa..d3a57d89 100644 --- a/tpu_sync/telemetry/BUILD +++ b/tpu_sync/telemetry/BUILD @@ -20,9 +20,7 @@ cc_library( name = "metrics_backend", hdrs = ["metrics_backend.h"], deps = [ - "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/strings", - "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", ], ) @@ -46,6 +44,12 @@ cc_library( ], ) +cc_library( + name = "test_util", + testonly = True, + hdrs = ["test_util.h"], +) + cc_library( name = "mock_metrics_backend", testonly = True, @@ -65,6 +69,7 @@ cc_test( ":metrics_api", ":metrics_backend", ":mock_metrics_backend", + ":test_util", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", @@ -76,11 +81,17 @@ cc_library( name = "metrics_3p_prometheus_exporter", srcs = ["prometheus_exporter.cc"], hdrs = ["prometheus_exporter.h"], + # -fexceptions is required because prometheus-cpp Exposer throws C++ + # exceptions (std::runtime_error) on socket binding and initialization failure. + copts = ["-fexceptions"], + # Disables Clang header modules due to incompatibility with prometheus-cpp / CivetWeb headers. + features = ["-use_header_modules"], deps = [ ":metrics_backend", "@com_github_jupp0r_prometheus_cpp//core", - "@com_google_absl//absl/base:no_destructor", + "@com_github_jupp0r_prometheus_cpp//pull", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/strings", ], ) @@ -91,6 +102,7 @@ cc_test( deps = [ ":metrics_3p_prometheus_exporter", ":metrics_backend", + ":test_util", "@com_github_jupp0r_prometheus_cpp//core", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", diff --git a/tpu_sync/telemetry/metrics_api.cc b/tpu_sync/telemetry/metrics_api.cc index 2ea463f9..b5fea913 100644 --- a/tpu_sync/telemetry/metrics_api.cc +++ b/tpu_sync/telemetry/metrics_api.cc @@ -28,6 +28,7 @@ #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/ascii.h" +#include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" @@ -39,6 +40,35 @@ namespace tpu_raiden::telemetry { +namespace { + +int ResolveExporterPort() { + const char* env_port = std::getenv(kPrometheusPortEnvVar); + if (env_port != nullptr && *env_port != '\0') { + int parsed_port = 0; + if (absl::SimpleAtoi(env_port, &parsed_port) && parsed_port >= kMinPort && + parsed_port <= kMaxPort) { + return parsed_port; + } + LOG(WARNING) << "Invalid port specified in " << kPrometheusPortEnvVar + << ": '" << env_port << "'. Expected integer in range [" + << kMinPort << ", " << kMaxPort + << "]. Falling back to default port (" << kDefaultExporterPort + << ")."; + } + return kDefaultExporterPort; +} + +std::string ResolveExporterHost() { + const char* env_host = std::getenv(kPrometheusHostEnvVar); + if (env_host != nullptr && *env_host != '\0') { + return std::string(env_host); + } + return std::string(kDefaultExporterHost); +} + +} // namespace + RaidenMetricStore& RaidenMetricStore::GetGlobalMetricStore() { static absl::NoDestructor global_store; static const bool initialized = [&] { @@ -80,7 +110,11 @@ absl::Status RaidenMetricStore::InitializeFromBackendNames( continue; } if (name == kPrometheus) { - new_backends.push_back(std::make_unique()); + new_backends.push_back( + std::make_unique(ExporterOptions{ + .bind_address = ResolveExporterHost(), + .port = ResolveExporterPort(), + })); } else if (name == kBuffered) { new_backends.push_back(std::make_unique()); } else { diff --git a/tpu_sync/telemetry/metrics_api.h b/tpu_sync/telemetry/metrics_api.h index 1e9f24cd..5d457059 100644 --- a/tpu_sync/telemetry/metrics_api.h +++ b/tpu_sync/telemetry/metrics_api.h @@ -34,6 +34,8 @@ namespace tpu_raiden::telemetry { // Environment Variables START. inline constexpr char kTelemetryBackendsEnvVar[] = "TPU_RAIDEN_TELEMETRY_BACKENDS"; +inline constexpr char kPrometheusPortEnvVar[] = "TPU_RAIDEN_PROMETHEUS_PORT"; +inline constexpr char kPrometheusHostEnvVar[] = "TPU_RAIDEN_PROMETHEUS_HOST"; // Environment Variables END. // Backend names START. diff --git a/tpu_sync/telemetry/metrics_api_test.cc b/tpu_sync/telemetry/metrics_api_test.cc index 1e253d73..d9005c5c 100644 --- a/tpu_sync/telemetry/metrics_api_test.cc +++ b/tpu_sync/telemetry/metrics_api_test.cc @@ -28,8 +28,10 @@ #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/test_util.h" #include "tpu_sync/telemetry/mock_metrics_backend.h" namespace tpu_raiden::telemetry { @@ -659,5 +661,36 @@ TEST_F(MetricsApiTest, InitializeFromEnvironmentNoOpIfAlreadyInitialized) { EXPECT_TRUE(store_.HasBackends()); } +TEST_F(MetricsApiTest, InitializeWithPrometheusPortEnvironmentVariable) { + int port = PickUnusedPort(); + std::string port_str = absl::StrCat(port); + ScopedEnvironmentVariable port_env(kPrometheusPortEnvVar, port_str.c_str()); + EXPECT_THAT(store_.InitializeFromBackendNames({"prometheus"}), IsOk()); + EXPECT_TRUE(store_.HasBackends()); +} + +TEST_F(MetricsApiTest, InitializeWithPrometheusHostEnvironmentVariable) { + int port = PickUnusedPort(); + std::string port_str = absl::StrCat(port); + ScopedEnvironmentVariable port_env(kPrometheusPortEnvVar, port_str.c_str()); + ScopedEnvironmentVariable host_env(kPrometheusHostEnvVar, "127.0.0.1"); + EXPECT_THAT(store_.InitializeFromBackendNames({"prometheus"}), IsOk()); + EXPECT_TRUE(store_.HasBackends()); +} + +TEST_F(MetricsApiTest, InitializeWithInvalidPrometheusPortFallsBackGracefully) { + ScopedEnvironmentVariable invalid_port_env(kPrometheusPortEnvVar, + "invalid_port"); + EXPECT_THAT(store_.InitializeFromBackendNames({"prometheus"}), IsOk()); + EXPECT_TRUE(store_.HasBackends()); +} + +TEST_F(MetricsApiTest, + InitializeWithOutOfRangePrometheusPortFallsBackGracefully) { + ScopedEnvironmentVariable out_of_range_env(kPrometheusPortEnvVar, "99999"); + EXPECT_THAT(store_.InitializeFromBackendNames({"prometheus"}), IsOk()); + EXPECT_TRUE(store_.HasBackends()); +} + } // namespace } // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/metrics_backend.h b/tpu_sync/telemetry/metrics_backend.h index afb0f799..4612e080 100644 --- a/tpu_sync/telemetry/metrics_backend.h +++ b/tpu_sync/telemetry/metrics_backend.h @@ -17,13 +17,10 @@ #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 { @@ -39,6 +36,22 @@ inline constexpr double kDefaultHistogramBuckets[] = { 25.0, 50.0, 100.0, 250.0, 500.0, 750.0, 1000.0, 2500.0, 5000.0, 7500.0, 10000.0, 25000.0, 50000.0}; +inline constexpr int kDefaultExporterPort = 0; +inline constexpr absl::string_view kDefaultExporterHost = "0.0.0.0"; +inline constexpr int kMinPort = 1; +inline constexpr int kMaxPort = 65535; + +// Configuration options for metric exporters. +struct ExporterOptions { + // Bind address for HTTP metric exporter. Defaults to "0.0.0.0". + std::string bind_address{kDefaultExporterHost}; + // TCP port for HTTP metric exporter. If <= 0, HTTP serving is disabled. + int port = kDefaultExporterPort; + // Non-owning view of histogram bucket boundaries. Defaults to + // kDefaultHistogramBuckets and is copied by the exporter during construction. + absl::Span custom_buckets = kDefaultHistogramBuckets; +}; + // Structure defining centralized metadata for a Raiden metric. struct MetricMetadata { absl::string_view name; diff --git a/tpu_sync/telemetry/prometheus_exporter.cc b/tpu_sync/telemetry/prometheus_exporter.cc index 7e3a9eb4..66882975 100644 --- a/tpu_sync/telemetry/prometheus_exporter.cc +++ b/tpu_sync/telemetry/prometheus_exporter.cc @@ -15,18 +15,22 @@ #include "tpu_sync/telemetry/prometheus_exporter.h" #include +#include #include // NOLINT: Required by prometheus-cpp client API. #include #include #include #include "prometheus/counter.h" +#include "prometheus/exposer.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/log/log.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "tpu_sync/telemetry/metrics_backend.h" @@ -48,6 +52,13 @@ std::map ConvertLabels(LabelSpan labels) { return result; } +std::string JoinHostPort(absl::string_view host, int port) { + if (absl::StrContains(host, ':') && !absl::StartsWith(host, "[")) { + return absl::StrCat("[", host, "]:", port); + } + return absl::StrCat(host, ":", port); +} + } // namespace void PrometheusExporter::RegisterKnownFamilies() { @@ -87,13 +98,37 @@ void PrometheusExporter::RegisterKnownFamilies() { } } -PrometheusExporter::PrometheusExporter( - const prometheus::Histogram::BucketBoundaries& custom_buckets) +PrometheusExporter::PrometheusExporter(const ExporterOptions& options) : registry_(std::make_shared()), - default_buckets_(custom_buckets) { + default_buckets_(options.custom_buckets.begin(), + options.custom_buckets.end()), + options_(options) { RegisterKnownFamilies(); + + if (options_.port >= kMinPort && options_.port <= kMaxPort) { + std::string endpoint = JoinHostPort(options_.bind_address, options_.port); + // prometheus-cpp Exposer throws std::runtime_error on socket binding or + // initialization failure. Catching here prevents abnormal termination and + // allows graceful degradation with IsServerRunning() reporting false. + try { + exposer_ = std::make_unique(endpoint); + exposer_->RegisterCollectable(registry_); + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to start Prometheus HTTP exporter on " << endpoint + << ": " << e.what(); + exposer_.reset(); + } + } else if (options_.port != 0) { + LOG(WARNING) << "Invalid port configured for Prometheus HTTP exporter: " + << options_.port << ". Expected port in range [" << kMinPort + << ", " << kMaxPort << "]."; + } } +PrometheusExporter::~PrometheusExporter() = default; + +bool PrometheusExporter::IsServerRunning() const { return exposer_ != nullptr; } + prometheus::Family* PrometheusExporter::GetCounterFamily( absl::string_view name) const { auto it = counter_families_.find(name); diff --git a/tpu_sync/telemetry/prometheus_exporter.h b/tpu_sync/telemetry/prometheus_exporter.h index 401a8c25..29a6a49a 100644 --- a/tpu_sync/telemetry/prometheus_exporter.h +++ b/tpu_sync/telemetry/prometheus_exporter.h @@ -16,30 +16,29 @@ #define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_EXPORTER_H_ #include -#include #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_sync/telemetry/metrics_backend.h" -namespace tpu_raiden::telemetry { +// Forward declaration of prometheus::Exposer. +// Note: "prometheus/exposer.h" transitively includes CivetWeb ("civetweb.h"), +// which requires -fexceptions and -use_header_modules. Forward-declaring +// Exposer here prevents forcing those compiler constraints onto any library +// that includes "prometheus_exporter.h". +namespace prometheus { +class Exposer; +} // namespace prometheus -inline const prometheus::Histogram::BucketBoundaries& -DefaultHistogramBuckets() { - static const absl::NoDestructor - kBuckets(prometheus::Histogram::BucketBoundaries( - std::begin(kDefaultHistogramBuckets), - std::end(kDefaultHistogramBuckets))); - return *kBuckets; -} +namespace tpu_raiden::telemetry { // Custom MetricsBackend that formats and exports TPU Raiden metrics to // prometheus-cpp. @@ -53,10 +52,10 @@ DefaultHistogramBuckets() { // metric families. class PrometheusExporter : public MetricsBackend { public: - explicit PrometheusExporter(const prometheus::Histogram::BucketBoundaries& - custom_buckets = DefaultHistogramBuckets()); + explicit PrometheusExporter( + const ExporterOptions& options = ExporterOptions{}); - ~PrometheusExporter() override = default; + ~PrometheusExporter() override; PrometheusExporter(const PrometheusExporter&) = delete; PrometheusExporter& operator=(const PrometheusExporter&) = delete; @@ -74,6 +73,11 @@ class PrometheusExporter : public MetricsBackend { std::string GetTextSnapshot() const override; + bool IsServerRunning() const; + // Returns the bound HTTP port if the server is running, or 0 if disabled or + // failed to bind. Useful for tests and runtime port inspection. + int GetBoundPort() const { return IsServerRunning() ? options_.port : 0; } + const std::shared_ptr& GetRegistry() const { return registry_; } @@ -88,7 +92,10 @@ class PrometheusExporter : public MetricsBackend { prometheus::Family* GetHistogramFamily( absl::string_view name) const; std::shared_ptr registry_; - prometheus::Histogram::BucketBoundaries default_buckets_; + std::vector default_buckets_; + ExporterOptions options_; + + std::unique_ptr exposer_; absl::flat_hash_map*> diff --git a/tpu_sync/telemetry/prometheus_exporter_test.cc b/tpu_sync/telemetry/prometheus_exporter_test.cc index e6f1c6d5..0a5e05c3 100644 --- a/tpu_sync/telemetry/prometheus_exporter_test.cc +++ b/tpu_sync/telemetry/prometheus_exporter_test.cc @@ -25,18 +25,20 @@ #include "absl/strings/str_cat.h" #include "third_party/prometheus_cpp_client/core/include/prometheus/histogram.h" #include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/test_util.h" namespace tpu_raiden::telemetry { namespace { using ::testing::HasSubstr; -TEST(PrometheusExporterTest, DefaultHistogramBucketsMatchesMetricsApi) { - const prometheus::Histogram::BucketBoundaries& buckets = - DefaultHistogramBuckets(); +TEST(PrometheusExporterTest, ExporterOptionsDefaultHistogramBuckets) { + ExporterOptions options; std::vector expected(std::begin(kDefaultHistogramBuckets), std::end(kDefaultHistogramBuckets)); - EXPECT_EQ(buckets, expected); + std::vector actual(options.custom_buckets.begin(), + options.custom_buckets.end()); + EXPECT_EQ(actual, expected); } TEST(PrometheusExporterTest, RecordAndExportFormat) { @@ -159,5 +161,50 @@ TEST(PrometheusExporterTest, ConcurrentMetricUpdates) { kNumThreads * kIterations))); } +TEST(PrometheusExporterTest, ServerDisabledWhenPortIsZero) { + PrometheusExporter exporter_no_port; + EXPECT_FALSE(exporter_no_port.IsServerRunning()); + EXPECT_EQ(exporter_no_port.GetBoundPort(), 0); +} + +TEST(PrometheusExporterTest, ServerStartsWhenPortConfigured) { + int port = PickUnusedPort(); + ExporterOptions options{ + .bind_address = "127.0.0.1", + .port = port, + }; + PrometheusExporter exporter_with_port(options); + EXPECT_TRUE(exporter_with_port.IsServerRunning()); + EXPECT_EQ(exporter_with_port.GetBoundPort(), port); +} + +TEST(PrometheusExporterTest, ServerDisabledWhenPortIsOutOfRange) { + ExporterOptions options{ + .bind_address = "127.0.0.1", + .port = 99999, + }; + PrometheusExporter exporter(options); + EXPECT_FALSE(exporter.IsServerRunning()); + EXPECT_EQ(exporter.GetBoundPort(), 0); +} + +TEST(PrometheusExporterTest, ServerHandlesPortCollisionGracefully) { + int port = PickUnusedPort(); + ExporterOptions options{ + .bind_address = "127.0.0.1", + .port = port, + }; + PrometheusExporter first_exporter(options); + EXPECT_TRUE(first_exporter.IsServerRunning()); + + // A second distinct exporter instance attempting to bind to the already + // occupied port will fail socket binding (EADDRINUSE). It catches the + // exception gracefully, leaves IsServerRunning() as false, and reports bound + // port as 0. + PrometheusExporter second_exporter(options); + EXPECT_FALSE(second_exporter.IsServerRunning()); + EXPECT_EQ(second_exporter.GetBoundPort(), 0); +} + } // namespace } // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/python/BUILD b/tpu_sync/telemetry/python/BUILD index ff47614e..33494195 100644 --- a/tpu_sync/telemetry/python/BUILD +++ b/tpu_sync/telemetry/python/BUILD @@ -54,6 +54,7 @@ py_test( srcs = ["telemetry_binding_test.py"], deps = [ ":_telemetry_binding_test_ext", + "//third_party/py/portpicker", "@com_google_absl_py//absl/testing:absltest", ], ) diff --git a/tpu_sync/telemetry/python/telemetry_binding_test.py b/tpu_sync/telemetry/python/telemetry_binding_test.py index c04a9d9f..787d417b 100644 --- a/tpu_sync/telemetry/python/telemetry_binding_test.py +++ b/tpu_sync/telemetry/python/telemetry_binding_test.py @@ -15,6 +15,7 @@ """Tests for TPU Raiden Python telemetry bindings.""" from absl.testing import absltest +import portpicker from tpu_sync.telemetry.python import _telemetry_binding_test_ext as telemetry_ext @@ -189,6 +190,15 @@ def test_configure_telemetry_buffered(self): samples = telemetry_ext.get_and_reset_metric_samples() self.assertEqual(samples, {}) + def test_configure_telemetry_with_prometheus_port_env(self): + port = str(portpicker.pick_unused_port()) + with absltest.mock.patch.dict( + "os.environ", {"TPU_RAIDEN_PROMETHEUS_PORT": port} + ): + telemetry_ext.configure_telemetry(["prometheus"]) + snapshot = telemetry_ext.get_raiden_metrics_prometheus_text() + self.assertIn("# TYPE tpu_raiden_sent_bytes_total counter", snapshot) + if __name__ == "__main__": absltest.main() diff --git a/tpu_sync/telemetry/test_util.h b/tpu_sync/telemetry/test_util.h new file mode 100644 index 00000000..e2f87967 --- /dev/null +++ b/tpu_sync/telemetry/test_util.h @@ -0,0 +1,66 @@ +// 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_TEST_UTIL_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_TEST_UTIL_H_ + +#include +#include +#include + +namespace tpu_raiden::telemetry { + +// Picks an available ephemeral loopback TCP port for test servers. +// Retries up to 10 times with SO_REUSEADDR to avoid port collisions under load. +inline int PickUnusedPort() { + for (int attempt = 0; attempt < 10; ++attempt) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) { + continue; + } + + int on = 1; + if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0) { + close(fd); + continue; + } + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + addr.sin_port = 0; + + if (bind(fd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + close(fd); + continue; + } + + socklen_t len = sizeof(addr); + if (getsockname(fd, reinterpret_cast(&addr), &len) != 0) { + close(fd); + continue; + } + + int port = ntohs(addr.sin_port); + close(fd); + if (port > 0) { + return port; + } + } + return 0; +} + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_TEST_UTIL_H_