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 #ifndef TENSORFLOW_CORE_PROFILER_LIB_PROFILER_CONTROLLER_H_ 16 #define TENSORFLOW_CORE_PROFILER_LIB_PROFILER_CONTROLLER_H_ 17 18 #include <memory> 19 20 #include "tensorflow/core/platform/status.h" 21 #include "tensorflow/core/profiler/lib/profiler_interface.h" 22 #include "tensorflow/core/profiler/protobuf/xplane.pb.h" 23 24 namespace tensorflow { 25 namespace profiler { 26 27 // Decorator for xprof profiler plugins. 28 // 29 // Tracks that calls to the underlying profiler interface functions are made 30 // in the expected order: Start, Stop and CollectData. Making the calls 31 // in a different order causes them to be aborted. 32 // 33 // Calls made in the right order will be aborted if one of the calls to the 34 // decorated profiler interface fails, and no more calls will be forwarded to 35 // the decorated profiler. 36 class ProfilerController : public ProfilerInterface { 37 public: 38 explicit ProfilerController(std::unique_ptr<ProfilerInterface> profiler); 39 ~ProfilerController() override; 40 41 Status Start() override; 42 43 Status Stop() override; 44 45 Status CollectData(XSpace* space) override; 46 47 private: 48 enum class ProfilerState { 49 kInit = 0, 50 kStart = 1, 51 kStop = 2, 52 kCollectData = 3, 53 }; 54 55 ProfilerState state_ = ProfilerState::kInit; 56 std::unique_ptr<ProfilerInterface> profiler_; 57 Status status_; // result of calls to profiler_ 58 }; 59 60 } // namespace profiler 61 } // namespace tensorflow 62 63 #endif // TENSORFLOW_CORE_PROFILER_LIB_PROFILER_CONTROLLER_H_ 64