xref: /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/batching_util/batch_scheduler.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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