TensorFlow中的通信机制——Rendezvous（二）gRPC传输

Status BaseRemoteRendezvous::Send(const Rendezvous::ParsedKey& parsed,
const Rendezvous::Args& args,
const Tensor& val, const bool is_dead) {
VLOG() << "BaseRemoteRendezvous Send " << this << " " << parsed.FullKey();
{
mutex_lock l(mu_);
if (!status_.ok()) return status_;
DCHECK(is_initialized_locked());
if (!IsLocalDevice(session_->worker_name, parsed.src_device)) {
return errors::InvalidArgument(
"Invalid rendezvous key (src): ", parsed.FullKey(), " @ ",
session_->worker_name);
}
}
// Buffers "val" and "device_context" in local_.
return local_->Send(parsed, args, val, is_dead);
}

// See worker.proto for details.
rpc RecvTensor(RecvTensorRequest) returns (RecvTensorResponse) {
// RecvTensor Method
}

// Names of worker methods.
enum class GrpcWorkerMethod {
kGetStatus,
kCreateWorkerSession,
kDeleteWorkerSession,
kRegisterGraph,
kDeregisterGraph,
kRunGraph,
kCleanupGraph,
kCleanupAll,
kRecvTensor,
kRecvBuf,
kLogging,
kTracing,
kCompleteGroup,
kCompleteInstance,
kGetStepSequence,
};

const char* GrpcWorkerMethodName(GrpcWorkerMethod id) {
switch (id) {
case GrpcWorkerMethod::kGetStatus:
return "/tensorflow.WorkerService/GetStatus";
case GrpcWorkerMethod::kCreateWorkerSession:
return "/tensorflow.WorkerService/CreateWorkerSession";
case GrpcWorkerMethod::kDeleteWorkerSession:
return "/tensorflow.WorkerService/DeleteWorkerSession";
case GrpcWorkerMethod::kRegisterGraph:
return "/tensorflow.WorkerService/RegisterGraph";
case GrpcWorkerMethod::kDeregisterGraph:
return "/tensorflow.WorkerService/DeregisterGraph";
case GrpcWorkerMethod::kRunGraph:
return "/tensorflow.WorkerService/RunGraph";
case GrpcWorkerMethod::kCleanupGraph:
return "/tensorflow.WorkerService/CleanupGraph";
case GrpcWorkerMethod::kCleanupAll:
return "/tensorflow.WorkerService/CleanupAll";
case GrpcWorkerMethod::kRecvTensor:
return "/tensorflow.WorkerService/RecvTensor";
case GrpcWorkerMethod::kRecvBuf:
return "/tensorflow.WorkerService/RecvBuf";
case GrpcWorkerMethod::kLogging:
return "/tensorflow.WorkerService/Logging";
case GrpcWorkerMethod::kTracing:
return "/tensorflow.WorkerService/Tracing";
case GrpcWorkerMethod::kCompleteGroup:
return "/tensorflow.WorkerService/CompleteGroup";
case GrpcWorkerMethod::kCompleteInstance:
return "/tensorflow.WorkerService/CompleteInstance";
case GrpcWorkerMethod::kGetStepSequence:
return "/tensorflow.WorkerService/GetStepSequence";
}
// Shouldn't be reached.
LOG(FATAL) << "Invalid id: this line shouldn't be reached.";
return "invalid id";
}

WorkerService::AsyncService::AsyncService() {
for (int i = ; i < kGrpcNumWorkerMethods; ++i) {
AddMethod(new ::grpc::internal::RpcServiceMethod(
GrpcWorkerMethodName(static_cast<GrpcWorkerMethod>(i)),
::grpc::internal::RpcMethod::NORMAL_RPC, nullptr));
::grpc::Service::MarkMethodAsync(i);
}
}

void RpcRemoteRendezvous::RecvFromRemoteAsync(
const Rendezvous::ParsedKey& parsed, const Rendezvous::Args& recv_args,
DoneCallback done) {
CHECK(is_initialized());
Status s; // Prepare a RecvTensor call that can handle being aborted.
RpcRecvTensorCall* call = get_call_freelist()->New(); // key.src_device identifies a remote device.
if (!DeviceNameUtils::SplitDeviceName(parsed.src_device, &call->src_worker_,
&call->src_rel_device_)) {
s = errors::Internal(parsed.src_device,
" is invalid remote source device.");
}
WorkerSession* sess = session();
WorkerInterface* rwi = sess->worker_cache->CreateWorker(call->src_worker_);
if (s.ok() && rwi == nullptr) {
s = errors::Internal("No worker known as ", call->src_worker_);
} Device* dst_device;
if (s.ok()) {
s = sess->device_mgr()->LookupDevice(parsed.dst_device, &dst_device);
}
if (!s.ok()) {
if (rwi != nullptr) {
sess->worker_cache->ReleaseWorker(call->src_worker_, rwi);
}
get_call_freelist()->Release(call, sess->worker_cache.get());
done(s, Args(), recv_args, Tensor{}, false);
return;
} call->Init(rwi, step_id_, parsed.FullKey(), recv_args.alloc_attrs, dst_device,
recv_args, std::move(done)); // Record "call" in active_ so that it can be aborted cleanly.
RegisterCall(call); // RendezvousMgr already aborted, shouldn't send RPC call any more
if (!call->status().ok()) {
call->done()(call->status(), Args(), Args(), Tensor(), false);
session()->worker_cache->ReleaseWorker(call->src_worker_, call->wi_);
call->wi_ = nullptr;
get_call_freelist()->Release(call, session()->worker_cache.get());
return;
} // Start "call".
Ref();
call->Start([this, call]() {
// Removes "call" from active_. Prevent StartAbort().
DeregisterCall(call);
// If StartAbort was called prior to DeregisterCall, then the
// current status should be bad.
Status s = call->status();
call->done()(s, Args(), call->recv_args(), call->tensor(), call->is_dead());
session()->worker_cache->ReleaseWorker(call->src_worker_, call->wi_);
call->wi_ = nullptr;
get_call_freelist()->Release(call, session()->worker_cache.get());
Unref();
});
}

