Skip to content

Commit

Permalink
#tensorflow enable double-typed gauge cell
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 586776669
  • Loading branch information
jimlinntu authored and tensorflower-gardener committed Nov 30, 2023
1 parent eba052c commit 77ffe8b
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 3 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/profiler/utils/tfstreamz_utils.cc
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,9 @@ Status SerializeToXPlane(const std::vector<TfStreamzSnapshot>& snapshots,
xevent.AddStatValue(*metadata, *xplane.GetOrCreateStatMetadata(
point->string_value));
break;
case monitoring::ValueType::kDouble:
xevent.AddStatValue(*metadata, point->double_value);
break;
case monitoring::ValueType::kHistogram:
xevent.AddStatValue(*metadata, point->histogram_value);
break;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ struct Point {
int64_t int64_value;
string string_value;
bool bool_value;
double double_value;
HistogramProto histogram_value;
Percentiles percentiles_value;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,18 @@ inline void CollectValue(Percentiles value, Point* const point) {
point->percentiles_value = std::move(value);
}

template <>
inline void CollectValue(double value, Point* const point) {
point->value_type = ValueType::kDouble;
point->double_value = value;
}

template <>
inline void CollectValue(std::function<double()> value_fn, Point* const point) {
point->value_type = ValueType::kDouble;
point->double_value = value_fn();
}

// Used by the CollectionRegistry class to collect all the values of all the
// metrics in the registry. This is an implementation detail of the
// CollectionRegistry class, please do not depend on this.
Expand Down
6 changes: 4 additions & 2 deletions third_party/xla/third_party/tsl/tsl/lib/monitoring/gauge.h
Original file line number Diff line number Diff line change
Expand Up @@ -296,8 +296,10 @@ Gauge<ValueType, NumLabels>* Gauge<ValueType, NumLabels>::New(
std::is_same<ValueType, bool>::value ||
std::is_same<ValueType, std::function<int64_t()> >::value ||
std::is_same<ValueType, std::function<std::string()> >::value ||
std::is_same<ValueType, std::function<bool()> >::value,
"Gauge only allows bool, int64, and string types.");
std::is_same<ValueType, std::function<bool()> >::value ||
std::is_same<ValueType, std::function<double()> >::value ||
std::is_same<ValueType, double>::value,
"Gauge only allows bool, int64, double, and string types.");
return new Gauge<ValueType, NumLabels>(
MetricDef<MetricKind::kGauge, ValueType, NumLabels>(
std::forward<MetricDefArgs>(metric_def_args)...));
Expand Down
13 changes: 12 additions & 1 deletion third_party/xla/third_party/tsl/tsl/lib/monitoring/metric_def.h
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ enum class ValueType : int {
kHistogram,
kString,
kBool,
kPercentiles
kPercentiles,
kDouble
};

// Everything in the internal namespace is implementation details. Do not depend
Expand Down Expand Up @@ -97,6 +98,16 @@ inline ValueType GetValueType<std::function<bool()>>() {
return ValueType::kBool;
}

template <>
inline ValueType GetValueType<double>() {
return ValueType::kDouble;
}

template <>
inline ValueType GetValueType<std::function<double()>>() {
return ValueType::kDouble;
}

} // namespace internal

// Abstract base class for a metric definition.
Expand Down

0 comments on commit 77ffe8b

Please sign in to comment.