Skip to content

Commit 6aeeedd

Browse files
committed
build: gate grpc eager server on android
1 parent dcdbc86 commit 6aeeedd

1 file changed

Lines changed: 74 additions & 0 deletions

File tree

scripts/patch_tfjava.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -806,6 +806,78 @@ def patch_workspace(path: Path) -> None:
806806
"""
807807

808808

809+
GRPC_SERVER_LIB_ANDROID_BUILD_PATCH = """--- a/tensorflow/core/distributed_runtime/rpc/BUILD
810+
+++ b/tensorflow/core/distributed_runtime/rpc/BUILD
811+
@@ -312,8 +312,12 @@
812+
"//tensorflow/core/distributed_runtime:session_mgr",
813+
"//tensorflow/core/distributed_runtime:worker_cache_wrapper",
814+
"//tensorflow/core/distributed_runtime:worker_env",
815+
- "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_service_impl",
816+
- ],
817+
+ ] + select({
818+
+ "//tensorflow:android": [],
819+
+ "//conditions:default": [
820+
+ "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_service_impl",
821+
+ ],
822+
+ }),
823+
alwayslink = 1,
824+
)
825+
826+
"""
827+
828+
829+
GRPC_SERVER_LIB_ANDROID_CC_PATCH = """--- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
830+
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
831+
@@ -34,7 +34,9 @@
832+
#include "tensorflow/core/distributed_runtime/master_env.h"
833+
#include "tensorflow/core/distributed_runtime/master_session.h"
834+
#include "tensorflow/core/distributed_runtime/rpc/async_service_interface.h"
835+
+#if !defined(IS_MOBILE_PLATFORM)
836+
#include "tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h"
837+
+#endif // !defined(IS_MOBILE_PLATFORM)
838+
#include "tensorflow/core/distributed_runtime/rpc/grpc_channel.h"
839+
#include "tensorflow/core/distributed_runtime/rpc/grpc_master_service.h"
840+
#include "tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.h"
841+
@@ -231,7 +233,9 @@
842+
worker_service_ = NewGrpcWorkerService(worker_impl_.get(), &builder,
843+
opts.worker_service_options)
844+
.release();
845+
+#if !defined(IS_MOBILE_PLATFORM)
846+
eager_service_ = new eager::GrpcEagerServiceImpl(&worker_env_, &builder);
847+
+#endif // !defined(IS_MOBILE_PLATFORM)
848+
849+
// extra service:
850+
if (opts.service_func != nullptr) {
851+
@@ -383,9 +387,11 @@
852+
worker_thread_.reset(
853+
env_->StartThread(ThreadOptions(), "TF_worker_service",
854+
[this] { worker_service_->HandleRPCsLoop(); }));
855+
+#if !defined(IS_MOBILE_PLATFORM)
856+
eager_thread_.reset(
857+
env_->StartThread(ThreadOptions(), "TF_eager_service",
858+
[this] { eager_service_->HandleRPCsLoop(); }));
859+
+#endif // !defined(IS_MOBILE_PLATFORM)
860+
state_ = STARTED;
861+
LOG(INFO) << "Started server with target: " << target();
862+
return Status::OK();
863+
@@ -402,9 +408,14 @@
864+
865+
Status GrpcServer::AddMasterEagerContextToEagerService(
866+
const tensorflow::uint64 context_id, tensorflow::EagerContext* context) {
867+
+#if defined(IS_MOBILE_PLATFORM)
868+
+ return errors::Unimplemented(
869+
+ "GRPC eager service is not supported on mobile builds.");
870+
+#else
871+
auto* eager_service =
872+
static_cast<eager::GrpcEagerServiceImpl*>(eager_service_);
873+
return eager_service->CreateMasterContext(context_id, context);
874+
+#endif // defined(IS_MOBILE_PLATFORM)
875+
}
876+
877+
Status GrpcServer::UpdateServerDef(const ServerDef& server_def) {
878+
"""
879+
880+
809881
SAVED_MODEL_ANDROID_LOADER_PATCH = """--- a/tensorflow/cc/saved_model/BUILD
810882
+++ b/tensorflow/cc/saved_model/BUILD
811883
@@ -92,7 +92,7 @@
@@ -871,6 +943,8 @@ def write_tensorflow_android_absl_patch(path: Path) -> None:
871943
text += TF_C_API_EXPERIMENTAL_BUILD_PATCH
872944
text += TF_C_API_EXPERIMENTAL_CC_PATCH
873945
text += DISTRIBUTED_EAGER_ANDROID_PATCH
946+
text += GRPC_SERVER_LIB_ANDROID_BUILD_PATCH
947+
text += GRPC_SERVER_LIB_ANDROID_CC_PATCH
874948
text += SAVED_MODEL_ANDROID_LOADER_PATCH
875949
text += TENSORFLOW_FRAMEWORK_ANDROID_PATCH
876950
if not text.endswith("\n"):

0 commit comments

Comments
 (0)