dynamic_batching_test.py

# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests dynamic_batching.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
from multiprocessing import pool
import time

import dynamic_batching

import tensorflow as tf

from six.moves import range


_SLEEP_TIME = 1.0


class DynamicBatchingTest(tf.test.TestCase):

  def test_one(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        batch_size = tf.shape(a)[0]
        return a + b, tf.tile([batch_size], [batch_size])

      output = f(tf.constant([[1, 3]]), tf.constant([2]))

      tf.train.start_queue_runners()

      result, batch_size = session.run(output)

      self.assertAllEqual([[3, 5]], result)
      self.assertAllEqual([1], batch_size)

  def test_two(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        batch_size = tf.shape(a)[0]
        return a + b, tf.tile([batch_size], [batch_size])

      output0 = f(tf.constant([1]), tf.constant([2]))
      output1 = f(tf.constant([2]), tf.constant([3]))

      tp = pool.ThreadPool(2)
      f0 = tp.apply_async(session.run, [output0])
      f1 = tp.apply_async(session.run, [output1])

      # Make sure both inputs are in the batcher before starting it.
      time.sleep(_SLEEP_TIME)

      tf.train.start_queue_runners()

      result0, batch_size0 = f0.get()
      result1, batch_size1 = f1.get()

      self.assertAllEqual([3], result0)
      self.assertAllEqual([2], batch_size0)
      self.assertAllEqual([5], result1)
      self.assertAllEqual([2], batch_size1)

  def test_many_small(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        return a + b

      outputs = []
      for i in range(200):
        outputs.append(f(tf.fill([1, 5], i), tf.fill([1, 5], i)))

      tf.train.start_queue_runners()

      tp = pool.ThreadPool(10)
      futures = []
      for output in outputs:
        futures.append(tp.apply_async(session.run, [output]))

      for i, future in enumerate(futures):
        result = future.get()
        self.assertAllEqual([[i * 2] * 5], result)

  def test_input_batch_size_should_be_one(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(a):
        return a

      output = f(tf.constant([1, 2]))

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'requires batch size 1'):
        coord.join()

  def test_run_after_error_should_be_cancelled(self):
    with self.test_session() as session:

      @dynamic_batching.batch_fn
      def f(a):
        return a

      output = f(tf.constant([1, 2]))

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

  def test_input_shapes_should_be_equal(self):
    with self.test_session() as session:

      @dynamic_batching.batch_fn
      def f(a, b):
        return a + b

      output0 = f(tf.constant([1]), tf.constant([2]))
      output1 = f(tf.constant([[2]]), tf.constant([3]))

      tp = pool.ThreadPool(2)
      f0 = tp.apply_async(session.run, [output0])
      f1 = tp.apply_async(session.run, [output1])

      time.sleep(_SLEEP_TIME)

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        f0.get()
        f1.get()

      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'Shapes of inputs much be equal'):
        coord.join()

  def test_output_must_have_batch_dimension(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(_):
        return tf.constant(1)

      output = f(tf.constant([1]))

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'Output shape must have a batch dimension'):
        coord.join()

  def test_output_must_have_same_batch_dimension_size_as_input(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(_):
        return tf.constant([1, 2, 3, 4])

      output = f(tf.constant([1]))

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          'Output shape must have the same batch dimension as the input batch '
          'size. Expected: 1 Observed: 4'):
        coord.join()

  def test_get_inputs_cancelled(self):
    with tf.Graph().as_default():

      @dynamic_batching.batch_fn
      def f(a):
        return a

      f(tf.constant([1]))

      # Intentionally using tf.Session() instead of self.test_session() to have
      # control over closing the session. test_session() is a cached session.
      with tf.Session():
        coord = tf.train.Coordinator()
        tf.train.start_queue_runners(coord=coord)
        # Sleep to make sure the queue runner has started the first run call.
        time.sleep(_SLEEP_TIME)

      # Session closed.
      coord.request_stop()
      coord.join()

  def test_batcher_closed(self):
    with tf.Graph().as_default():
      @dynamic_batching.batch_fn
      def f(a):
        return a

      f(tf.constant([1]))

      # Intentionally using tf.Session() instead of self.test_session() to have
      # control over closing the session. test_session() is a cached session.
      with tf.Session():
        coord = tf.train.Coordinator()
        tf.train.start_queue_runners(coord=coord)
        time.sleep(_SLEEP_TIME)
        coord.request_stop()  # Calls close operation.
        coord.join()
      # Session closed.

  def test_minimum_batch_size(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn_with_options(
          minimum_batch_size=2, timeout_ms=1000)
      def f(a, b):
        batch_size = tf.shape(a)[0]
        return a + b, tf.tile([batch_size], [batch_size])

      output = f(tf.constant([[1, 3]]), tf.constant([2]))

      tf.train.start_queue_runners()

      start = datetime.datetime.now()
      session.run(output)
      duration = datetime.datetime.now() - start

      # There should have been a timeout here because only one sample was added
      # and the minimum batch size is 2.
      self.assertLessEqual(.9, duration.total_seconds())
      self.assertGreaterEqual(1.5, duration.total_seconds())

      outputs = [
          f(tf.constant([[1, 3]]), tf.constant([2])),
          f(tf.constant([[1, 3]]), tf.constant([2]))
      ]

      start = datetime.datetime.now()
      (_, batch_size), _ = session.run(outputs)
      duration = datetime.datetime.now() - start

      # The outputs should be executed immediately because two samples are
      # added.
      self.assertGreaterEqual(.5, duration.total_seconds())
      self.assertEqual(2, batch_size)

  def test_maximum_batch_size(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn_with_options(maximum_batch_size=2)
      def f(a, b):
        batch_size = tf.shape(a)[0]
        return a + b, tf.tile([batch_size], [batch_size])

      outputs = [
          f(tf.constant([1]), tf.constant([2])),
          f(tf.constant([1]), tf.constant([2])),
          f(tf.constant([1]), tf.constant([2])),
          f(tf.constant([1]), tf.constant([2])),
          f(tf.constant([1]), tf.constant([2])),
      ]

      tf.train.start_queue_runners()

      results = session.run(outputs)

      for value, batch_size in results:
        self.assertEqual(3, value)
        self.assertGreaterEqual(2, batch_size)

  def test_static_shape(self):
    assertions_triggered = [0]

    @dynamic_batching.batch_fn_with_options(minimum_batch_size=1,
                                            maximum_batch_size=2)
    def f0(a):
      self.assertEqual(None, a.shape[0].value)
      assertions_triggered[0] += 1
      return a

    @dynamic_batching.batch_fn_with_options(minimum_batch_size=2,
                                            maximum_batch_size=2)
    def f1(a):
      # Even though minimum_batch_size and maximum_batch_size are equal, the
      # timeout can cause a batch with less than mininum_batch_size.
      self.assertEqual(None, a.shape[0].value)
      assertions_triggered[0] += 1
      return a

    @dynamic_batching.batch_fn_with_options(minimum_batch_size=2,
                                            maximum_batch_size=2,
                                            timeout_ms=None)
    def f2(a):
      # When timeout is disabled and minimum/maximum batch size are equal, the
      # shape is statically known.
      self.assertEqual(2, a.shape[0].value)
      assertions_triggered[0] += 1
      return a

    f0(tf.constant([1]))
    f1(tf.constant([1]))
    f2(tf.constant([1]))
    self.assertEqual(3, assertions_triggered[0])

  def test_out_of_order_execution1(self):
    with self.test_session() as session:
      batcher = dynamic_batching._Batcher(minimum_batch_size=1,
                                          maximum_batch_size=1,
                                          timeout_ms=None)

      tp = pool.ThreadPool(10)
      r0 = tp.apply_async(session.run, batcher.compute([[1]], [tf.int32]))
      (input0,), computation_id0 = session.run(batcher.get_inputs([tf.int32]))
      r1 = tp.apply_async(session.run, batcher.compute([[2]], [tf.int32]))
      (input1,), computation_id1 = session.run(batcher.get_inputs([tf.int32]))

      self.assertAllEqual([1], input0)
      self.assertAllEqual([2], input1)

      session.run(batcher.set_outputs([input0 + 42], computation_id0))
      session.run(batcher.set_outputs([input1 + 42], computation_id1))

      self.assertAllEqual([43], r0.get())
      self.assertAllEqual([44], r1.get())

  def test_out_of_order_execution2(self):
    with self.test_session() as session:
      batcher = dynamic_batching._Batcher(minimum_batch_size=1,
                                          maximum_batch_size=1,
                                          timeout_ms=None)

      tp = pool.ThreadPool(10)
      r0 = tp.apply_async(session.run, batcher.compute([[1]], [tf.int32]))
      (input0,), computation_id0 = session.run(batcher.get_inputs([tf.int32]))
      r1 = tp.apply_async(session.run, batcher.compute([[2]], [tf.int32]))
      (input1,), computation_id1 = session.run(batcher.get_inputs([tf.int32]))

      self.assertAllEqual([1], input0)
      self.assertAllEqual([2], input1)

      # These two runs are switched from testOutOfOrderExecution1.
      session.run(batcher.set_outputs([input1 + 42], computation_id1))
      session.run(batcher.set_outputs([input0 + 42], computation_id0))

      self.assertAllEqual([43], r0.get())
      self.assertAllEqual([44], r1.get())

  def test_invalid_computation_id(self):
    with self.test_session() as session:
      batcher = dynamic_batching._Batcher(minimum_batch_size=1,
                                          maximum_batch_size=1,
                                          timeout_ms=None)

      tp = pool.ThreadPool(10)
      tp.apply_async(session.run, batcher.compute([[1]], [tf.int32]))
      (input0,), _ = session.run(batcher.get_inputs([tf.int32]))

      self.assertAllEqual([1], input0)

      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'Invalid computation id'):
        session.run(batcher.set_outputs([input0], 42))

  def test_op_shape(self):
    with self.test_session():
      batcher = dynamic_batching._Batcher(minimum_batch_size=1,
                                          maximum_batch_size=1,
                                          timeout_ms=None)

      _, computation_id = batcher.get_inputs([tf.int32])

      self.assertEqual([], computation_id.shape)


class DynamicBatchingBenchmarks(tf.test.Benchmark):

  def benchmark_batching_small(self):
    with tf.Session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        return a + b

      outputs = []
      for _ in range(1000):
        outputs.append(f(tf.ones([1, 10]), tf.ones([1, 10])))
      op_to_benchmark = tf.group(*outputs)

      tf.train.start_queue_runners()

      self.run_op_benchmark(
          name='batching_many_small',
          sess=session,
          op_or_tensor=op_to_benchmark,
          burn_iters=10,
          min_iters=50)

  def benchmark_batching_large(self):
    with tf.Session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        return a + b

      outputs = []
      for _ in range(1000):
        outputs.append(f(tf.ones([1, 100000]), tf.ones([1, 100000])))
      op_to_benchmark = tf.group(*outputs)

      tf.train.start_queue_runners()

      self.run_op_benchmark(
          name='batching_many_large',
          sess=session,
          op_or_tensor=op_to_benchmark,
          burn_iters=10,
          min_iters=50)


if __name__ == '__main__':
  tf.test.main()