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_FRAMEWORK_QUEUE_INTERFACE_H_ 17 #define TENSORFLOW_FRAMEWORK_QUEUE_INTERFACE_H_ 18 19 #include <string> 20 #include <vector> 21 22 #include "tensorflow/core/framework/op_kernel.h" 23 #include "tensorflow/core/framework/resource_mgr.h" 24 #include "tensorflow/core/framework/tensor.h" 25 #include "tensorflow/core/framework/types.h" 26 #include "tensorflow/core/platform/types.h" 27 28 namespace tensorflow { 29 30 // All implementations must be thread-safe. 31 class QueueInterface : public ResourceBase { 32 public: 33 typedef std::vector<Tensor> Tuple; 34 typedef AsyncOpKernel::DoneCallback DoneCallback; 35 typedef std::function<void(const Tuple&)> CallbackWithTuple; 36 37 virtual Status ValidateTuple(const Tuple& tuple) = 0; 38 virtual Status ValidateManyTuple(const Tuple& tuple) = 0; 39 40 // Stashes a function object for future execution, that will eventually 41 // enqueue the tuple of tensors into the queue, and returns immediately. The 42 // function object is guaranteed to call 'callback'. 43 virtual void TryEnqueue(const Tuple& tuple, OpKernelContext* ctx, 44 DoneCallback callback) = 0; 45 46 // Same as above, but the component tensors are sliced along the 0th dimension 47 // to make multiple queue-element components. 48 virtual void TryEnqueueMany(const Tuple& tuple, OpKernelContext* ctx, 49 DoneCallback callback) = 0; 50 51 // Stashes a function object for future execution, that will eventually 52 // dequeue an element from the queue and call 'callback' with that tuple 53 // element as argument. 54 virtual void TryDequeue(OpKernelContext* ctx, CallbackWithTuple callback) = 0; 55 56 // Same as above, but the stashed function object will attempt to dequeue 57 // num_elements items. If allow_small_batch is true, and the Queue is 58 // closed but at least 1 element is available, there is no blocking 59 // and between 1 and num_elements items are immediately returned. 60 // If the queue does not support the allow_small_batch flag will 61 // return an Unimplemented error. 62 virtual void TryDequeueMany(int num_elements, OpKernelContext* ctx, 63 bool allow_small_batch, 64 CallbackWithTuple callback) = 0; 65 66 // Signals that no more elements will be enqueued, and optionally 67 // cancels pending Enqueue(Many) operations. 68 // 69 // After calling this function, subsequent calls to Enqueue(Many) 70 // will fail. If `cancel_pending_enqueues` is true, all pending 71 // calls to Enqueue(Many) will fail as well. 72 // 73 // After calling this function, all current and subsequent calls to 74 // Dequeue(Many) will fail instead of blocking (though they may 75 // succeed if they can be satisfied by the elements in the queue at 76 // the time it was closed). 77 virtual void Close(OpKernelContext* ctx, bool cancel_pending_enqueues, 78 DoneCallback callback) = 0; 79 80 // Returns true if a given queue is closed and false if it is open. 81 virtual bool is_closed() const = 0; 82 83 // Assuming *this represents a shared queue, verify that it matches 84 // another instantiation indicated by node_def. 85 virtual Status MatchesNodeDef(const NodeDef& node_def) = 0; 86 87 // Returns the number of elements in the queue. 88 virtual int32 size() = 0; 89 90 virtual const DataTypeVector& component_dtypes() const = 0; 91 92 string DebugString() override { 93 return strings::StrCat("A Queue of size: ", size()); 94 } 95 96 protected: 97 virtual ~QueueInterface() {} 98 }; 99 100 } // namespace tensorflow 101 102 #endif // TENSORFLOW_FRAMEWORK_QUEUE_INTERFACE_H_ 103
