Start a new profiling session that include all the hosts included in hostnames, for the time interval of duration_ms. Possibly save the profiling result in the directory specified by repository_root and session_id.
| 131 | // hostnames, for the time interval of duration_ms. Possibly save the profiling |
| 132 | // result in the directory specified by repository_root and session_id. |
| 133 | Status NewSession(const string& service_addr, |
| 134 | const std::vector<tensorflow::string>& hostnames, |
| 135 | int duration_ms, const string& repository_root, |
| 136 | const string& session_id, const ProfileOptions& opts) { |
| 137 | NewProfileSessionRequest new_session_request; |
| 138 | *new_session_request.mutable_request() = |
| 139 | PopulateProfileRequest(duration_ms, repository_root, session_id, opts); |
| 140 | new_session_request.set_repository_root(repository_root); |
| 141 | new_session_request.set_session_id(session_id); |
| 142 | for (const auto& hostname : hostnames) { |
| 143 | new_session_request.add_hosts(hostname); |
| 144 | } |
| 145 | |
| 146 | ::grpc::ClientContext context; |
| 147 | ::grpc::ChannelArguments channel_args; |
| 148 | // TODO(qiuminxu): use `NewHostPortGrpcChannel` instead once their |
| 149 | // `ValidateHostPortPair` checks for empty host string case. |
| 150 | channel_args.SetMaxReceiveMessageSize(std::numeric_limits<int32>::max()); |
| 151 | // TODO(jiesun): GRPC support following relevant naming scheme: |
| 152 | // 1. dns:///host:port |
| 153 | // 2. ipv4:host:port or ipv6:[host]:port |
| 154 | // We might need to change the prefix which depends on what TPU name resolver |
| 155 | // will give us. |
| 156 | std::unique_ptr<grpc::ProfileAnalysis::Stub> stub = |
| 157 | grpc::ProfileAnalysis::NewStub(::grpc::CreateCustomChannel( |
| 158 | "dns:///" + service_addr, ::grpc::InsecureChannelCredentials(), |
| 159 | channel_args)); |
| 160 | NewProfileSessionResponse new_session_response; |
| 161 | TF_RETURN_IF_ERROR(FromGrpcStatus( |
| 162 | stub->NewSession(&context, new_session_request, &new_session_response))); |
| 163 | |
| 164 | std::cout << "Profile session succeed for host(s):" |
| 165 | << absl::StrJoin(hostnames, ",") << std::endl; |
| 166 | if (new_session_response.empty_trace()) { |
| 167 | return Status(tensorflow::error::Code::UNAVAILABLE, |
| 168 | "No trace event is collected"); |
| 169 | } |
| 170 | return Status::OK(); |
| 171 | } |
| 172 | |
| 173 | // Creates an empty event file if not already exists, which indicates that we |
| 174 | // have a plugins/profile/ directory in the current logdir. |