| 1033 | } |
| 1034 | |
| 1035 | Status GrpcTpuDriver::Close() { |
| 1036 | auto stub = CreateTpuDriverStub(config_, creds_); |
| 1037 | ::grpc::ClientContext ctx; |
| 1038 | ctx.set_fail_fast(false); |
| 1039 | ctx.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10)); |
| 1040 | CloseRequest req; |
| 1041 | req.set_client_id(client_id_); |
| 1042 | CloseResponse resp; |
| 1043 | ::grpc::Status status = stub->Close(&ctx, req, &resp); |
| 1044 | if (!status.ok()) { |
| 1045 | return xla::Status(tensorflow::error::Code(status.error_code()), |
| 1046 | absl::StrCat("Failed to close TPU driver. Error was: ", |
| 1047 | status.error_message(), |
| 1048 | ". Details: ", status.error_details())); |
| 1049 | } |
| 1050 | closed_ = true; |
| 1051 | return Status::OK(); |
| 1052 | } |
| 1053 | } // namespace |
| 1054 | |
| 1055 | xla::StatusOr<std::unique_ptr<TpuDriver>> CreateGrpcTpuDriver( |
no test coverage detected