Skip to content

Commit

Permalink
Move tensorflow/core/lib/monitoring/{counter.h,gauge.h,sampler.h/.cc}…
Browse files Browse the repository at this point in the history 
… to tensorflow/tsl/lib/monitoring/

PiperOrigin-RevId: 474397486
  • Loading branch information
@tensorflower-gardener
tensorflower-gardener committed Sep 14, 2022
1 parent 5be9a59 commit 678eeed
Show file tree
Hide file tree
Showing 11 changed files with 1,068 additions and 795 deletions.
2 changes: 2 additions & 0 deletions tensorflow/core/framework/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ package(
default_visibility = [
"//tensorflow/core:__subpackages__",
"//tensorflow/security/fuzzing:__subpackages__",
# TODO(pedaveeraiah): to be removed when summary.proto.h deps moves to TSL
"//tensorflow/tsl/lib/monitoring:__subpackages__",
],
licenses = ["notice"],
)
Expand Down
2 changes: 2 additions & 0 deletions tensorflow/core/lib/histogram/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ package(
"//tensorflow/core:__pkg__",
# tensorflow/core/lib/monitoring:sampler uses histogram
"//tensorflow/core/lib/monitoring:__pkg__",
# TODO(pedaveeraiah): to be removed when histogram.h deps moves to TSL
"//tensorflow/tsl/lib/monitoring:__subpackages__",
],
licenses = ["notice"],
)
Expand Down
24 changes: 17 additions & 7 deletions tensorflow/core/lib/monitoring/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ package(
"//tensorflow/compiler/xla/service/gpu:__subpackages__",
# tensorflow/compiler/mlir/tfrt:tf_jitrt depends on this package
"//tensorflow/compiler/mlir/tfrt:__subpackages__",
"//tensorflow/tsl/lib/monitoring:__subpackages__",
],
licenses = ["notice"],
)
Expand Down Expand Up @@ -90,6 +91,11 @@ cc_library(
name = "collection_registry",
srcs = ["collection_registry.cc"],
hdrs = ["collection_registry.h"],
visibility = [
"//learning/brain/google/monitoring:__subpackages__",
"//tensorflow/core:__subpackages__",
"//tensorflow/tsl:__subpackages__",
],
deps = [
":collected_metrics",
":metric_def",
Expand Down Expand Up @@ -118,7 +124,7 @@ cc_library(
"//tensorflow/core/platform:macros",
"//tensorflow/core/platform:mutex",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/platform:types",
"//tensorflow/tsl/lib/monitoring:counter",
],
)

Expand All @@ -129,17 +135,18 @@ cc_library(
":collection_registry",
":metric_def",
"//tensorflow/core/lib/core:status",
"//tensorflow/core/platform",
"//tensorflow/core/platform:macros",
"//tensorflow/core/platform:mutex",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/platform:types",
"//tensorflow/tsl/lib/monitoring:gauge",
],
)

