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 
16 #ifndef TENSORFLOW_CORE_PROFILER_RPC_CLIENT_REMOTE_PROFILER_SESSION_MANAGER_H_
17 #define TENSORFLOW_CORE_PROFILER_RPC_CLIENT_REMOTE_PROFILER_SESSION_MANAGER_H_
18 
19 #include <functional>
20 #include <memory>
21 #include <vector>
22 
23 #include "absl/strings/string_view.h"
24 #include "tensorflow/core/platform/macros.h"
25 #include "tensorflow/core/platform/mutex.h"
26 #include "tensorflow/core/platform/status.h"
27 #include "tensorflow/core/platform/thread_annotations.h"
28 #include "tensorflow/core/platform/types.h"
29 #include "tensorflow/core/profiler/rpc/client/profiler_client.h"
30 
31 namespace tensorflow {
32 namespace profiler {
33 
34 using AddressResolver = std::function<std::string(absl::string_view)>;
35 
36 // Manages one or more remote profiling sessions.
37 class RemoteProfilerSessionManager {
38  public:
39   struct Response {
40     std::string service_address;
41     std::unique_ptr<ProfileResponse> profile_response;
42     Status status;
43   };
44   // Instantiates a collection of RemoteProfilerSessions starts profiling on
45   // each of them immediately. Assumes that options have already been validated.
46   static std::unique_ptr<RemoteProfilerSessionManager> Create(
47       const RemoteProfilerSessionManagerOptions& options,
48       const ProfileRequest& request, tensorflow::Status& out_status,
49       AddressResolver resolver = nullptr);
50 
51   // Awaits for responses from remote profiler sessions and returns them as a
52   // list. Subsequent calls beyond the first will yield a list of errors.
53   std::vector<Response> WaitForCompletion();
54 
55   // Not copyable or movable.
56   RemoteProfilerSessionManager(const RemoteProfilerSessionManager&) = delete;
57   RemoteProfilerSessionManager operator=(const RemoteProfilerSessionManager&) =
58       delete;
59 
60   ~RemoteProfilerSessionManager();
61 
62  private:
63   explicit RemoteProfilerSessionManager(
64       RemoteProfilerSessionManagerOptions options, ProfileRequest request,
65       AddressResolver resolver);
66 
67   // Initialization of all client contexts.
68   Status Init();
69 
70   mutex mutex_;
71   // Remote profiler session options.
72   RemoteProfilerSessionManagerOptions options_ TF_GUARDED_BY(mutex_);
73   ProfileRequest request_ TF_GUARDED_BY(mutex_);
74   // List of clients, each connects to a profiling service.
75   std::vector<std::unique_ptr<RemoteProfilerSession>> clients_
76       TF_GUARDED_BY(mutex_);
77   // Resolves an address into a format that gRPC understands.
78   AddressResolver resolver_ TF_GUARDED_BY(mutex_);
79 };
80 
81 }  // namespace profiler
82 }  // namespace tensorflow
83 
84 #endif  // TENSORFLOW_CORE_PROFILER_RPC_CLIENT_REMOTE_PROFILER_SESSION_MANAGER_H_
85