Method ProcessFuncBatch

tensorflow/core/kernels/batch_kernels.cc:480–564

  }

  void ProcessFuncBatch(std::unique_ptr<Batch> batch) const {
    if (batch->empty()) {
      return;
    }

    // We use the 'propagated_context' from one of the threads which setup one
    // of the tasks. This will propagate any common context over all the threads
    // which are running this Session, of which this BatchOp is a part.
    WithContext wc(batch->task(batch->num_tasks() - 1).propagated_context);

    OpKernelContext* last_task_context =
        batch->task(batch->num_tasks() - 1).context;

    // Regardless of the outcome, we need to propagate the status to the
    // individual tasks and signal that they are done. We use MakeCleanup() to
    // ensure that this happens no matter how we exit the method below.
    Status status;
    bool cleanup_done = false;
    auto cleanup_fn = [&cleanup_done, &batch](const Status& status) {
      if (cleanup_done) {
        return;
      }
      for (int i = 0; i < batch->num_tasks(); ++i) {
        batch->mutable_task(i)->context->SetStatus(status);
        batch->mutable_task(i)->done_callback();
      }
      cleanup_done = true;
    };
    auto finally =
        gtl::MakeCleanup([&cleanup_fn, &status] { cleanup_fn(status); });

    status = ValidateBatch(*batch);
    if (!status.ok()) {
      return;
    }

    std::vector<Tensor> concatenated_tensors;
    status =
        ConcatInputTensors(*batch, last_task_context, &concatenated_tensors);
    if (!status.ok()) {
      return;
    }
    FunctionLibraryRuntime::Options opts;
    opts.step_container = last_task_context->step_container();
    opts.cancellation_manager = last_task_context->cancellation_manager();
    opts.stats_collector = last_task_context->stats_collector();
    opts.rendezvous = last_task_context->rendezvous();
    opts.runner = last_task_context->runner();

    auto* flib = last_task_context->function_library();
    std::vector<Tensor> combined_outputs;
    Notification done;
    std::vector<Tensor> args(concatenated_tensors.begin(),
                             concatenated_tensors.end());
    const auto& captured_inputs =
        batch->task(batch->num_tasks() - 1).captured_inputs;
    args.insert(args.end(), captured_inputs.begin(), captured_inputs.end());