Tests that the unbatch timeout works.
(self)
| 373 | self.assertEqual(main_results[0], [3]) |
| 374 | |
| 375 | def testUnbatchTimeout(self): |
| 376 | """Tests that the unbatch timeout works.""" |
| 377 | if context.executing_eagerly(): |
| 378 | return |
| 379 | with self.cached_session() as sess: |
| 380 | inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1]) |
| 381 | batched, index, id_t = batch_ops.batch( |
| 382 | [inp], num_batch_threads=1, max_batch_size=2, |
| 383 | batch_timeout_micros=36000000, grad_timeout_micros=0, |
| 384 | batching_queue="") |
| 385 | computation = batched[0] + 1 |
| 386 | timeout_micros = 10 |
| 387 | result = batch_ops.unbatch(computation, index, id_t, timeout_micros, |
| 388 | shared_name="shared_unbatch") |
| 389 | # Set up a parallel pipeline that delays the computation, but uses the |
| 390 | # same unbatch resource object as the non-delayed pipeline. |
| 391 | computation_delayed = script_ops.py_func(delayed_plus1, |
| 392 | [batched[0]], |
| 393 | dtypes.int32) |
| 394 | result_delayed = batch_ops.unbatch(computation_delayed, |
| 395 | index, |
| 396 | id_t, |
| 397 | timeout_micros, |
| 398 | shared_name="shared_unbatch") |
| 399 | |
| 400 | thread_results = [] |
| 401 | def worker(): |
| 402 | # A first call using the non-delayed pipeline. The batcher will send an |
| 403 | # empty tensor along the non-delayed pipeline. |
| 404 | thread_results.extend(sess.run([result], feed_dict={inp: [1]})) |
| 405 | worker_thread = threading.Thread(target=worker) |
| 406 | worker_thread.start() |
| 407 | time.sleep(0.1) # Ensure the thread's call starts first. |
| 408 | # A second call using the delayed pipeline. The batcher will send the |
| 409 | # batched tensor along the delayed pipeline, thus delaying the arrival of |
| 410 | # the batched tensor at the unbatch op, relative to the empty tensor. |
| 411 | # |
| 412 | # TODO(olston, apassos): Avoid relying on the order in which the batch op |
| 413 | # emits the empty tensor versus the batched one. |
| 414 | _ = sess.run([result_delayed], feed_dict={inp: [2]}) |
| 415 | worker_thread.join() |
| 416 | # The thread's call should hit the timeout, and thus get 0 results. |
| 417 | self.assertEqual(len(thread_results), 0) |
| 418 | |
| 419 | |
| 420 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected