| 84 | }; |
| 85 | |
| 86 | void RunGraphOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { |
| 87 | OP_REQUIRES(ctx, ctx->num_inputs() == feed_names_.size(), |
| 88 | errors::Internal("Op input size must equal to feed_names size, ", |
| 89 | ctx->num_inputs(), " .vs ", feed_names_.size())); |
| 90 | |
| 91 | OP_REQUIRES(ctx, ctx->num_outputs() == fetch_names_.size(), |
| 92 | errors::Internal("Op input size must equal to fetch_names size, ", |
| 93 | ctx->num_inputs(), " .vs ", fetch_names_.size())); |
| 94 | |
| 95 | CallContext *call_ctx = new CallContext(); |
| 96 | call_ctx->req = worker_->CreateRunGraphRequest(); |
| 97 | call_ctx->res = worker_->CreateRunGraphResponse(); |
| 98 | |
| 99 | call_ctx->req->set_graph_handle(graph_handle_); |
| 100 | call_ctx->req->set_step_id(ctx->round_step_id()); |
| 101 | call_ctx->req->set_run_graph_mode(true); |
| 102 | |
| 103 | for(int i = 0; i < feed_names_.size(); i++) { |
| 104 | CHECK(!IsRefType(data_type_[i])) << "Data type is ref type."; |
| 105 | call_ctx->req->add_send(feed_names_[i], ctx->input(i)); |
| 106 | } |
| 107 | |
| 108 | std::unordered_map<std::string, int> fetch_name_idxs; |
| 109 | for (int i = 0; i < fetch_names_.size(); i++) { |
| 110 | call_ctx->req->add_recv_key(fetch_names_[i]); |
| 111 | fetch_name_idxs[fetch_names_[i]] = i; |
| 112 | } |
| 113 | |
| 114 | worker_->RunGraphAsync(&call_ctx->opts, call_ctx->req, call_ctx->res, |
| 115 | [ctx, call_ctx, this, done, fetch_name_idxs] |
| 116 | (const Status &status) { |
| 117 | |
| 118 | ctx->SetStatus(status); |
| 119 | if (status.ok()) { |
| 120 | if (this->fetch_names_.size() != call_ctx->res->num_recvs()) { |
| 121 | ctx->SetStatus(errors::Internal("run graph failed. num_recvs[", |
| 122 | call_ctx->res->num_recvs(), "] != num_fetchs[", |
| 123 | this->fetch_names_.size(), "]")); |
| 124 | delete call_ctx; |
| 125 | done(); |
| 126 | return; |
| 127 | } |
| 128 | |
| 129 | for (size_t i = 0; i < this->fetch_names_.size(); i++) { |
| 130 | Tensor tensor; |
| 131 | Status s = call_ctx->res->RecvValue(i, &tensor); |
| 132 | if (!s.ok()) { |
| 133 | ctx->SetStatus(s); |
| 134 | break; |
| 135 | } |
| 136 | const string &fetch_name = call_ctx->res->recv_key(i); |
| 137 | auto it = fetch_name_idxs.find(fetch_name); |
| 138 | if (it == fetch_name_idxs.end()) { |
| 139 | ctx->SetStatus(errors::Internal("can not find recv_key in fetch names, recv_key:", fetch_name)); |
| 140 | break; |
| 141 | } |
| 142 | int idx = it->second; |
| 143 |
nothing calls this directly
no test coverage detected