Implementation:Tensorflow Serving Streaming Batch Scheduler Test
| Knowledge Sources | |
|---|---|
| Domains | Testing, Batching |
| Last Updated | 2026-02-13 00:00 GMT |
Overview
Test suite validating the StreamingBatchScheduler which groups incoming tasks into batches with configurable size constraints and timeout behavior.
Description
This test file exercises the StreamingBatchScheduler class, a batch scheduler that forms batches from a stream of incoming tasks. The tests use a FakeTask helper class (implementing BatchTask) with configurable sizes to control batching behavior. Tests cover basic batch formation, batch size constraints, timeout-based batch closing (both with fake and real clocks), final underfull batch processing on scheduler destruction, callback timing, and const method behavior.
Usage
Run these tests to verify that the streaming batch scheduler correctly forms batches respecting size limits, triggers batch processing on timeout, and properly handles cleanup on destruction. Critical after changes to scheduler logic or timeout mechanisms.
Code Reference
Source Location
- Repository: Tensorflow_Serving
- File:
tensorflow_serving/batching/streaming_batch_scheduler_test.cc - Lines: 1-321
Test Fixture
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
~FakeTask() override = default;
size_t size() const override { return size_; }
private:
const size_t size_;
TF_DISALLOW_COPY_AND_ASSIGN(FakeTask);
};
// Creates a FakeTask of size 'task_size', and calls 'scheduler->Schedule()'
// on that task. Returns the resulting status.
absl::Status ScheduleTask(size_t task_size,
BatchScheduler<FakeTask>* scheduler) {
std::unique_ptr<FakeTask> task(new FakeTask(task_size));
absl::Status status = scheduler->Schedule(&task);
CHECK_EQ(status.ok(), task == nullptr);
return status;
}
Build Target
bazel test //tensorflow_serving/batching:streaming_batch_scheduler_test
Test Coverage
Key Test Cases
| Test Name | Category | Description |
|---|---|---|
Basic |
Core | Verifies two tasks are grouped into a single batch with correct sizes |
ObeyBatchSizeConstraint |
Size Constraint | Tests that tasks exceeding max_batch_size are split into separate batches |
Timeout |
Timeout | Validates batch processing triggers on timeout using a FakeClockEnv |
RealClockTimeout |
Timeout | Tests timeout behavior with the real system clock |
FinalUnderfullBatchProcessedUponDeletion |
Lifecycle | Ensures incomplete batches are flushed when the scheduler is destroyed |
BatchHandedToCallbackWhenFirstCreated |
Callback | Verifies batch is handed to callback immediately upon first task arrival |
ConstMethods |
API | Tests const methods (NumEnqueuedTasks, SchedulingCapacity) return correct values |
Usage Examples
Test Pattern
TEST(StreamingBatchSchedulerTest, Basic) {
bool callback_called = false;
auto callback = [&callback_called](std::unique_ptr<Batch<FakeTask>> batch) {
callback_called = true;
batch->WaitUntilClosed();
ASSERT_EQ(2, batch->num_tasks());
EXPECT_EQ(3, batch->task(0).size());
EXPECT_EQ(5, batch->task(1).size());
};
{
StreamingBatchScheduler<FakeTask>::Options options;
options.max_batch_size = 10;
options.batch_timeout_micros = 100 * 1000; // 100 milliseconds
options.num_batch_threads = 1;
std::unique_ptr<StreamingBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(StreamingBatchScheduler<FakeTask>::Create(options, callback,
&scheduler));
TF_ASSERT_OK(ScheduleTask(3, scheduler.get()));
TF_ASSERT_OK(ScheduleTask(5, scheduler.get()));
}
EXPECT_TRUE(callback_called);
}