cc_library(
name = "metric_def",
hdrs = ["metric_def.h"],
visibility = [
"//learning/brain/google/monitoring:__subpackages__",
"//tensorflow/core:__subpackages__",
"//tensorflow/tsl/lib/monitoring:__subpackages__",
],
deps = [
":types",
"//tensorflow/core/framework:summary_proto_cc",
Expand All @@ -150,7 +157,6 @@ cc_library(

cc_library(
name = "sampler",
srcs = ["sampler.cc"],
hdrs = ["sampler.h"],
deps = [
":collection_registry",
Expand All @@ -163,6 +169,7 @@ cc_library(
"//tensorflow/core/platform:mutex",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/platform:types",
"//tensorflow/tsl/lib/monitoring:sampler",
],
)

Expand Down Expand Up @@ -212,6 +219,7 @@ filegroup(
"test_utils.h",
"timed.h",
"types.h",
"//tensorflow/tsl/lib/monitoring:mobile_srcs_only_runtime",
],
visibility = ["//tensorflow/core:__pkg__"],
)
Expand All @@ -230,6 +238,7 @@ filegroup(
"test_utils.h",
"timed.h",
"types.h",
"//tensorflow/tsl/lib/monitoring:legacy_lib_monitoring_lib_headers",
],
visibility = ["//tensorflow/core:__pkg__"],
)
Expand All @@ -247,6 +256,7 @@ filegroup(
"sampler.h",
"test_utils.h",
"types.h",
"//tensorflow/tsl/lib/monitoring:legacy_lib_monitoring_all_headers",
],
visibility = ["//tensorflow/core:__pkg__"],
)
Expand Down
211 changes: 7 additions & 204 deletions tensorflow/core/lib/monitoring/counter.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,225 +16,28 @@ limitations under the License.
#ifndef TENSORFLOW_CORE_LIB_MONITORING_COUNTER_H_
#define TENSORFLOW_CORE_LIB_MONITORING_COUNTER_H_

// clang-format off
// Required for IS_MOBILE_PLATFORM
#include "tensorflow/core/platform/platform.h"
// clang-format on

// We replace this implementation with a null implementation for mobile
// platforms.
#include "tensorflow/tsl/lib/monitoring/counter.h"
#ifdef IS_MOBILE_PLATFORM

#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"

namespace tensorflow {
namespace monitoring {

// CounterCell which has a null implementation.
class CounterCell {
public:
CounterCell() {}
~CounterCell() {}

void IncrementBy(int64 step) {}
int64 value() const { return 0; }

private:
TF_DISALLOW_COPY_AND_ASSIGN(CounterCell);
};

// Counter which has a null implementation.
template <int NumLabels>
class Counter {
public:
~Counter() {}

template <typename... MetricDefArgs>
static Counter* New(MetricDefArgs&&... metric_def_args) {
return new Counter<NumLabels>();
}

template <typename... Labels>
CounterCell* GetCell(const Labels&... labels) {
return &default_counter_cell_;
}

Status GetStatus() { return Status::OK(); }

private:
Counter() {}

CounterCell default_counter_cell_;

TF_DISALLOW_COPY_AND_ASSIGN(Counter);
};

} // namespace monitoring
} // namespace tensorflow

#else // IS_MOBILE_PLATFORM

#include <array>
#include <atomic>
#include <map>

#else
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/monitoring/collection_registry.h"
#include "tensorflow/core/lib/monitoring/metric_def.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"

#endif
// NOLINTBEGIN(misc-unused-using-decls)
namespace tensorflow {
namespace monitoring {

// CounterCell stores each value of an Counter.
//
// A cell can be passed off to a module which may repeatedly update it without
// needing further map-indexing computations. This improves both encapsulation
// (separate modules can own a cell each, without needing to know about the map
// to which both cells belong) and performance (since map indexing and
// associated locking are both avoided).
//
// This class is thread-safe.
class CounterCell {
public:
explicit CounterCell(int64_t value) : value_(value) {}
~CounterCell() {}

// Atomically increments the value by step.
// REQUIRES: Step be non-negative.
void IncrementBy(int64_t step);

// Retrieves the current value.
int64_t value() const;

private:
std::atomic<int64_t> value_;

TF_DISALLOW_COPY_AND_ASSIGN(CounterCell);
};

// A stateful class for updating a cumulative integer metric.
//
// This class encapsulates a set of values (or a single value for a label-less
// metric). Each value is identified by a tuple of labels. The class allows the
// user to increment each value.
//
// Counter allocates storage and maintains a cell for each value. You can
// retrieve an individual cell using a label-tuple and update it separately.
// This improves performance since operations related to retrieval, like
// map-indexing and locking, are avoided.
//
// This class is thread-safe.
template <int NumLabels>
class Counter {
public:
~Counter() {
// Deleted here, before the metric_def is destroyed.
registration_handle_.reset();
}

// Creates the metric based on the metric-definition arguments.
//
// Example;
// auto* counter_with_label = Counter<1>::New("/tensorflow/counter",
// "Tensorflow counter", "MyLabelName");
template <typename... MetricDefArgs>
static Counter* New(MetricDefArgs&&... metric_def_args);

// Retrieves the cell for the specified labels, creating it on demand if
// not already present.
template <typename... Labels>
CounterCell* GetCell(const Labels&... labels) TF_LOCKS_EXCLUDED(mu_);

Status GetStatus() { return status_; }

private:
explicit Counter(
const MetricDef<MetricKind::kCumulative, int64_t, NumLabels>& metric_def)
: metric_def_(metric_def),
registration_handle_(CollectionRegistry::Default()->Register(
&metric_def_, [&](MetricCollectorGetter getter) {
auto metric_collector = getter.Get(&metric_def_);

mutex_lock l(mu_);
for (const auto& cell : cells_) {
metric_collector.CollectValue(cell.first, cell.second.value());
}
})) {
if (registration_handle_) {
status_ = OkStatus();
} else {
status_ = Status(tensorflow::error::Code::ALREADY_EXISTS,
"Another metric with the same name already exists.");
}
}

mutable mutex mu_;

Status status_;

// The metric definition. This will be used to identify the metric when we
// register it for collection.
const MetricDef<MetricKind::kCumulative, int64_t, NumLabels> metric_def_;

std::unique_ptr<CollectionRegistry::RegistrationHandle> registration_handle_;

using LabelArray = std::array<string, NumLabels>;
std::map<LabelArray, CounterCell> cells_ TF_GUARDED_BY(mu_);

TF_DISALLOW_COPY_AND_ASSIGN(Counter);
};

////
// Implementation details follow. API readers may skip.
////

inline void CounterCell::IncrementBy(const int64_t step) {
DCHECK_LE(0, step) << "Must not decrement cumulative metrics.";
value_ += step;
}

inline int64_t CounterCell::value() const { return value_; }

template <int NumLabels>
template <typename... MetricDefArgs>
Counter<NumLabels>* Counter<NumLabels>::New(
MetricDefArgs&&... metric_def_args) {
return new Counter<NumLabels>(
MetricDef<MetricKind::kCumulative, int64_t, NumLabels>(
std::forward<MetricDefArgs>(metric_def_args)...));
}

template <int NumLabels>
template <typename... Labels>
CounterCell* Counter<NumLabels>::GetCell(const Labels&... labels)
TF_LOCKS_EXCLUDED(mu_) {
// Provides a more informative error message than the one during array
// construction below.
static_assert(sizeof...(Labels) == NumLabels,
"Mismatch between Counter<NumLabels> and number of labels "
"provided in GetCell(...).");

const LabelArray& label_array = {{labels...}};
mutex_lock l(mu_);
const auto found_it = cells_.find(label_array);
if (found_it != cells_.end()) {
return &(found_it->second);
}
return &(cells_
.emplace(std::piecewise_construct,
std::forward_as_tuple(label_array),
std::forward_as_tuple(0))
.first->second);
}
using tsl::monitoring::Counter;
using tsl::monitoring::CounterCell;

} // namespace monitoring
} // namespace tensorflow

#endif // IS_MOBILE_PLATFORM
// NOLINTEND(misc-unused-using-decls)
#endif // TENSORFLOW_CORE_LIB_MONITORING_COUNTER_H_

0 comments on commit 678eeed

Please sign in to comment.