1 /* Copyright 2017 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_CONTRIB_GDR_GDR_MEMORY_MANAGER_H_ 17 #define TENSORFLOW_CONTRIB_GDR_GDR_MEMORY_MANAGER_H_ 18 19 #include "google/protobuf/any.pb.h" 20 #include "tensorflow/core/lib/core/status.h" 21 22 namespace tensorflow { 23 24 class Device; 25 class DeviceContext; 26 class Tensor; 27 28 // Abstract interface that handles out-of-band tensor transport. 29 // 30 // The transport options are encoded into a protocol buffer and transmitted via 31 // some other communication channels like RPC. 32 // See RecvTensorRequest in tensorflow/core/protobuf/worker.proto 33 class RemoteMemoryManager { 34 public: ~RemoteMemoryManager()35 virtual ~RemoteMemoryManager() {} 36 virtual Status Init() = 0; 37 virtual void Run() = 0; 38 virtual void Stop() = 0; 39 40 // Encodes the tensor information to an arbitrary protocol buffer 41 // The protocol buffer needs to be transmitted via some other channel 42 virtual void TransportOptionsFromTensor( 43 ::google::protobuf::Any* mutable_transport_options, const Tensor& tensor, 44 Device* device, DeviceContext* device_context, bool on_host, 45 StatusCallback done) = 0; 46 47 // Retrieve the tensor from the encoded protocol buffer 48 // Note that the tensor has to be allocated, but not initialized 49 virtual void TensorFromTransportOptions( 50 Tensor* tensor, const ::google::protobuf::Any& transport_options, 51 Device* device, DeviceContext* device_context, bool on_host, 52 StatusCallback done) = 0; 53 }; 54 55 RemoteMemoryManager* CreateRemoteMemoryManager(const string& host, 56 const string& port); 57 58 } // namespace tensorflow 59 60 #endif // TENSORFLOW_CONTRIB_GDR_GDR_MEMORY_MANAGER_H_ 61