Implementation:Tensorflow Serving Batching Session Test

Knowledge Sources	Tensorflow_Serving
Domains	Testing, Batching
Last Updated	2026-02-13 00:00 GMT

Overview

Test suite validating the BatchingSession wrapper that adds transparent batching to TensorFlow Session inference calls.

Description

This test file exercises the BatchingSession class, which wraps a standard TensorFlow Session to batch multiple concurrent Run() requests together for improved throughput. The tests use a BatchSizeCapturingSession helper fixture that wraps a real session (loaded with the half-plus-two SavedModel) and records the actual batch sizes submitted to the underlying session. Tests validate correct batching behavior, padding logic, allowed batch sizes, cost propagation, multiple signatures, timeout handling, and thread pool options.

Usage

Run these tests to verify that batching session correctly groups individual inference requests into batches, pads tensors when needed, respects allowed batch size constraints, and correctly splits output tensors back to individual callers. Essential after any changes to the batching session layer or scheduler integration.

Code Reference

Source Location

• Repository: Tensorflow_Serving
• File: tensorflow_serving/batching/batching_session_test.cc
• Lines: 1-1108

Test Fixture

// A wrapper around a Session that captures the batch size.
class BatchSizeCapturingSession : public ServingSession {
 public:
  explicit BatchSizeCapturingSession(std::unique_ptr<Session> wrapped)
      : wrapped_(std::move(wrapped)) {}
  ~BatchSizeCapturingSession() override = default;

  absl::Status Run(const RunOptions& run_options,
                   const std::vector<std::pair<string, Tensor>>& inputs,
                   const std::vector<string>& output_tensor_names,
                   const std::vector<string>& target_node_names,
                   std::vector<Tensor>* outputs, RunMetadata* run_metadata,
                   const thread::ThreadPoolOptions& thread_pool_options)
      override TF_LOCKS_EXCLUDED(latest_batch_size_mu_) {
    {
      mutex_lock l(latest_batch_size_mu_);
      latest_batch_size_ = inputs[0].second.shape().dim_size(0);
    }
    // ...
  }

  int latest_batch_size() const;
  CostGraphDef* mutable_cost_graph();

 private:
  std::unique_ptr<Session> wrapped_;
  mutable mutex latest_batch_size_mu_;
  int latest_batch_size_ TF_GUARDED_BY(latest_batch_size_mu_) = -1;
  CostGraphDef cost_graph_;
};

Build Target

bazel test //tensorflow_serving/batching:batching_session_test

Test Coverage

Key Test Cases

Test Name	Category	Description
TensorSignatureFromSignatureDef	Signature	Validates extraction of tensor signatures from a single SignatureDef
TensorSignatureFromSignatureDefs	Signature	Validates extraction of tensor signatures from multiple SignatureDefs
Basic	Core	Verifies basic batching of two concurrent requests into one batch
BatchingWithPadding	Padding	Tests batching with variable-length tensor padding enabled
BatchingWithLargeBatch	Core	Tests batching behavior with large input batches
BatchHandlesSplitError	Error Handling	Verifies correct error propagation when batch splitting fails
BatchingLazySplit	Core	Tests lazy splitting mode for batch processing
BatchingWithPaddingAndCost	Cost	Validates cost graph propagation with padded batches
BatchingWithCost	Cost	Verifies cost graph metadata is correctly returned per request
UnequalTensorShapesWithPaddingTurnedOff	Validation	Ensures error on unequal tensor shapes when padding is disabled
SingletonBatch	Core	Tests batching with a single request (batch size 1)
RequestThatDoesntMatchSignatureGetsRunAnyway	Fallback	Verifies non-batched requests bypass the batching path
RequestWithIncompatibleInputTensorSizes	Validation	Tests rejection of requests with mismatched input tensor dimensions
AllowedBatchSizesNoPaddingNeeded	Configuration	Tests allowed batch sizes when no padding is needed
AllowedBatchSizesRequirePadding	Configuration	Tests allowed batch sizes when padding up to the next allowed size
UnsortedAllowedBatchSizesRejected	Validation	Verifies rejection of unsorted allowed batch size configurations
DifferentOrderForInputAndOutputTensors	Core	Tests that input/output tensor ordering is preserved correctly
MultipleSignatures	Multi-Signature	Validates batching across multiple model signatures
EnqueuedLongerThanTimeout	Timeout	Tests behavior when batch timeout expires before batch is full
ThreadPoolOptions	Threading	Validates custom thread pool options are passed through
SubsetOutputTensors	Core	Tests requesting only a subset of output tensors

Usage Examples

Test Pattern

// Helper to create a session loaded with the half-plus-two SavedModel
std::unique_ptr<Session> CreateHalfPlusTwoSession() {
  tensorflow::SessionOptions session_options;
  tensorflow::RunOptions run_options;
  const string export_dir = test_util::TensorflowTestSrcDirPath(
      "cc/saved_model/testdata/half_plus_two/00000123");
  SavedModelBundle bundle;
  TF_CHECK_OK(LoadSavedModel(session_options, run_options, export_dir,
                             {kSavedModelTagServe}, &bundle));
  return std::move(bundle.session);
}

// Test validates basic batching behavior
void TestRequest(const std::vector<float>& x_values, TensorShape x_shape,
                 const std::vector<float>& y_values, TensorShape y_shape,
                 Session* session,
                 test_util::CountingThreadPool* inter_op_threadpool = nullptr,
                 test_util::CountingThreadPool* intra_op_threadpool = nullptr) {
  Tensor input = test::AsTensor<float>(x_values, x_shape);
  Tensor expected_output = test::AsTensor<float>(y_values, y_shape);
  // ...
}

Related Pages

• Implementation:Tensorflow_Serving_Streaming_Batch_Scheduler_Test
• Implementation:Tensorflow_Serving_Tfrt_Saved_Model_With_Batching_Test