[gRPC] Migrate gcs data structures to protobuf (#5024)

ray-project · Jun 25, 2019 · 0131353 · 0131353
1 parent bd8aceb
commit 0131353
Show file tree

Hide file tree

Showing 52 changed files with 1,465 additions and 1,642 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -1,22 +1,55 @@
 # Bazel build
 # C/C++ documentation: https://docs.bazel.build/versions/master/be/c-cpp.html
 
-load("@com_github_grpc_grpc//bazel:grpc_build_system.bzl", "grpc_proto_library")
+load("@com_github_grpc_grpc//bazel:cc_grpc_library.bzl", "cc_grpc_library")
+load("@build_stack_rules_proto//python:python_proto_compile.bzl", "python_proto_compile")
 load("@com_github_google_flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
 load("@//bazel:ray.bzl", "flatbuffer_py_library")
 load("@//bazel:cython_library.bzl", "pyx_library")
 
 COPTS = ["-DRAY_USE_GLOG"]
 
-# Node manager gRPC lib.
-grpc_proto_library(
-    name = "node_manager_grpc_lib",
+# === Begin of protobuf definitions ===
+
+proto_library(
+    name = "gcs_proto",
+    srcs = ["src/ray/protobuf/gcs.proto"],
+    visibility = ["//java:__subpackages__"],
+)
+
+cc_proto_library(
+    name = "gcs_cc_proto",
+    deps = [":gcs_proto"],
+)
+
+python_proto_compile(
+    name = "gcs_py_proto",
+    deps = [":gcs_proto"],
+)
+
+proto_library(
+    name = "node_manager_proto",
     srcs = ["src/ray/protobuf/node_manager.proto"],
 )
 
+cc_proto_library(
+    name = "node_manager_cc_proto",
+    deps = ["node_manager_proto"],
+)
+
+# === End of protobuf definitions ===
+
+# Node manager gRPC lib.
+cc_grpc_library(
+    name = "node_manager_cc_grpc",
+    srcs = [":node_manager_proto"],
+    grpc_only = True,
+    deps = [":node_manager_cc_proto"],
+)
+
 # Node manager server and client.
 cc_library(
-    name = "node_manager_rpc_lib",
+    name = "node_manager_rpc",
     srcs = glob([
         "src/ray/rpc/*.cc",
     ]),
@@ -25,7 +58,7 @@ cc_library(
     ]),
     copts = COPTS,
     deps = [
-        ":node_manager_grpc_lib",
+        ":node_manager_cc_grpc",
         ":ray_common",
         "@boost//:asio",
         "@com_github_grpc_grpc//:grpc++",
@@ -114,7 +147,7 @@ cc_library(
         ":gcs",
         ":gcs_fbs",
         ":node_manager_fbs",
-        ":node_manager_rpc_lib",
+        ":node_manager_rpc",
         ":object_manager",
         ":ray_common",
         ":ray_util",
@@ -422,9 +455,11 @@ cc_library(
         "src/ray/gcs/format",
     ],
     deps = [
+        ":gcs_cc_proto",
         ":gcs_fbs",
         ":hiredis",
         ":node_manager_fbs",
+        ":node_manager_rpc",
         ":ray_common",
         ":ray_util",
         ":stats_lib",
@@ -555,46 +590,6 @@ filegroup(
     visibility = ["//java:__subpackages__"],
 )
 
-flatbuffer_py_library(
-    name = "python_gcs_fbs",
-    srcs = [
-        ":gcs_fbs_file",
-    ],
-    outs = [
-        "ActorCheckpointIdData.py",
-        "ActorState.py",
-        "ActorTableData.py",
-        "Arg.py",
-        "ClassTableData.py",
-        "ClientTableData.py",
-        "ConfigTableData.py",
-        "CustomSerializerData.py",
-        "DriverTableData.py",
-        "EntryType.py",
-        "ErrorTableData.py",
-        "ErrorType.py",
-        "FunctionTableData.py",
-        "GcsEntry.py",
-        "HeartbeatBatchTableData.py",
-        "HeartbeatTableData.py",
-        "Language.py",
-        "ObjectTableData.py",
-        "ProfileEvent.py",
-        "ProfileTableData.py",
-        "RayResource.py",
-        "ResourcePair.py",
-        "SchedulingState.py",
-        "TablePrefix.py",
-        "TablePubsub.py",
-        "TaskInfo.py",
-        "TaskLeaseData.py",
-        "TaskReconstructionData.py",
-        "TaskTableData.py",
-        "TaskTableTestAndUpdate.py",
-    ],
-    out_prefix = "python/ray/core/generated/",
-)
-
 flatbuffer_py_library(
     name = "python_node_manager_fbs",
     srcs = [
@@ -679,6 +674,7 @@ cc_binary(
     linkstatic = 1,
     visibility = ["//java:__subpackages__"],
     deps = [
+        ":gcs_cc_proto",
         ":ray_common",
     ],
 )
@@ -688,7 +684,7 @@ genrule(
     srcs = [
         "python/ray/_raylet.so",
         "//:python_sources",
-        "//:python_gcs_fbs",
+        "//:gcs_py_proto",
         "//:python_node_manager_fbs",
         "//:redis-server",
         "//:redis-cli",
@@ -710,11 +706,13 @@ genrule(
         cp -f $(location //:raylet_monitor) $$WORK_DIR/python/ray/core/src/ray/raylet/ &&
         cp -f $(location @plasma//:plasma_store_server) $$WORK_DIR/python/ray/core/src/plasma/ &&
         cp -f $(location //:raylet) $$WORK_DIR/python/ray/core/src/ray/raylet/ &&
-        for f in $(locations //:python_gcs_fbs); do cp -f $$f $$WORK_DIR/python/ray/core/generated/; done &&
         mkdir -p $$WORK_DIR/python/ray/core/generated/ray/protocol/ &&
         for f in $(locations //:python_node_manager_fbs); do
             cp -f $$f $$WORK_DIR/python/ray/core/generated/ray/protocol/;
         done &&
+        for f in $(locations //:gcs_py_proto); do
+            cp -f $$f $$WORK_DIR/python/ray/core/generated/;
+        done &&
         echo $$WORK_DIR > $@
     """,
     local = 1,

diff --git a/bazel/ray_deps_build_all.bzl b/bazel/ray_deps_build_all.bzl
@@ -4,6 +4,8 @@ load("@com_github_jupp0r_prometheus_cpp//:repositories.bzl", "prometheus_cpp_rep
 load("@com_github_ray_project_ray//bazel:python_configure.bzl", "python_configure")
 load("@com_github_checkstyle_java//:repo.bzl", "checkstyle_deps")
 load("@com_github_grpc_grpc//bazel:grpc_deps.bzl", "grpc_deps")
+load("@build_stack_rules_proto//java:deps.bzl", "java_proto_compile")
+load("@build_stack_rules_proto//python:deps.bzl", "python_proto_compile")
 
 
 def ray_deps_build_all():
@@ -13,4 +15,6 @@ def ray_deps_build_all():
   prometheus_cpp_repositories()
   python_configure(name = "local_config_python")
   grpc_deps()
+  java_proto_compile()
+  python_proto_compile()
 
diff --git a/bazel/ray_deps_setup.bzl b/bazel/ray_deps_setup.bzl
@@ -105,7 +105,14 @@ def ray_deps_setup():
     http_archive(
         name = "com_github_grpc_grpc",
         urls = [
-            "https://github.com/grpc/grpc/archive/7741e806a213cba63c96234f16d712a8aa101a49.tar.gz",
+            "https://github.com/grpc/grpc/archive/76a381869413834692b8ed305fbe923c0f9c4472.tar.gz",
         ],
-        strip_prefix = "grpc-7741e806a213cba63c96234f16d712a8aa101a49",
+        strip_prefix = "grpc-76a381869413834692b8ed305fbe923c0f9c4472",
+    )
+
+    http_archive(
+        name = "build_stack_rules_proto",
+        urls = ["https://github.com/stackb/rules_proto/archive/b93b544f851fdcd3fc5c3d47aee3b7ca158a8841.tar.gz"],
+        sha256 = "c62f0b442e82a6152fcd5b1c0b7c4028233a9e314078952b6b04253421d56d61",
+        strip_prefix = "rules_proto-b93b544f851fdcd3fc5c3d47aee3b7ca158a8841",
     )
diff --git a/doc/source/conf.py b/doc/source/conf.py
@@ -23,20 +23,7 @@
     "gym.spaces",
     "ray._raylet",
     "ray.core.generated",
-    "ray.core.generated.ActorCheckpointIdData",
-    "ray.core.generated.ClientTableData",
-    "ray.core.generated.DriverTableData",
-    "ray.core.generated.EntryType",
-    "ray.core.generated.ErrorTableData",
-    "ray.core.generated.ErrorType",
-    "ray.core.generated.GcsEntry",
-    "ray.core.generated.HeartbeatBatchTableData",
-    "ray.core.generated.HeartbeatTableData",
-    "ray.core.generated.Language",
-    "ray.core.generated.ObjectTableData",
-    "ray.core.generated.ProfileTableData",
-    "ray.core.generated.TablePrefix",
-    "ray.core.generated.TablePubsub",
+    "ray.core.generated.gcs_pb2",
     "ray.core.generated.ray.protocol.Task",
     "scipy",
     "scipy.signal",

diff --git a/java/BUILD.bazel b/java/BUILD.bazel
@@ -1,4 +1,5 @@
 load("//bazel:ray.bzl", "flatbuffer_java_library", "define_java_module")
+load("@build_stack_rules_proto//java:java_proto_compile.bzl", "java_proto_compile")
 
 exports_files([
     "testng.xml",
@@ -50,6 +51,7 @@ define_java_module(
     name = "runtime",
     additional_srcs = [
         ":generate_java_gcs_fbs",
+        ":gcs_java_proto",
     ],
     additional_resources = [
         ":java_native_deps",
@@ -68,6 +70,7 @@ define_java_module(
         "@plasma//:org_apache_arrow_arrow_plasma",
         "@maven//:com_github_davidmoten_flatbuffers_java",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_typesafe_config",
         "@maven//:commons_io_commons_io",
         "@maven//:de_ruedigermoeller_fst",
@@ -148,38 +151,16 @@ java_binary(
     ],
 )
 
+java_proto_compile(
+    name = "gcs_java_proto",
+    deps = ["@//:gcs_proto"],
+)
+
 flatbuffers_generated_files = [
-    "ActorCheckpointData.java",
-    "ActorCheckpointIdData.java",
-    "ActorState.java",
-    "ActorTableData.java",
     "Arg.java",
-    "ClassTableData.java",
-    "ClientTableData.java",
-    "ConfigTableData.java",
-    "CustomSerializerData.java",
-    "DriverTableData.java",
-    "EntryType.java",
-    "ErrorTableData.java",
-    "ErrorType.java",
-    "FunctionTableData.java",
-    "GcsEntry.java",
-    "HeartbeatBatchTableData.java",
-    "HeartbeatTableData.java",
     "Language.java",
-    "ObjectTableData.java",
-    "ProfileEvent.java",
-    "ProfileTableData.java",
-    "RayResource.java",
-    "ResourcePair.java",
-    "SchedulingState.java",
-    "TablePrefix.java",
-    "TablePubsub.java",
     "TaskInfo.java",
-    "TaskLeaseData.java",
-    "TaskReconstructionData.java",
-    "TaskTableData.java",
-    "TaskTableTestAndUpdate.java",
+    "ResourcePair.java",
 ]
 
 flatbuffer_java_library(
@@ -198,7 +179,7 @@ genrule(
     cmd = """
         for f in $(locations //java:java_gcs_fbs); do
             chmod +w $$f
-            cp -f $$f $(@D)/runtime/src/main/java/org/ray/runtime/generated
+            mv -f $$f $(@D)/runtime/src/main/java/org/ray/runtime/generated
         done
         python $$(pwd)/java/modify_generated_java_flatbuffers_files.py $(@D)/..
     """,
@@ -221,8 +202,10 @@ filegroup(
 genrule(
     name = "gen_maven_deps",
     srcs = [
-        ":java_native_deps",
+        ":gcs_java_proto",
         ":generate_java_gcs_fbs",
+        ":java_native_deps",
+        ":copy_pom_file",
         "@plasma//:org_apache_arrow_arrow_plasma",
     ],
     outs = ["gen_maven_deps.out"],
@@ -237,10 +220,15 @@ genrule(
             chmod +w $$f
             cp $$f $$NATIVE_DEPS_DIR
         done
-        # Copy flatbuffers-generated files
+        # Copy protobuf-generated files.
         GENERATED_DIR=$$WORK_DIR/java/runtime/src/main/java/org/ray/runtime/generated
         rm -rf $$GENERATED_DIR
         mkdir -p $$GENERATED_DIR
+        for f in $(locations //java:gcs_java_proto); do
+            unzip $$f
+            mv org/ray/runtime/generated/* $$GENERATED_DIR
+        done
+        # Copy flatbuffers-generated files
         for f in $(locations //java:generate_java_gcs_fbs); do
             cp $$f $$GENERATED_DIR
         done
@@ -250,6 +238,7 @@ genrule(
         echo $$(date) > $@
     """,
     local = 1,
+    tags = ["no-cache"],
 )
 
 genrule(

diff --git a/java/dependencies.bzl b/java/dependencies.bzl
@@ -6,6 +6,7 @@ def gen_java_deps():
             "com.beust:jcommander:1.72",
             "com.github.davidmoten:flatbuffers-java:1.9.0.1",
             "com.google.guava:guava:27.0.1-jre",
+            "com.google.protobuf:protobuf-java:3.8.0",
             "com.puppycrawl.tools:checkstyle:8.15",
             "com.sun.xml.bind:jaxb-core:2.3.0",
             "com.sun.xml.bind:jaxb-impl:2.3.0",

diff --git a/java/modify_generated_java_flatbuffers_files.py b/java/modify_generated_java_flatbuffers_files.py
@@ -4,7 +4,6 @@
 
 import os
 import sys
-
 """
 This script is used for modifying the generated java flatbuffer
 files for the reason: The package declaration in Java is different
@@ -21,35 +20,30 @@
 PACKAGE_DECLARATION = "package org.ray.runtime.generated;"
 
 
-def add_new_line(file, line_num, text):
+def add_package(file):
     with open(file, "r") as file_handler:
         lines = file_handler.readlines()
-        if (line_num <= 0) or (line_num > len(lines) + 1):
-            return False
 
-    lines.insert(line_num - 1, text + os.linesep)
+    if "FlatBuffers" not in lines[0]:
+        return
+
+    lines.insert(1, PACKAGE_DECLARATION + os.linesep)
     with open(file, "w") as file_handler:
         for line in lines:
             file_handler.write(line)
 
-    return True
-
 
 def add_package_declarations(generated_root_path):
     file_names = os.listdir(generated_root_path)
     for file_name in file_names:
         if not file_name.endswith(".java"):
             continue
         full_name = os.path.join(generated_root_path, file_name)
-        success = add_new_line(full_name, 2, PACKAGE_DECLARATION)
-        if not success:
-            raise RuntimeError("Failed to add package declarations, "
-                               "file name is %s" % full_name)
+        add_package(full_name)
 
 
 if __name__ == "__main__":
     ray_home = sys.argv[1]
     root_path = os.path.join(
-        ray_home,
-        "java/runtime/src/main/java/org/ray/runtime/generated")
+        ray_home, "java/runtime/src/main/java/org/ray/runtime/generated")
     add_package_declarations(root_path)