1/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15#ifndef TENSORFLOW_CORE_ACTIVITY_WATCHER_ACTIVITY_H_
16#define TENSORFLOW_CORE_ACTIVITY_WATCHER_ACTIVITY_H_
17
18#include <atomic>
19#include <functional>
20#include <memory>
21#include <utility>
22
23#include "absl/container/flat_hash_map.h"
24#include "tensorflow/core/platform/macros.h"
25#include "tensorflow/core/platform/types.h"
26
27namespace tensorflow {
28class CoordinationServiceAgent;
29
30namespace activity_watcher {
31
32using ActivityId = uint64;
33constexpr ActivityId kActivityNotRecorded = 0;
34constexpr int kWatcherDisabled = 0;
35
36enum ActivityCategory {
37 kCollective = 0,
38 kRemoteFunction = 1,
39 kMisc = 2,
40 kDatasetOp = 3,
41 kTpuOp = 4,
42};
43
44static tensorflow::string ToString(ActivityCategory category) {
45 switch (category) {
46 case ActivityCategory::kCollective:
47 return "Collective";
48 case ActivityCategory::kRemoteFunction:
49 return "Remote Function";
50 case ActivityCategory::kMisc:
51 return "Miscellaneous";
52 case ActivityCategory::kDatasetOp:
53 return "Dataset Op";
54 case ActivityCategory::kTpuOp:
55 return "TPU Op";
56 }
57}
58
59// An activity to be recorded.
60struct Activity {
61 using Attributes =
62 absl::flat_hash_map<tensorflow::string, tensorflow::string>;
63 // A human readable title of the activity.
64 tensorflow::string title;
65 // The category of the activity.
66 ActivityCategory category = ActivityCategory::kMisc;
67 // Key/value pairs that are attached to the activity.
68 Attributes attributes;
69 Activity() = default;
70 Activity(tensorflow::string title, ActivityCategory category)
71 : title(std::move(title)), category(category) {}
72 Activity(tensorflow::string title, ActivityCategory category,
73 Attributes attributes)
74 : title(std::move(title)),
75 category(category),
76 attributes(std::move(attributes)) {}
77};
78
79// Enable activity wathcer to send own workers activities to coordination
80// service and also fetch all workers' activities.
81void MaybeEnableMultiWorkersWatching(CoordinationServiceAgent* agent);
82
83namespace tfw_internal {
84
85#if !defined(IS_MOBILE_PLATFORM)
86
87// Records an activity start without checking whether the watcher is enabled.
88ActivityId RecordActivityStart(std::unique_ptr<Activity> activity);
89// Records an activity end without checking whether the activity_id is valid.
90void RecordActivityEnd(ActivityId activity_id);
91
92TF_EXPORT extern std::atomic<int> g_watcher_level;
93
94// Returns whether the activitity watcher is enabled.
95inline bool WatcherEnabled(int level = 1) {
96 return g_watcher_level.load(std::memory_order_acquire) >= level;
97}
98
99#endif
100
101// NOTE: Borrowed from boost C++ libraries because std::is_invocable_r is not
102// available in Android NDK.
103template <typename R, typename F, typename... Args>
104struct is_invocable_r
105 : std::is_constructible<
106 std::function<R(Args...)>,
107 std::reference_wrapper<typename std::remove_reference<F>::type>> {};
108
109} // namespace tfw_internal
110
111template <typename F>
112constexpr bool is_activity_generator =
113 tfw_internal::is_invocable_r<std::unique_ptr<Activity>, F>::value;
114
115// Records an activity explicitly. Useful when the start and end of an activity
116// happen in different threads. Generates the Activity only if activity
117// watching is enabled, useful for avoiding expensive operations when activity
118// watching is disabled.
119// Example Usage:
120// auto aid = ActivityStart([&]() {
121// return std::make_unique<Activity>(
122// op_name, category,
123// Activity::Attributes{{"key1", value1}, {"key2", value2}});
124// }, /*level=*/2);
125// DoSomething();
126// ActivityEnd(aid);
127template <
128 typename ActivityGenerator,
129 std::enable_if_t<is_activity_generator<ActivityGenerator>, bool> = true>
130inline ActivityId ActivityStart(ActivityGenerator&& gen, int level = 1) {
131#if !defined(IS_MOBILE_PLATFORM)
132 if (TF_PREDICT_FALSE(tfw_internal::WatcherEnabled(level))) {
133 return tfw_internal::RecordActivityStart(
134 std::forward<ActivityGenerator>(gen)());
135 }
136#endif
137 return kActivityNotRecorded;
138}
139
140inline void ActivityEnd(ActivityId id) {
141#if !defined(IS_MOBILE_PLATFORM)
142 if (TF_PREDICT_FALSE(id != kActivityNotRecorded)) {
143 tfw_internal::RecordActivityEnd(id);
144 }
145#endif
146}
147
148// ActivityScope marks a scope as an activity and record it with a global
149// ActivityRecorder.
150// Example Usage:
151// {
152// ActivityScope activity_scope([&]() {
153// return std::make_unique<Activity>(
154// op_name, ActivityCategory::kMisc,
155// Activity::Attributes{{"key1", value1}, {"key2", value2}});
156// }, /*level=*/2);
157// DoSomething();
158// }
159class ActivityScope {
160 public:
161 template <
162 typename ActivityGenerator,
163 std::enable_if_t<is_activity_generator<ActivityGenerator>, bool> = true>
164 explicit ActivityScope(ActivityGenerator&& gen, int level = 1) {
165 activity_id_ = ActivityStart(std::forward<ActivityGenerator>(gen), level);
166 }
167 ~ActivityScope() { ActivityEnd(activity_id_); }
168
169 private:
170 ActivityId activity_id_;
171 TF_DISALLOW_COPY_AND_ASSIGN(ActivityScope);
172};
173
174} // namespace activity_watcher
175} // namespace tensorflow
176
177#endif // TENSORFLOW_CORE_ACTIVITY_WATCHER_ACTIVITY_H_
178