Implementation:Tensorflow Serving Tfrt Saved Model With Batching Test
| Knowledge Sources | |
|---|---|
| Domains | Testing, Batching |
| Last Updated | 2026-02-13 00:00 GMT |
Overview
Test suite validating the SavedModelWithBatching wrapper that adds transparent batching to TFRT SavedModel inference calls.
Description
This test file exercises the SavedModelWithBatching class, which wraps a TFRT SavedModel to batch concurrent function invocations. The tests use a SavedModelWithBatchingTest fixture class with a MockSavedModel to control and verify batching behavior. Custom matchers (MatchesTensor, TFStatusIs) are used for assertion clarity. Tests cover null model validation, duplicate scheduler detection, unknown function handling, batching with and without padding, deadline expiry, multiple function support, and input splitting.
Usage
Run these tests to verify that TFRT SavedModel batching correctly groups concurrent function calls, pads variable-length inputs, handles errors from partial tasks, and routes different functions to separate batch schedulers. Essential when modifying TFRT batching integration.
Code Reference
Source Location
- Repository: Tensorflow_Serving
- File:
tensorflow_serving/batching/tfrt_saved_model_with_batching_test.cc - Lines: 1-699
Test Fixture
class SavedModelWithBatchingTest : public ::testing::Test {
protected:
SavedModelWithBatchingTest() = default;
std::unique_ptr<test_util::MockSavedModel> InitializeMockSavedModel() {
auto wrapped_saved_model = absl::make_unique<test_util::MockSavedModel>();
wrapped_saved_model_ = wrapped_saved_model.get();
ON_CALL(*wrapped_saved_model_, GetFunctionMetadata(_))
.WillByDefault(Return(tfrt::FunctionMetadata(&signature)));
return wrapped_saved_model;
}
void Initialize(
const BasicBatchScheduler<SavedModelBatchingTask>::Options
&scheduler_options,
const SavedModelBatchingOptions &options) {
// Creates mock model, binds scheduler creators for func1 and func2,
// and creates the SavedModelWithBatching wrapper
}
Tensor MakeTensor(const std::vector<float> &tensor_vec,
const TensorShape &shape);
std::vector<Tensor> MakeTensors(...);
std::vector<std::vector<Tensor>> MakeTensorsBatch(...);
std::unique_ptr<tfrt::SavedModel> saved_model_with_batching_;
test_util::MockSavedModel *wrapped_saved_model_;
};
Build Target
bazel test //tensorflow_serving/batching:tfrt_saved_model_with_batching_test
Test Coverage
Key Test Cases
| Test Name | Category | Description |
|---|---|---|
NullWrappedSavedModel |
Validation | Verifies error when wrapping a null SavedModel |
MultipleBatchSchedulersForOneFunction |
Validation | Ensures error when duplicate schedulers are registered for one function |
FailedCreatingBatchScheduler |
Error Handling | Tests graceful handling of batch scheduler creation failure |
FunctionNameNotFound |
Validation | Verifies error when calling an unregistered function name |
BatchingWithoutPadding |
Core | Tests batching of multiple requests without tensor padding |
BatchingWithPadding |
Padding | Tests batching with variable-length input padding enabled |
UnequalShapesWhenPaddingIsTurnedOff |
Validation | Ensures error on unequal tensor shapes without padding |
AllTasksExceededDeadline |
Timeout | Tests behavior when all tasks in a batch exceed their deadline |
MultipleFunctions |
Multi-Function | Validates independent batching across different function names |
SplitInputBasic |
Splitting | Tests input splitting across multiple sub-batches |
PartialTaskFails |
Error Handling | Verifies correct error propagation when a partial task fails |
Usage Examples
Test Pattern
// Helper to build scheduler options with a specific max batch size
BasicBatchScheduler<SavedModelBatchingTask>::Options BuildSchedulerOptions(
int max_batch_size) {
BasicBatchScheduler<SavedModelBatchingTask>::Options options;
options.max_batch_size = max_batch_size;
options.batch_timeout_micros = 1000 * 1000 * 1000; // 1000s.
options.num_batch_threads = 1;
return options;
}
// Custom tensor matcher for test assertions
MATCHER_P(MatchesTensor, p, "") {
const Tensor &x = arg;
const Tensor &y = *p;
const float *Tx = x.unaligned_flat<float>().data();
const float *Ty = y.unaligned_flat<float>().data();
auto size = x.NumElements();
for (decltype(size) i = 0; i < size; ++i) {
if (Tx[i] != Ty[i]) return false;
}
return true;
}