Method RunFunction

tensorflow/core/kernels/partitioned_function_ops.cc:226–282

}

void PartitionedCallOp::RunFunction(FunctionLibraryRuntime::Handle handle,
                                    const std::vector<Tensor>& inputs,
                                    FunctionLibraryRuntime* lib,
                                    OpKernelContext* ctx, DoneCallback done) {
  FunctionLibraryRuntime::Options run_opts;
  ResourceMgr* resource_mgr = lib->device()->resource_manager();
  ScopedStepContainer* step_container = new ScopedStepContainer(
      run_opts.step_id, [resource_mgr](const string& name) {
        resource_mgr->Cleanup(name).IgnoreError();
      });
  run_opts.step_container = step_container;
  run_opts.cancellation_manager = ctx->cancellation_manager();
  run_opts.stats_collector = ctx->stats_collector();
  run_opts.collective_executor = ctx->collective_executor();
  // TODO(akshayka): Consider selecting a runner on a per-device basis,
  // i.e., using device-specific threadpools when available.
  run_opts.runner = ctx->runner();
  run_opts.source_device =
      lib->device() == nullptr ? "" : lib->device()->name();
  run_opts.allow_dead_tensors = true;

  Rendezvous* rendez;
  OP_REQUIRES_OK_ASYNC(
      ctx,
      ctx->create_rendezvous(run_opts.step_id,
                             ctx->function_library()->device_mgr(), &rendez),
      done);
  run_opts.rendezvous = rendez;

  std::vector<Tensor>* rets = new std::vector<Tensor>;
  const string& func_name = func_->name();
  profiler::TraceMe trace_me(
      [&] {
        return absl::StrCat(
            "PartitionedCallOp #parent_step_id=", ctx->step_id(),
            ",function_step_id=", run_opts.step_id, "#");
      },
      /*level=*/2);
  lib->Run(run_opts, handle, inputs, rets,
           [rets, rendez, done, ctx, func_name,
            step_container](const Status& status) {
             if (!status.ok()) {
               const string function_and_msg =
                   strings::StrCat(errors::FormatFunctionForError(func_name),
                                   " ", status.error_message());
               ctx->SetStatus(Status(status.code(), function_and_msg));
             } else {
               for (int i = 0; i < rets->size(); ++i) {
                 ctx->set_output(i, (*rets)[i]);
               }
             }
             delete rets;
             delete step_container;
             rendez->Unref();
             done();
           });
}