1 /* Copyright 2019 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 #ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_TENSOR_HANDLE_DATA_H_ 16 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_TENSOR_HANDLE_DATA_H_ 17 18 #include "tensorflow/core/common_runtime/eager/context.h" 19 #include "tensorflow/core/framework/tensor.h" 20 #include "tensorflow/core/lib/core/status.h" 21 22 namespace tensorflow { 23 24 // Remote Tensor Handle: A handle to a Tensor on a remote host. Note that only 25 // the shape is known. 26 class RemoteTensorHandleData { 27 public: 28 // Constructor for lazy remote handles. A lazy remote handle is created on 29 // a remote worker with an op_id and an output_num. It doesn't control the 30 // lifetime of a remote handle that it refers to. If it refers to a remote 31 // function input, it's sent by a client which won't serialize it until 32 // the corresponding remote tensor is ready. So the remote tensor should be 33 // ready when we create a lazy remote handle. If it refers to a remote output, 34 // it's not ready until the shape is set. 35 RemoteTensorHandleData(int64 op_id, int output_num, uint64 context_view_id, 36 bool is_ready); 37 // Constructor for unshaped remote handles. It controls the lifetime of a 38 // remote handel that it refers to. 39 RemoteTensorHandleData(int64 op_id, int output_num, const string& remote_task, 40 EagerContext* ctx); 41 ~RemoteTensorHandleData(); 42 43 // A remote tensor handle does not have a Tensor object, hence it can only 44 // support the shape requests. 45 Status Shape(TensorShape* shape) const; 46 Status NumDims(int* num_dims) const; 47 Status Dim(int dim_index, int64* dim) const; 48 Status NumElements(int64* num_elements) const; Unprotect()49 Status Unprotect() { return Status::OK(); } 50 51 bool IsReady() const; 52 Status WaitReady(const char* caller) const; 53 Status SetShape(const TensorShape& shape); 54 Status SetShapeAndRemoteTask(const TensorShape& shape, 55 const string& remote_task); 56 void Poison(Status status); 57 Status IsPoisoned() const; 58 59 string DebugString() const; 60 61 // Return the op id and output num. If wait_util_ready is true, block until 62 // the remote tensor is ready on a remote worker. 63 Status OpIdAndOutputNum(const bool wait_util_ready, int64* op_id, 64 int32* output_num) const; 65 context_view_id()66 uint64 context_view_id() const { return context_view_id_; } 67 68 private: 69 mutable mutex mu_; 70 bool is_ready_ TF_GUARDED_BY(mu_); 71 Status is_poisoned_ TF_GUARDED_BY(mu_); 72 TensorShape shape_ TF_GUARDED_BY(mu_); 73 74 // IDs required when this class is representing a remote tensor handle. 75 const int64 op_id_; 76 const int32 output_num_; 77 string remote_task_ TF_GUARDED_BY(mu_); 78 uint64 context_id_; 79 uint64 context_view_id_; 80 EagerContext* ctx_; 81 }; 82 83 } // namespace tensorflow 84 85 #endif // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_TENSOR_HANDLE_DATA_H_ 86