1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations 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
27 namespace tensorflow {
28 class CoordinationServiceAgent;
29
30 namespace activity_watcher {
31
32 using ActivityId = uint64;
33 constexpr ActivityId kActivityNotRecorded = 0;
34 constexpr int kWatcherDisabled = 0;
35
36 enum ActivityCategory {
37 kCollective = 0,
38 kRemoteFunction = 1,
39 kMisc = 2,
40 kDatasetOp = 3,
41 kTpuOp = 4,
42 };
43
ToString(ActivityCategory category)44 static 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.
60 struct 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;
ActivityActivity70 Activity(tensorflow::string title, ActivityCategory category)
71 : title(std::move(title)), category(category) {}
ActivityActivity72 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.
81 void MaybeEnableMultiWorkersWatching(CoordinationServiceAgent* agent);
82
83 namespace tfw_internal {
84
85 #if !defined(IS_MOBILE_PLATFORM)
86
87 // Records an activity start without checking whether the watcher is enabled.
88 ActivityId RecordActivityStart(std::unique_ptr<Activity> activity);
89 // Records an activity end without checking whether the activity_id is valid.
90 void RecordActivityEnd(ActivityId activity_id);
91
92 TF_EXPORT extern std::atomic<int> g_watcher_level;
93
94 // Returns whether the activitity watcher is enabled.
95 inline 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.
103 template <typename R, typename F, typename... Args>
104 struct 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
111 template <typename F>
112 constexpr 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);
127 template <
128 typename ActivityGenerator,
129 std::enable_if_t<is_activity_generator<ActivityGenerator>, bool> = true>
130 inline 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
ActivityEnd(ActivityId id)140 inline 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 // }
159 class 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 }
~ActivityScope()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