| 73 | } |
| 74 | |
| 75 | void StarRunGraphOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { |
| 76 | #if TENSORFLOW_USE_STAR |
| 77 | // NOTE(jiankeng.pt) Only grpc++ support zero copy now. |
| 78 | if (ctx->num_inputs() != feed_names_.size()) { |
| 79 | LOG(FATAL) << "Op input size must equal to feed_names size, " |
| 80 | << ctx->num_inputs() << " .vs " << feed_names_.size(); |
| 81 | } |
| 82 | if (ctx->num_outputs() != fetch_names_.size()) { |
| 83 | LOG(FATAL) << "Op input size must equal to fetch_names size, " |
| 84 | << ctx->num_inputs() << " .vs " << fetch_names_.size(); |
| 85 | } |
| 86 | StarRunGraphRequest* req = new StarRunGraphRequest(); |
| 87 | req->graph_handle_ = graph_handle_; |
| 88 | req->step_id_ = ctx->round_step_id(); |
| 89 | req->ps_graph_count_ = ps_graph_count_; |
| 90 | req->feed_names_.resize(feed_names_.size()); |
| 91 | req->feed_tensors_.resize(feed_names_.size()); |
| 92 | req->is_dead_.resize(feed_names_.size()); |
| 93 | |
| 94 | for (unsigned i = 0; i < feed_names_.size(); ++i) { |
| 95 | req->feed_names_[i] = feed_names_[i]; |
| 96 | CHECK(!IsRefType(data_type_[i])) << "Data type is ref type."; |
| 97 | req->feed_tensors_[i] = ctx->input(i); |
| 98 | req->is_dead_[i] = ctx->is_input_dead(i); |
| 99 | } |
| 100 | |
| 101 | std::unordered_map<std::string, int> fetch_name_idxs; |
| 102 | req->fetch_names_.resize(fetch_names_.size()); |
| 103 | for (unsigned i = 0; i < fetch_names_.size(); ++i) { |
| 104 | req->fetch_names_[i] = fetch_names_[i]; |
| 105 | fetch_name_idxs[fetch_names_[i]] = i; |
| 106 | } |
| 107 | |
| 108 | StarRunGraphResponse* resp = new StarRunGraphResponse(); |
| 109 | resp->device_ = ctx->device(); |
| 110 | |
| 111 | StarWorkerInterface* seastar_worker = dynamic_cast<StarWorkerInterface*>(worker_); |
| 112 | if (seastar_worker == nullptr) { |
| 113 | LOG(FATAL) << "Error worker in star run graph op. Required StarWorkerInterface."; |
| 114 | } |
| 115 | |
| 116 | seastar_worker->StarRunGraphAsync(req, resp, [this, fetch_name_idxs, |
| 117 | done, ctx, req, resp] (const Status& status) { |
| 118 | ctx->SetStatus(status); |
| 119 | if (status.ok()) { |
| 120 | if (this->fetch_names_.size() != resp->fetch_names_.size() |
| 121 | || this->fetch_names_.size() != resp->fetch_tensors_.size()) { |
| 122 | ctx->SetStatus(errors::Internal("run graph failed. num_recv_fetch_names[", |
| 123 | resp->fetch_names_.size() , "], num_recv_fetch_tensors[", |
| 124 | resp->fetch_tensors_.size(), "], num_fetch_count[" |
| 125 | , this->fetch_names_.size(), "]")); |
| 126 | delete req; |
| 127 | delete resp; |
| 128 | done(); |
| 129 | return; |
| 130 | } |
| 131 | |
| 132 | for (size_t i = 0; i < this->fetch_names_.size(); ++i) { |
nothing calls this directly
no test coverage detected