| 754 | |
| 755 | namespace { |
| 756 | tensorflow::Status EnableCollectiveOps(const tensorflow::ServerDef& server_def, |
| 757 | TFE_Context* ctx) { |
| 758 | // We don't use the TF_RETURN_IF_ERROR macro directly since that destroys the |
| 759 | // server object (which currently CHECK-fails) and we miss the error, instead, |
| 760 | // we log the error, and then return to allow the user to see the error |
| 761 | // message. |
| 762 | #define LOG_AND_RETURN_IF_ERROR(...) \ |
| 763 | do { \ |
| 764 | const ::tensorflow::Status _status = (__VA_ARGS__); \ |
| 765 | if (TF_PREDICT_FALSE(!_status.ok())) { \ |
| 766 | LOG(ERROR) << _status.error_message(); \ |
| 767 | return _status; \ |
| 768 | } \ |
| 769 | } while (0); |
| 770 | |
| 771 | std::unique_ptr<tensorflow::ServerInterface> server; |
| 772 | LOG_AND_RETURN_IF_ERROR(tensorflow::NewServer(server_def, &server)); |
| 773 | |
| 774 | tensorflow::GrpcServer* grpc_server = |
| 775 | dynamic_cast<tensorflow::GrpcServer*>(server.get()); |
| 776 | if (grpc_server == nullptr) { |
| 777 | LOG_AND_RETURN_IF_ERROR(tensorflow::errors::Internal( |
| 778 | "Currently, TFE_NewContext only supports tensorflow::GrpcServer.")); |
| 779 | } |
| 780 | |
| 781 | LOG_AND_RETURN_IF_ERROR(grpc_server->Start()); |
| 782 | |
| 783 | LOG_AND_RETURN_IF_ERROR(ctx->context->StoreCollectiveOpsServer( |
| 784 | std::move(server), grpc_server->worker_env()->device_mgr, |
| 785 | grpc_server->worker_env()->collective_executor_mgr)); |
| 786 | |
| 787 | return tensorflow::Status::OK(); |
| 788 | #undef LOG_AND_RETURN_IF_ERROR |
| 789 | } |
| 790 | } // namespace |
| 791 | |
| 792 | // Set server_def on the context, possibly updating it. |
no test coverage detected