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 // Abstractions for processing small tasks in a batched fashion, to reduce
17 // processing times and costs that can be amortized across multiple tasks.
18 //
19 // The core class is BatchScheduler, which groups tasks into batches.
20 //
21 // BatchScheduler encapsulates logic for aggregating multiple tasks into a
22 // batch, and kicking off processing of a batch on a thread pool it manages.
23 //
24 // This file defines an abstract BatchScheduler class.
25
26 #ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
27 #define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
28
29 #include <stddef.h>
30
31 #include <algorithm>
32 #include <functional>
33 #include <memory>
34 #include <utility>
35 #include <vector>
36
37 #include "tensorflow/core/lib/core/notification.h"
38 #include "tensorflow/core/lib/core/status.h"
39 #include "tensorflow/core/platform/logging.h"
40 #include "tensorflow/core/platform/macros.h"
41 #include "tensorflow/core/platform/mutex.h"
42 #include "tensorflow/core/platform/thread_annotations.h"
43 #include "tensorflow/core/platform/types.h"
44 #include "tensorflow/core/profiler/lib/traceme.h"
45
46 namespace tensorflow {
47 namespace serving {
48
49 // The abstract superclass for a unit of work to be done as part of a batch.
50 //
51 // An implementing subclass typically contains (or points to):
52 // (a) input data;
53 // (b) a thread-safe completion signal (e.g. a Notification);
54 // (c) a place to store the outcome (success, or some error), upon completion;
55 // (d) a place to store the output data, upon success.
56 //
57 // Items (b), (c) and (d) are typically non-owned pointers to data homed
58 // elsewhere, because a task's ownership gets transferred to a BatchScheduler
59 // (see below) and it may be deleted as soon as it is done executing.
60 class BatchTask {
61 public:
62 virtual ~BatchTask() = default;
63
64 // Returns the size of the task, in terms of how much it contributes to the
65 // size of a batch. (A batch's size is the sum of its task sizes.)
66 virtual size_t size() const = 0;
67 };
68
69 // A thread-safe collection of BatchTasks, to be executed together in some
70 // fashion.
71 //
72 // At a given time, a batch is either "open" or "closed": an open batch can
73 // accept new tasks; a closed one cannot. A batch is monotonic: initially it is
74 // open and tasks can be added to it; then it is closed and its set of tasks
75 // remains fixed for the remainder of its life. A closed batch cannot be re-
76 // opened. Tasks can never be removed from a batch.
77 //
78 // Type parameter TaskType must be a subclass of BatchTask.
79 template <typename TaskType>
80 class Batch {
81 public:
82 Batch();
83 explicit Batch(uint64 traceme_context_id);
84 virtual ~Batch(); // Blocks until the batch is closed.
85
86 // Appends 'task' to the batch. After calling AddTask(), the newly-added task
87 // can be accessed via task(num_tasks()-1) or mutable_task(num_tasks()-1).
88 // Dies if the batch is closed.
89 void AddTask(std::unique_ptr<TaskType> task);
90
91 // Removes the most recently added task. Returns nullptr if the batch is
92 // empty.
93 std::unique_ptr<TaskType> RemoveTask();
94
95 // Caller takes ownership of returned tasks.
96 // Must be called after a batch is closed.
97 std::vector<std::unique_ptr<TaskType>> RemoveAllTasks();
98
99 // Returns the number of tasks in the batch.
100 int num_tasks() const;
101
102 // Returns true iff the batch contains 0 tasks.
103 bool empty() const;
104
105 // Returns a reference to the ith task (in terms of insertion order).
106 const TaskType& task(int i) const;
107
108 // Returns a pointer to the ith task (in terms of insertion order).
109 //
110 // Caller doesn't take ownership.
111 TaskType* mutable_task(int i);
112
113 // Returns the sum of the task sizes.
114 size_t size() const;
115
116 // Returns true iff the batch is currently closed.
117 bool IsClosed() const;
118
119 // Blocks until the batch is closed.
120 void WaitUntilClosed() const;
121
122 // Marks the batch as closed. Dies if called more than once.
123 void Close();
124
125 // Returns the TraceMe context id of this batch.
126 uint64 traceme_context_id() const;
127
128 private:
129 mutable mutex mu_;
130
131 // The tasks in the batch.
132 std::vector<std::unique_ptr<TaskType>> tasks_ TF_GUARDED_BY(mu_);
133
134 // The sum of the sizes of the tasks in 'tasks_'.
135 size_t size_ TF_GUARDED_BY(mu_) = 0;
136
TF_GUARDED_BY(mu_)137 std::atomic<bool> empty_ TF_GUARDED_BY(mu_){true};
138
139 // Whether the batch has been closed.
140 Notification closed_;
141
142 // The TracMe context id.
143 const uint64 traceme_context_id_;
144
145 TF_DISALLOW_COPY_AND_ASSIGN(Batch);
146 };
147
148 // An abstract batch scheduler class. Collects individual tasks into batches,
149 // and processes each batch on a pool of "batch threads" that it manages. The
150 // actual logic for processing a batch is accomplished via a callback.
151 //
152 // Type parameter TaskType must be a subclass of BatchTask.
153 template <typename TaskType>
154 class BatchScheduler {
155 public:
156 virtual ~BatchScheduler() = default;
157
158 // Submits a task to be processed as part of a batch.
159 //
160 // Ownership of '*task' is transferred to the callee iff the method returns
161 // Status::OK. In that case, '*task' is left as nullptr. Otherwise, '*task' is
162 // left as-is.
163 //
164 // If no batch processing capacity is available to process this task at the
165 // present time, and any task queue maintained by the implementing subclass is
166 // full, this method returns an UNAVAILABLE error code. The client may retry
167 // later.
168 //
169 // Other problems, such as the task size being larger than the maximum batch
170 // size, yield other, permanent error types.
171 //
172 // In all cases, this method returns "quickly" without blocking for any
173 // substantial amount of time. If the method returns Status::OK, the task is
174 // processed asynchronously, and any errors that occur during the processing
175 // of the batch that includes the task can be reported to 'task'.
176 virtual Status Schedule(std::unique_ptr<TaskType>* task) = 0;
177
178 // Returns the number of tasks that have been scheduled (i.e. accepted by
179 // Schedule()), but have yet to be handed to a thread for execution as part of
180 // a batch. Note that this returns the number of tasks, not the aggregate task
181 // size (so if there is one task of size 3 and one task of size 5, this method
182 // returns 2 rather than 8).
183 virtual size_t NumEnqueuedTasks() const = 0;
184
185 // Returns a guaranteed number of size 1 tasks that can be Schedule()d without
186 // getting an UNAVAILABLE error. In a typical implementation, returns the
187 // available space on a queue.
188 //
189 // There are two important caveats:
190 // 1. The guarantee does not extend to varying-size tasks due to possible
191 // internal fragmentation of batches.
192 // 2. The guarantee only holds in a single-thread environment or critical
193 // section, i.e. if an intervening thread cannot call Schedule().
194 //
195 // This method is useful for monitoring, or for guaranteeing a future slot in
196 // the schedule (but being mindful about the caveats listed above).
197 virtual size_t SchedulingCapacity() const = 0;
198
199 // Returns the maximum allowed size of tasks submitted to the scheduler. (This
200 // is typically equal to a configured maximum batch size.)
201 virtual size_t max_task_size() const = 0;
202 };
203
204 //////////
205 // Implementation details follow. API users need not read.
206
207 template <typename TaskType>
Batch()208 Batch<TaskType>::Batch() : Batch(0) {}
209
210 template <typename TaskType>
Batch(uint64 traceme_context_id)211 Batch<TaskType>::Batch(uint64 traceme_context_id)
212 : traceme_context_id_(traceme_context_id) {}
213
214 template <typename TaskType>
~Batch()215 Batch<TaskType>::~Batch() {
216 WaitUntilClosed();
217 }
218
219 template <typename TaskType>
AddTask(std::unique_ptr<TaskType> task)220 void Batch<TaskType>::AddTask(std::unique_ptr<TaskType> task) {
221 DCHECK(!IsClosed());
222 {
223 mutex_lock l(mu_);
224 size_ += task->size();
225 tasks_.push_back(std::move(task));
226 empty_.store(false);
227 }
228 }
229
230 template <typename TaskType>
RemoveAllTasks()231 std::vector<std::unique_ptr<TaskType>> Batch<TaskType>::RemoveAllTasks() {
232 DCHECK(IsClosed());
233 {
234 mutex_lock l(mu_);
235 size_ = 0;
236 empty_.store(true);
237 std::vector<std::unique_ptr<TaskType>> tasks_to_return;
238
239 // Swapping vector takes constant time.
240 tasks_to_return.swap(tasks_);
241 return std::move(tasks_to_return);
242 }
243 }
244
245 template <typename TaskType>
RemoveTask()246 std::unique_ptr<TaskType> Batch<TaskType>::RemoveTask() {
247 {
248 mutex_lock l(mu_);
249 if (tasks_.empty()) {
250 return nullptr;
251 }
252 std::unique_ptr<TaskType> task = std::move(tasks_.back());
253 size_ -= task->size();
254 tasks_.pop_back();
255 if (tasks_.empty()) {
256 empty_.store(true);
257 }
258 return task;
259 }
260 }
261
262 template <typename TaskType>
num_tasks()263 int Batch<TaskType>::num_tasks() const {
264 {
265 mutex_lock l(mu_);
266 return tasks_.size();
267 }
268 }
269
270 template <typename TaskType>
empty()271 bool Batch<TaskType>::empty() const TF_NO_THREAD_SAFETY_ANALYSIS {
272 // tracer is added to zoom in about this method.
273 // TODO(b/160249203): Remove tracer after evaluating a change to reduce
274 // lock contention and cpu usage (which is observed in profiler and
275 // very data-driven).
276 tensorflow::profiler::TraceMe tracer("BatchTask::empty");
277 return empty_.load();
278 }
279
280 template <typename TaskType>
task(int i)281 const TaskType& Batch<TaskType>::task(int i) const {
282 DCHECK_GE(i, 0);
283 {
284 mutex_lock l(mu_);
285 DCHECK_LT(i, tasks_.size());
286 return *tasks_[i].get();
287 }
288 }
289
290 template <typename TaskType>
mutable_task(int i)291 TaskType* Batch<TaskType>::mutable_task(int i) {
292 DCHECK_GE(i, 0);
293 {
294 mutex_lock l(mu_);
295 DCHECK_LT(i, tasks_.size());
296 return tasks_[i].get();
297 }
298 }
299
300 template <typename TaskType>
size()301 size_t Batch<TaskType>::size() const {
302 {
303 mutex_lock l(mu_);
304 return size_;
305 }
306 }
307
308 template <typename TaskType>
IsClosed()309 bool Batch<TaskType>::IsClosed() const {
310 return const_cast<Notification*>(&closed_)->HasBeenNotified();
311 }
312
313 template <typename TaskType>
WaitUntilClosed()314 void Batch<TaskType>::WaitUntilClosed() const {
315 const_cast<Notification*>(&closed_)->WaitForNotification();
316 }
317
318 template <typename TaskType>
Close()319 void Batch<TaskType>::Close() {
320 closed_.Notify();
321 }
322
323 template <typename TaskType>
traceme_context_id()324 uint64 Batch<TaskType>::traceme_context_id() const {
325 return traceme_context_id_;
326 }
327
328 } // namespace serving
329 } // namespace tensorflow
330
331 #endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
332