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 | |
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; |
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. |
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 | |
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 | } |
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 | |