1 | /* Copyright 2015 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_COMMON_RUNTIME_RENDEZVOUS_MGR_H_ |
17 | #define TENSORFLOW_CORE_COMMON_RUNTIME_RENDEZVOUS_MGR_H_ |
18 | |
19 | #include <string> |
20 | #include <unordered_map> |
21 | |
22 | #include "tensorflow/core/common_runtime/device_mgr.h" |
23 | #include "tensorflow/core/framework/local_rendezvous.h" |
24 | #include "tensorflow/core/framework/rendezvous.h" |
25 | #include "tensorflow/core/framework/tensor.h" |
26 | #include "tensorflow/core/lib/core/status.h" |
27 | #include "tensorflow/core/platform/macros.h" |
28 | #include "tensorflow/core/platform/mutex.h" |
29 | #include "tensorflow/core/platform/types.h" |
30 | |
31 | namespace tensorflow { |
32 | |
33 | // The IntraProcessRendezvous classes are implementations of a Rendezvous that |
34 | // expects all producers and consumers to be devices immediately accessible |
35 | // within the process. That is, it will never be necessary to perform an RPC to |
36 | // communicate with either. |
37 | // |
38 | // Buffering of Tensor values is delegated to a `LocalRendezvous`. An |
39 | // IntraProcessRendezvous. just adds functionality to coordinate multiple |
40 | // process-local devices. |
41 | |
42 | // Reference-counted implementation that may be shared between multiple threads. |
43 | class RefCountedIntraProcessRendezvous : public Rendezvous { |
44 | public: |
45 | explicit RefCountedIntraProcessRendezvous(const DeviceMgr* device_mgr); |
46 | |
47 | // Implementation of RendezvousInterface methods. |
48 | Status Send(const ParsedKey& key, const Rendezvous::Args& args, |
49 | const Tensor& val, const bool is_dead) override; |
50 | void RecvAsync(const ParsedKey& key, const Rendezvous::Args& args, |
51 | DoneCallback done) override; |
52 | void StartAbort(const Status& status) override; |
53 | |
54 | // Returns the member LocalRendezvous' status. |
55 | Status GetLocalRendezvousStatus(); |
56 | |
57 | inline void UpdateDeviceManager(DeviceMgr* device_mgr) { |
58 | device_mgr_ = device_mgr; |
59 | } |
60 | |
61 | private: |
62 | const DeviceMgr* device_mgr_; // Not owned. |
63 | LocalRendezvous local_; |
64 | |
65 | ~RefCountedIntraProcessRendezvous() override; |
66 | |
67 | TF_DISALLOW_COPY_AND_ASSIGN(RefCountedIntraProcessRendezvous); |
68 | }; |
69 | |
70 | // RefCountedIntraProcessRendezvous is aliased to IntraProcessRendezvous for |
71 | // backwards compatibility with existing users. |
72 | using IntraProcessRendezvous = RefCountedIntraProcessRendezvous; |
73 | |
74 | // Non-reference-counted implementation that may be stack-allocated for |
75 | // performance. |
76 | // |
77 | // Prefer to use PrivateIntraProcessRendezvous in new code. |
78 | class PrivateIntraProcessRendezvous : public RendezvousInterface { |
79 | public: |
80 | explicit PrivateIntraProcessRendezvous(const DeviceMgr* device_mgr); |
81 | ~PrivateIntraProcessRendezvous() override; |
82 | |
83 | // Implementation of RendezvousInterface methods. |
84 | Status Send(const ParsedKey& key, const Rendezvous::Args& args, |
85 | const Tensor& val, const bool is_dead) override; |
86 | void RecvAsync(const ParsedKey& key, const Rendezvous::Args& args, |
87 | DoneCallback done) override; |
88 | void StartAbort(const Status& status) override; |
89 | |
90 | private: |
91 | const DeviceMgr* device_mgr_; |
92 | LocalRendezvous local_; |
93 | |
94 | TF_DISALLOW_COPY_AND_ASSIGN(PrivateIntraProcessRendezvous); |
95 | }; |
96 | |
97 | } // end namespace tensorflow |
98 | |
99 | #endif // TENSORFLOW_CORE_COMMON_RUNTIME_RENDEZVOUS_MGR_H_ |
100 | |