Skip to content

Commit fed0513

Browse files
itaruncopybara-github
itarun
authored and
copybara-github
committed
Add multi-dimensional telemetry labels support to TPU Raiden.
- Add multi-dimensional label support to MetricMetadata and BufferedMetricsExporter. - Expose metric label metadata and labeled sample maps in Python bindings. PiperOrigin-RevId: 969218712
1 parent 152cfae commit fed0513

8 files changed

Lines changed: 514 additions & 49 deletions

‎tpu_sync/telemetry/BUILD‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ cc_library(
117117
":metrics_backend",
118118
"@com_google_absl//absl/base:core_headers",
119119
"@com_google_absl//absl/container:flat_hash_map",
120+
"@com_google_absl//absl/container:inlined_vector",
120121
"@com_google_absl//absl/strings",
121122
"@com_google_absl//absl/synchronization",
122123
"@com_google_absl//absl/types:span",
@@ -129,6 +130,9 @@ cc_test(
129130
deps = [
130131
":buffered_metrics_exporter",
131132
":metrics_backend",
133+
"@com_google_absl//absl/container:inlined_vector",
134+
"@com_google_absl//absl/strings",
135+
"@com_google_absl//absl/types:span",
132136
"@com_google_googletest//:gtest_main",
133137
],
134138
)

‎tpu_sync/telemetry/buffered_metrics_exporter.cc‎

Lines changed: 107 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,15 @@
1414

1515
#include "tpu_sync/telemetry/buffered_metrics_exporter.h"
1616

17+
#include <algorithm>
1718
#include <cstdint>
1819
#include <map>
1920
#include <memory>
2021
#include <string>
2122
#include <utility>
2223
#include <vector>
2324

25+
#include "absl/container/inlined_vector.h"
2426
#include "absl/strings/str_cat.h"
2527
#include "absl/strings/string_view.h"
2628
#include "absl/types/span.h"
@@ -32,21 +34,82 @@ namespace {
3234

3335
constexpr absl::string_view kMetricPrefix = "tpu_raiden_";
3436

37+
void AppendEscapedLabelValue(
38+
absl::string_view value,
39+
absl::InlinedVector<char, kDefaultInlinedLabelBufferSize>* out) {
40+
for (char c : value) {
41+
switch (c) {
42+
case '\\':
43+
out->push_back('\\');
44+
out->push_back('\\');
45+
break;
46+
case '"':
47+
out->push_back('\\');
48+
out->push_back('"');
49+
break;
50+
case '\n':
51+
out->push_back('\\');
52+
out->push_back('n');
53+
break;
54+
default:
55+
out->push_back(c);
56+
break;
57+
}
58+
}
59+
}
60+
3561
} // namespace
3662

63+
void FormatCanonicalLabels(
64+
LabelSpan labels,
65+
absl::InlinedVector<char, kDefaultInlinedLabelBufferSize>* out) {
66+
out->clear();
67+
if (labels.empty()) {
68+
return;
69+
}
70+
absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> sorted_labels(
71+
labels.begin(), labels.end());
72+
std::sort(sorted_labels.begin(), sorted_labels.end());
73+
74+
out->push_back('{');
75+
bool first = true;
76+
for (const auto& [key, value] : sorted_labels) {
77+
if (!first) {
78+
out->push_back(',');
79+
}
80+
first = false;
81+
out->insert(out->end(), key.begin(), key.end());
82+
out->push_back('=');
83+
out->push_back('"');
84+
AppendEscapedLabelValue(value, out);
85+
out->push_back('"');
86+
}
87+
out->push_back('}');
88+
}
89+
90+
std::string FormatCanonicalLabels(LabelSpan labels) {
91+
absl::InlinedVector<char, kDefaultInlinedLabelBufferSize> buf;
92+
FormatCanonicalLabels(labels, &buf);
93+
return std::string(buf.data(), buf.size());
94+
}
95+
3796
BufferedMetricsExporter::BufferedMetricsExporter(
3897
absl::Span<const MetricMetadata> metrics) {
3998
for (const auto& meta : metrics) {
4099
switch (meta.type) {
41100
case MetricType::kCounter:
42-
counters_.emplace(meta.name,
43-
std::make_unique<LockFreeCounterAccumulator>());
101+
counters_.emplace(
102+
meta.name,
103+
std::make_unique<
104+
MetricFamilyBuffer<LockFreeCounterAccumulator>>());
44105
break;
45106
case MetricType::kGauge:
46-
gauges_.emplace(meta.name, std::make_unique<QueueBuffer<>>());
107+
gauges_.emplace(
108+
meta.name, std::make_unique<MetricFamilyBuffer<QueueBuffer<>>>());
47109
break;
48110
case MetricType::kHistogram:
49-
histograms_.emplace(meta.name, std::make_unique<QueueBuffer<>>());
111+
histograms_.emplace(
112+
meta.name, std::make_unique<MetricFamilyBuffer<QueueBuffer<>>>());
50113
break;
51114
}
52115
}
@@ -57,15 +120,19 @@ void BufferedMetricsExporter::IncrementCounter(absl::string_view name,
57120
uint64_t val) const {
58121
auto it = counters_.find(name);
59122
if (it != counters_.end()) {
60-
it->second->Add(val);
123+
if (auto* acc = it->second->GetOrCreate(labels)) {
124+
acc->Add(val);
125+
}
61126
}
62127
}
63128

64129
void BufferedMetricsExporter::SetGauge(absl::string_view name, LabelSpan labels,
65130
double val) const {
66131
auto it = gauges_.find(name);
67132
if (it != gauges_.end()) {
68-
it->second->Push(val);
133+
if (auto* buf = it->second->GetOrCreate(labels)) {
134+
buf->Push(val);
135+
}
69136
}
70137
}
71138

@@ -74,36 +141,51 @@ void BufferedMetricsExporter::ObserveHistogram(absl::string_view name,
74141
double val) const {
75142
auto it = histograms_.find(name);
76143
if (it != histograms_.end()) {
77-
it->second->Push(val);
144+
if (auto* buf = it->second->GetOrCreate(labels)) {
145+
buf->Push(val);
146+
}
78147
}
79148
}
80149

81150
std::map<std::string, std::vector<double>>
82151
BufferedMetricsExporter::GetAndResetMetricSamples() {
83152
std::map<std::string, std::vector<double>> result;
84153

85-
for (const auto& [name, counter] : counters_) {
86-
uint64_t delta = counter->ExchangeAndReset();
87-
if (delta > 0) {
88-
std::string full_name = absl::StrCat(kMetricPrefix, name);
89-
result[full_name].push_back(static_cast<double>(delta));
90-
}
154+
for (const auto& [name, family_buffer] : counters_) {
155+
family_buffer->ForEachAccumulator(
156+
[&](absl::string_view canonical_labels,
157+
LockFreeCounterAccumulator* counter) {
158+
uint64_t delta = counter->ExchangeAndReset();
159+
if (delta > 0) {
160+
std::string full_name =
161+
absl::StrCat(kMetricPrefix, name, canonical_labels);
162+
result[full_name].push_back(static_cast<double>(delta));
163+
}
164+
});
91165
}
92166

93-
for (const auto& [name, gauge] : gauges_) {
94-
std::vector<double> samples = gauge->ExtractAndReset();
95-
if (!samples.empty()) {
96-
std::string full_name = absl::StrCat(kMetricPrefix, name);
97-
result[full_name] = std::move(samples);
98-
}
167+
for (const auto& [name, family_buffer] : gauges_) {
168+
family_buffer->ForEachAccumulator(
169+
[&](absl::string_view canonical_labels, QueueBuffer<>* gauge) {
170+
std::vector<double> samples = gauge->ExtractAndReset();
171+
if (!samples.empty()) {
172+
std::string full_name =
173+
absl::StrCat(kMetricPrefix, name, canonical_labels);
174+
result[full_name] = std::move(samples);
175+
}
176+
});
99177
}
100178

101-
for (const auto& [name, histogram] : histograms_) {
102-
std::vector<double> samples = histogram->ExtractAndReset();
103-
if (!samples.empty()) {
104-
std::string full_name = absl::StrCat(kMetricPrefix, name);
105-
result[full_name] = std::move(samples);
106-
}
179+
for (const auto& [name, family_buffer] : histograms_) {
180+
family_buffer->ForEachAccumulator(
181+
[&](absl::string_view canonical_labels, QueueBuffer<>* histogram) {
182+
std::vector<double> samples = histogram->ExtractAndReset();
183+
if (!samples.empty()) {
184+
std::string full_name =
185+
absl::StrCat(kMetricPrefix, name, canonical_labels);
186+
result[full_name] = std::move(samples);
187+
}
188+
});
107189
}
108190

109191
return result;

‎tpu_sync/telemetry/buffered_metrics_exporter.h‎

Lines changed: 91 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
#include "absl/base/thread_annotations.h"
2727
#include "absl/container/flat_hash_map.h"
28+
#include "absl/container/inlined_vector.h"
2829
#include "absl/strings/string_view.h"
2930
#include "absl/synchronization/mutex.h"
3031
#include "absl/types/span.h"
@@ -33,15 +34,30 @@
3334
namespace tpu_raiden::telemetry {
3435

3536
inline constexpr size_t kDefaultQueueBufferSize = 4096;
37+
inline constexpr size_t kMaxLabeledSeries = 512;
38+
inline constexpr size_t kDefaultInlinedLabelBufferSize = 64;
39+
inline constexpr size_t kDefaultInlinedLabelCapacity = 4;
40+
41+
// Formats canonical Prometheus label string for in-memory series
42+
// identification into an inlined stack buffer. Clears `out` and appends
43+
// "{key1=\"val1\",key2=\"val2\"}" sorted by key, with Prometheus character
44+
// escaping. Leaves `out` empty if labels are empty.
45+
void FormatCanonicalLabels(
46+
LabelSpan labels,
47+
absl::InlinedVector<char, kDefaultInlinedLabelBufferSize>* out);
48+
49+
// Formats canonical Prometheus label string for in-memory series
50+
// identification. Returns "{key1=\"val1\",key2=\"val2\"}" sorted by key, with
51+
// Prometheus character escaping. Returns "" if labels are empty.
52+
std::string FormatCanonicalLabels(LabelSpan labels);
3653

3754
// Fixed-capacity sample buffer for step-level metric observations.
3855
// Uses an absl::Mutex to ensure thread safety with O(1) buffer swapping
3956
// on extraction.
4057
template <size_t N = kDefaultQueueBufferSize>
4158
class QueueBuffer {
4259
public:
43-
explicit QueueBuffer(size_t max_capacity = N)
44-
: max_capacity_(max_capacity) {}
60+
explicit QueueBuffer(size_t max_capacity = N) : max_capacity_(max_capacity) {}
4561

4662
// Adds a value to the buffer. If the buffer is full, the value is dropped.
4763
void Push(double val) {
@@ -84,6 +100,70 @@ class LockFreeCounterAccumulator {
84100
std::atomic<uint64_t> value_{0};
85101
};
86102

103+
// Encapsulates all time-series buffers for a single metric family.
104+
template <typename AccumulatorType>
105+
class MetricFamilyBuffer {
106+
public:
107+
MetricFamilyBuffer() = default;
108+
109+
MetricFamilyBuffer(const MetricFamilyBuffer&) = delete;
110+
MetricFamilyBuffer& operator=(const MetricFamilyBuffer&) = delete;
111+
MetricFamilyBuffer(MetricFamilyBuffer&&) = delete;
112+
MetricFamilyBuffer& operator=(MetricFamilyBuffer&&) = delete;
113+
114+
// Returns accumulator for the given label span.
115+
// When labels are empty, returns unlabeled_ fast path without locks.
116+
AccumulatorType* GetOrCreate(LabelSpan labels) const {
117+
if (labels.empty()) {
118+
return &unlabeled_;
119+
}
120+
return GetOrCreateLabeled(labels);
121+
}
122+
123+
AccumulatorType* unlabeled() const { return &unlabeled_; }
124+
125+
template <typename Fn>
126+
void ForEachAccumulator(Fn&& fn) const {
127+
fn(/*canonical_labels=*/"", &unlabeled_);
128+
absl::ReaderMutexLock lock(labeled_mu_);
129+
for (const auto& [canonical_labels, acc] : labeled_) {
130+
fn(canonical_labels, acc.get());
131+
}
132+
}
133+
134+
private:
135+
AccumulatorType* GetOrCreateLabeled(LabelSpan labels) const {
136+
absl::InlinedVector<char, kDefaultInlinedLabelBufferSize>
137+
canonical_labels_buf;
138+
FormatCanonicalLabels(labels, &canonical_labels_buf);
139+
absl::string_view canonical_labels(canonical_labels_buf.data(),
140+
canonical_labels_buf.size());
141+
{
142+
absl::ReaderMutexLock lock(labeled_mu_);
143+
auto it = labeled_.find(canonical_labels);
144+
if (it != labeled_.end()) {
145+
return it->second.get();
146+
}
147+
}
148+
absl::MutexLock lock(labeled_mu_);
149+
auto it = labeled_.find(canonical_labels);
150+
if (it != labeled_.end()) {
151+
return it->second.get();
152+
}
153+
if (labeled_.size() >= kMaxLabeledSeries) {
154+
return nullptr;
155+
}
156+
auto [insert_it, _] = labeled_.try_emplace(
157+
std::string(canonical_labels), std::make_unique<AccumulatorType>());
158+
return insert_it->second.get();
159+
}
160+
161+
mutable AccumulatorType unlabeled_;
162+
mutable absl::Mutex labeled_mu_;
163+
mutable absl::flat_hash_map<std::string, std::unique_ptr<AccumulatorType>>
164+
labeled_ ABSL_GUARDED_BY(labeled_mu_);
165+
};
166+
87167
// MetricsBackend for step-level sample buffering.
88168
//
89169
// Thread Safety:
@@ -116,10 +196,16 @@ class BufferedMetricsExporter : public MetricsBackend {
116196
override;
117197

118198
private:
119-
absl::flat_hash_map<std::string, std::unique_ptr<LockFreeCounterAccumulator>>
199+
absl::flat_hash_map<
200+
std::string,
201+
std::unique_ptr<MetricFamilyBuffer<LockFreeCounterAccumulator>>>
120202
counters_;
121-
absl::flat_hash_map<std::string, std::unique_ptr<QueueBuffer<>>> gauges_;
122-
absl::flat_hash_map<std::string, std::unique_ptr<QueueBuffer<>>> histograms_;
203+
absl::flat_hash_map<std::string,
204+
std::unique_ptr<MetricFamilyBuffer<QueueBuffer<>>>>
205+
gauges_;
206+
absl::flat_hash_map<std::string,
207+
std::unique_ptr<MetricFamilyBuffer<QueueBuffer<>>>>
208+
histograms_;
123209
};
124210

125211
} // namespace tpu_raiden::telemetry

0 commit comments

Comments
 (0)