| 1011 | } |
| 1012 | |
| 1013 | Status GrpcTpuDriver::Reset() { |
| 1014 | auto stub = CreateTpuDriverStub(config_, creds_); |
| 1015 | ::grpc::ClientContext ctx; |
| 1016 | ctx.set_fail_fast(false); |
| 1017 | ctx.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(10)); |
| 1018 | ResetRequest req; |
| 1019 | ResetResponse resp; |
| 1020 | ::grpc::Status status = stub->Reset(&ctx, req, &resp); |
| 1021 | if (!status.ok()) { |
| 1022 | LOG(ERROR) << "Failed to reset the gRPC driver: " << status.error_code() |
| 1023 | << ": " << status.error_message() << ": " |
| 1024 | << status.error_details(); |
| 1025 | return xla::Status(tensorflow::error::Code(status.error_code()), |
| 1026 | absl::StrCat("Failed to reset TPU driver. Error was: ", |
| 1027 | status.error_message(), |
| 1028 | ". Details: ", status.error_details())); |
| 1029 | } |
| 1030 | streams_.clear(); |
| 1031 | host_stream_.reset(); |
| 1032 | return Close(); |
| 1033 | } |
| 1034 | |
| 1035 | Status GrpcTpuDriver::Close() { |
| 1036 | auto stub = CreateTpuDriverStub(config_, creds_); |