1 | /* Copyright 2022 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 | |
16 | #include "tensorflow/core/activity_watcher/activity_utils.h" |
17 | |
18 | #include <memory> |
19 | #include <utility> |
20 | |
21 | namespace tensorflow { |
22 | namespace activity_watcher { |
23 | |
24 | std::unique_ptr<Activity> ActivityFromContext( |
25 | OpKernelContext* context, tensorflow::string name, |
26 | ActivityCategory category, Activity::Attributes additional_attributes) { |
27 | Activity::Attributes attributes(std::move(additional_attributes)); |
28 | if (context) { |
29 | attributes.merge(Activity::Attributes({ |
30 | {"node_name" , context->op_kernel().def().name()}, |
31 | {"step_id" , absl::StrCat(context->step_id())}, |
32 | {"device" , context->device()->name()}, |
33 | {"op" , context->op_kernel().def().op()}, |
34 | {"iter_num" , absl::StrCat(context->frame_iter().iter_id)}, |
35 | })); |
36 | } |
37 | |
38 | return std::make_unique<Activity>(name, category, std::move(attributes)); |
39 | } |
40 | |
41 | } // namespace activity_watcher |
42 | } // namespace tensorflow |
43 | |