1 /* Copyright 2020 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 // GRPC client to perform on-demand profiling 16 17 #ifndef TENSORFLOW_CORE_PROFILER_RPC_CLIENT_PROFILER_CLIENT_H_ 18 #define TENSORFLOW_CORE_PROFILER_RPC_CLIENT_PROFILER_CLIENT_H_ 19 20 #include <memory> 21 #include <string> 22 23 #include "absl/strings/string_view.h" 24 #include "absl/time/time.h" 25 #include "tensorflow/core/platform/status.h" 26 #include "tensorflow/core/profiler/profiler_analysis.grpc.pb.h" 27 #include "tensorflow/core/profiler/profiler_service.grpc.pb.h" 28 29 namespace tensorflow { 30 namespace profiler { 31 32 // Note that tensorflow/tools/def_file_filter/symbols_pybind.txt is incompatible 33 // with absl::string_view. 34 Status ProfileGrpc(const std::string& service_address, 35 const ProfileRequest& request, ProfileResponse* response); 36 37 Status NewSessionGrpc(const std::string& service_address, 38 const NewProfileSessionRequest& request, 39 NewProfileSessionResponse* response); 40 41 Status MonitorGrpc(const std::string& service_address, 42 const MonitorRequest& request, MonitorResponse* response); 43 44 class RemoteProfilerSession { 45 public: 46 // Creates an instance and starts a remote profiling session immediately. 47 // This is a non-blocking call and does not wait for a response. 48 // Response must outlive the instantiation. 49 static std::unique_ptr<RemoteProfilerSession> Create( 50 const std::string& service_address, absl::Time deadline, 51 const ProfileRequest& profile_request); 52 53 // Not copyable or movable. 54 RemoteProfilerSession(const RemoteProfilerSession&) = delete; 55 RemoteProfilerSession operator=(const RemoteProfilerSession&) = delete; 56 57 ~RemoteProfilerSession(); 58 GetServiceAddress()59 absl::string_view GetServiceAddress() const { return service_address_; } 60 61 // Blocks until a response has been received or until deadline expiry, 62 // whichever is first. Subsequent calls after the first will yield nullptr and 63 // an error status. 64 std::unique_ptr<ProfileResponse> WaitForCompletion(Status& out_status); 65 66 private: 67 explicit RemoteProfilerSession(const std::string& service_addr, 68 absl::Time deadline, 69 const ProfileRequest& profile_request); 70 71 // Starts a remote profiling session. This is a non-blocking call. 72 // Will be called exactly once during instantiation. 73 // RPC will write to response.profile_response eagerly. However, since 74 // response.status requires a conversion from grpc::Status, it can only be 75 // evaluated lazily at WaitForCompletion() time. 76 void ProfileAsync(); 77 78 Status status_on_completion_; 79 std::unique_ptr<ProfileResponse> response_; 80 // Client address and connection attributes. 81 std::string service_address_; 82 std::unique_ptr<grpc::ProfilerService::Stub> stub_; 83 absl::Time deadline_; 84 ::grpc::ClientContext grpc_context_; 85 std::unique_ptr<::grpc::ClientAsyncResponseReader<ProfileResponse>> rpc_; 86 ::grpc::Status grpc_status_ = ::grpc::Status::OK; 87 88 // Asynchronous completion queue states. 89 ::grpc::CompletionQueue cq_; 90 91 ProfileRequest profile_request_; 92 }; 93 94 } // namespace profiler 95 } // namespace tensorflow 96 97 #endif // TENSORFLOW_CORE_PROFILER_RPC_CLIENT_PROFILER_CLIENT_H_ 98