Method ExecuteAsync

tensorflow/core/common_runtime/base_collective_executor.cc:220–281

}

void BaseCollectiveExecutor::ExecuteAsync(OpKernelContext* ctx,
                                          const CollectiveParams& col_params,
                                          const string& exec_key,
                                          StatusCallback done) {
  // On any individual collective Op failure we need to abort the
  // BufRendezvous so that other Ops in the instance don't hang
  // waiting for transmissions that will never happen.  Do so after a
  // delay so that the original error status is more likely to
  // propagate up, and peers are unlikely to re-create the purged
  // BufRendezvous by late-arriving requests.
  StatusCallback done_safe = [this, done](const Status& s) {
    if (!s.ok()) {
      Ref();  // Ensure this lasts until the closure executes.
      SchedNonBlockingClosureAfter(1000000, [this, s] {
        remote_access_->buf_rendezvous()->StartAbort(s);
        Unref();
      });
    }
    done(s);
  };

  Tensor* output = ctx->mutable_output(0);
  const Tensor* input = (col_params.instance.type == REDUCTION_COLLECTIVE ||
                         col_params.instance.type == GATHER_COLLECTIVE ||
                         (col_params.instance.type == BROADCAST_COLLECTIVE &&
                          col_params.is_source))
                            ? &ctx->input(0)
                            : nullptr;
  CollectiveImplementationInterface* col_impl = nullptr;
  Status status = CreateCollective(col_params, &col_impl);
  if (!status.ok()) {
    done_safe(status);
    DCHECK_EQ(nullptr, col_impl);
    return;
  }
  CollectiveContext* col_ctx =
      new CollectiveContext(this, dev_mgr_, ctx, CtxParams(ctx), col_params,
                            exec_key, step_id_, input, output);
  status = col_impl->InitializeCollectiveContext(col_ctx);
  if (!status.ok()) {
    done_safe(status);
    delete col_ctx;
    delete col_impl;
    return;
  }
  // Run on an unbounded work queue that can handle blocking work so as to not
  // starve executor threads.
  remote_access_->RunClosure([col_impl, col_ctx, done_safe, ctx]() {
    profiler::TraceMe activity(
        [&] {
          return strings::StrCat(ctx->op_kernel().name(), ":",
                                 ctx->op_kernel().type_string(),
                                 "#id=", ctx->step_id(), "#");
        },
        profiler::TraceMeLevel::kInfo);
    col_impl->Run([col_impl, col_ctx, done_safe](const Status& s) {
      done_safe(s);
      delete col_ctx;