void RecvTensorAsync(CallOptions* call_opts, const RecvTensorRequest* request,
TensorResponse* response, StatusCallback done) override {
VLOG() << "RecvTensorAsync req: " << request->DebugString();
int64 start_usec = Env::Default()->NowMicros();
// Type-specialized logging for this method.
bool logging_active = logger_->LoggingActive() || VLOG_IS_ON();
StatusCallback wrapper_done;
const StatusCallback* cb_to_use;
if (!logging_active) {
cb_to_use = &done; // No additional work to do, so just use done directly
} else {
wrapper_done = [this, request, response, done, start_usec](Status s) {
if (logger_->LoggingActive()) {
int64 end_usec = Env::Default()->NowMicros();
int64 step_id = request->step_id();
int64 bytes = response->tensor().TotalBytes();
int64 send_start_usec = start_usec;
// If a send start time was reported by the other side, use
// that instead. Maybe we should mark the display if we're using
// our local time instead of the remote start time?
if (response->metadata().send_start_micros()) {
// send_start_micros is the timestamp taken when the
// remote machine began to send the RecvTensor response.
// Due to clock skew between source and dest machines, it
// is possible that send_start_micros can be larger than
// end_usec or less than start_usec.
//
// To respect causality, we enforce the invariants that
// the RecvTensor response can not have been sent before
// the RecvTensor request, and must have been sent before
// it was received.
send_start_usec = std::max(
start_usec,
static_cast<int64>(response->metadata().send_start_micros()));
send_start_usec = std::min(send_start_usec, end_usec - );
}
const string& key = request->rendezvous_key();
std::vector<string> key_parts = str_util::Split(key, ';');
if (key_parts.size() != ) {
LOG(WARNING) << "Bad key: " << key;
} else {
logger_->RecordRecvTensor(step_id, send_start_usec, end_usec,
key_parts[], // tensor name
key_parts[], // src_device
key_parts[], // dst_device
bytes);
}
}
VLOG() << "done callback, req: " << request->DebugString()
<< " response " << response->metadata().DebugString();
done(s);
};
cb_to_use = &wrapper_done;
} IssueRequest(request, response, recvtensor_, *cb_to_use, call_opts);
}

void StartCall() {
context_.reset(new ::grpc::ClientContext());
context_->set_fail_fast(fail_fast_); if (timeout_in_ms_ > ) {
context_->set_deadline(
gpr_time_from_millis(timeout_in_ms_, GPR_TIMESPAN));
}
if (call_opts_) {
call_opts_->SetCancelCallback([this]() { context_->TryCancel(); });
} VLOG() << "Starting call: " << method_; call_ = std::move(
stub_->PrepareUnaryCall(context_.get(), method_, request_buf_, cq_));
call_->StartCall();
call_->Finish(&response_buf_, &status_, this);
}

env_->rendezvous_mgr->RecvLocalAsync(
step_id, parsed,
[opts, response, done, src_dev, request](
const Status& status, const Rendezvous::Args& send_args,
const Rendezvous::Args& recv_args, const Tensor& val,
const bool is_dead) {
opts->ClearCancelCallback();
if (status.ok()) {
// DMA can only be used for Tensors that do not fall into
// the following three odd edge cases: 1) a zero-size
// buffer, 2) a dead tensor which has an uninit value, and
// 3) the tensor has the on_host allocation attribute,
// i.e. it's in CPU RAM *independent of its assigned
// device type*.
const bool on_host = send_args.alloc_attrs.on_host();
{
// Non-DMA cases.
if (src_dev->tensorflow_gpu_device_info() && (!on_host)) {
DeviceContext* send_dev_context = send_args.device_context;
AllocatorAttributes alloc_attrs;
alloc_attrs.set_gpu_compatible(true);
alloc_attrs.set_on_host(true);
Allocator* alloc = src_dev->GetAllocator(alloc_attrs);
Tensor* copy = new Tensor(alloc, val.dtype(), val.shape());
CHECK(send_dev_context)
<< "send dev name: " << src_dev->name()
<< " gpu_info: " << src_dev->tensorflow_gpu_device_info();
// "val" is on an accelerator device. Uses the device_context to
// fill the copy on host.
StatusCallback copy_ready = [response, done, copy,
is_dead](const Status& s) {
// The value is now ready to be returned on the wire.
grpc::EncodeTensorToByteBuffer(is_dead, *copy, response);
done(s);
delete copy;
}; send_dev_context->CopyDeviceTensorToCPU(
&val, request->rendezvous_key(), src_dev, copy, copy_ready);
} else {
grpc::EncodeTensorToByteBuffer(is_dead, val, response);
done(Status::OK());
}
}
} else {
// !s.ok()
done(status);
}
});