1 | /* Copyright 2016 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_DISTRIBUTED_RUNTIME_MASTER_INTERFACE_H_ |
17 | #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_MASTER_INTERFACE_H_ |
18 | |
19 | #include "tensorflow/core/distributed_runtime/call_options.h" |
20 | #include "tensorflow/core/distributed_runtime/message_wrappers.h" |
21 | #include "tensorflow/core/distributed_runtime/request_id.h" |
22 | #include "tensorflow/core/lib/core/errors.h" |
23 | #include "tensorflow/core/lib/core/status.h" |
24 | #include "tensorflow/core/protobuf/master.pb.h" |
25 | |
26 | namespace tensorflow { |
27 | |
28 | // Abstract interface for communicating with the TensorFlow Master service. |
29 | // |
30 | // This interface supports both RPC-based master implementations, and |
31 | // in-process master implementations that do not require an RPC |
32 | // roundtrip. |
33 | class MasterInterface { |
34 | public: |
35 | virtual ~MasterInterface() {} |
36 | virtual Status CreateSession(CallOptions* call_options, |
37 | const CreateSessionRequest* request, |
38 | CreateSessionResponse* response) = 0; |
39 | |
40 | virtual Status ExtendSession(CallOptions* call_options, |
41 | const ExtendSessionRequest* request, |
42 | ExtendSessionResponse* response) = 0; |
43 | |
44 | virtual Status PartialRunSetup(CallOptions* call_options, |
45 | const PartialRunSetupRequest* request, |
46 | PartialRunSetupResponse* response) { |
47 | return errors::Unimplemented("Partial run not implemented for this master" ); |
48 | } |
49 | |
50 | virtual Status RunStep(CallOptions* call_options, |
51 | RunStepRequestWrapper* request, |
52 | MutableRunStepResponseWrapper* response) = 0; |
53 | |
54 | virtual Status RunStep(CallOptions* call_options, |
55 | const RunStepRequest* request, |
56 | RunStepResponse* response) { |
57 | std::unique_ptr<RunStepRequestWrapper> wrapped_request( |
58 | new ProtoRunStepRequest(request)); |
59 | std::unique_ptr<MutableRunStepResponseWrapper> wrapped_response( |
60 | new NonOwnedProtoRunStepResponse(response)); |
61 | return RunStep(call_options, wrapped_request.get(), wrapped_response.get()); |
62 | } |
63 | |
64 | // Returns a request object for use in calls to |
65 | // `RunStep()`. Ownership is transferred to the caller. |
66 | // |
67 | // The message returned from this method must only be used in a |
68 | // `RunStep()` call on the same `MasterInterface` instance. |
69 | virtual MutableRunStepRequestWrapper* CreateRunStepRequest() { |
70 | MutableProtoRunStepRequest* ret = new MutableProtoRunStepRequest; |
71 | ret->request_.set_request_id(GetUniqueRequestId()); |
72 | return ret; |
73 | } |
74 | |
75 | // Returns a response object for use in calls to |
76 | // `RunStep()`. Ownership is transferred to the caller. |
77 | // |
78 | // The message returned from this method must only be used in a |
79 | // `RunStep()` call on the same `MasterInterface` instance. |
80 | virtual MutableRunStepResponseWrapper* CreateRunStepResponse() { |
81 | return new OwnedProtoRunStepResponse; |
82 | } |
83 | |
84 | virtual Status CloseSession(CallOptions* call_options, |
85 | const CloseSessionRequest* request, |
86 | CloseSessionResponse* response) = 0; |
87 | |
88 | virtual Status ListDevices(CallOptions* call_options, |
89 | const ListDevicesRequest* request, |
90 | ListDevicesResponse* response) = 0; |
91 | |
92 | virtual Status Reset(CallOptions* call_options, const ResetRequest* request, |
93 | ResetResponse* response) = 0; |
94 | |
95 | virtual Status MakeCallable(CallOptions* call_options, |
96 | const MakeCallableRequest* request, |
97 | MakeCallableResponse* response) = 0; |
98 | virtual Status RunCallable(CallOptions* call_options, |
99 | const RunCallableRequest* request, |
100 | RunCallableResponse* response) = 0; |
101 | virtual Status ReleaseCallable(CallOptions* call_options, |
102 | const ReleaseCallableRequest* request, |
103 | ReleaseCallableResponse* response) = 0; |
104 | |
105 | protected: |
106 | // NOTE: This should only be called by implementations of this |
107 | // interface whose CreateRunStepResponse() method returns a |
108 | // proto-based wrappers for the RunStepResponse message. |
109 | RunStepResponse* get_proto_from_wrapper( |
110 | MutableRunStepResponseWrapper* wrapper) { |
111 | return wrapper->get_proto(); |
112 | } |
113 | }; |
114 | |
115 | } // namespace tensorflow |
116 | |
117 | #endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_MASTER_INTERFACE_H_ |
118 | |