message: In BatchingSession, optionally re-use a shared batch scheduler with other batching session instances.
diff: @@ -141,6 +141,39 @@ class BasicBatchScheduler : public BatchScheduler<TaskType> {
  // (Keep them mirrored to the ones in SharedBatchScheduler::QueueOptions and
  // SharedBatchScheduler::Options.)
  struct Options {
  + // Options related with (underlying) shared batch scheduler.
  + // 'thread_pool_name' and 'num_batch_threads' are used to initialize
  + // a shared batch scheduler underlyingly iff 'shared_batch_scheduler' is
  + // nullptr.
  + //
  + // There are two ways to specify threading:
  + // 1) Have each session create its own pool.
  + // 2) Have multiple sessions share the same pool.
  + //
  + // In general, the number of threads should be tied to roughly the number of
  + // compute resources (CPU cores or accelerator cores) backing the threads.
  + // Sharing a thread pool helps alleviate potential over allocation of
  + // threads to limited compute resources.
  +
  + // To have each session create its own thread pool (1) set
  + // thread_pool_name/num_batch_threads.
  +
  + // To share a thread pool (2) create a scheduler and pass it in.
  +
  + // The name to use for the pool of batch threads.
  + string thread_pool_name = {"batch_threads"};
  +
  + // The number of threads to use to process batches.
  + // Must be >= 1, and should be tuned carefully.
  + int num_batch_threads = port::MaxParallelism();
  +
  + // If specified, this scheduler will be used underlyingly to schedule
  + // batches. Note setting this means `thread_pool_name` and
  + // `num_batch_threads` are ignored.
  + std::shared_ptr<SharedBatchScheduler<TaskType>> shared_batch_scheduler =
  + nullptr;
  +
  + // Options for queue.
  // The maximum size of each batch.
  //
  // The scheduler may form batches of any size between 1 and this number
  @@ -163,12 +196,6 @@ class BasicBatchScheduler : public BatchScheduler<TaskType> {
  // avoid latency spikes.
  int64 batch_timeout_micros = 0;
  - // The name to use for the pool of batch threads.
  - string thread_pool_name = {"batch_threads"};
  -
  - // The number of threads to use to process batches.
  - // Must be >= 1, and should be tuned carefully.
  - int num_batch_threads = port::MaxParallelism();
  // The maximum allowable number of enqueued (accepted by Schedule() but
  // not yet being processed on a batch thread) tasks in terms of batches.
  @@ -272,13 +299,19 @@ Status BasicBatchScheduler<TaskType>::Create(
  std::function<void(std::unique_ptr<Batch<TaskType>>)>
  process_batch_callback,
  std::unique_ptr<BasicBatchScheduler>* scheduler) {
  + std::shared_ptr<SharedBatchScheduler<TaskType>> shared_scheduler;
  +
  + if (options.shared_batch_scheduler == nullptr) {
  typename SharedBatchScheduler<TaskType>::Options shared_scheduler_options;
  shared_scheduler_options.thread_pool_name = options.thread_pool_name;
  shared_scheduler_options.num_batch_threads = options.num_batch_threads;
  shared_scheduler_options.env = options.env;
  - std::shared_ptr<SharedBatchScheduler<TaskType>> shared_scheduler;
  +
  TF_RETURN_IF_ERROR(SharedBatchScheduler<TaskType>::Create(
  shared_scheduler_options, &shared_scheduler));
  + } else {
  + shared_scheduler = options.shared_batch_scheduler;
  + }
  typename SharedBatchScheduler<TaskType>::QueueOptions
  shared_scheduler_queue_options;
  