[Streaming] fix streaming ci

BUILD.bazel

@@ -367,6 +367,27 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "exported_lib",
+    srcs = glob(
+        [
+            "src/ray/exported/exported.cc",
+        ],
+    ),
+    hdrs = glob(
+        [
+            "src/ray/exported/exported.h",
+        ],
+    ),
+    copts = COPTS,
+    alwayslink = 1,
+    strip_include_prefix = "src",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":core_worker_lib",
+    ],
+)
+
 cc_binary(
     name = "raylet",
     srcs = ["src/ray/raylet/main.cc"],
@@ -1701,6 +1722,7 @@ pyx_library(
         linkstatic = 1,
     ),
     deps = [
+        "//:exported_lib",
         "//:core_worker_lib",
         "//:global_state_accessor_lib",
         "//:ray_util",
@@ -1755,6 +1777,7 @@ cc_binary(
     linkstatic = 1,
     visibility = ["//java:__subpackages__"],
     deps = [
+        "//:exported_lib",
         "//:core_worker_lib",
         "//:global_state_accessor_lib",
         "//:src/ray/ray_exported_symbols.lds",

src/ray/common/id.h

@@ -75,7 +75,8 @@ class BaseID {
  protected:
   BaseID(const std::string &binary) {
     RAY_CHECK(binary.size() == Size() || binary.size() == 0)
-        << "expected size is " << Size() << ", but got " << binary.size();
+        << "expected size is " << Size() << ", but got data " << binary << " of size "
+        << binary.size();
     std::memcpy(const_cast<uint8_t *>(this->Data()), binary.data(), binary.size());
   }
   // All IDs are immutable for hash evaluations. MutableData is only allow to use
@@ -341,24 +342,25 @@ std::ostream &operator<<(std::ostream &os, const TaskID &id);
 std::ostream &operator<<(std::ostream &os, const ObjectID &id);
 std::ostream &operator<<(std::ostream &os, const PlacementGroupID &id);
 
-#define DEFINE_UNIQUE_ID(type)                                                 \
-  class RAY_EXPORT type : public UniqueID {                                    \
-   public:                                                                     \
-    explicit type(const UniqueID &from) {                                      \
-      std::memcpy(&id_, from.Data(), kUniqueIDSize);                           \
-    }                                                                          \
-    type() : UniqueID() {}                                                     \
-    static type FromRandom() { return type(UniqueID::FromRandom()); }          \
-    static type FromBinary(const std::string &binary) { return type(binary); } \
-    static type Nil() { return type(UniqueID::Nil()); }                        \
-    static size_t Size() { return kUniqueIDSize; }                             \
-                                                                               \
-   private:                                                                    \
-    explicit type(const std::string &binary) {                                 \
-      RAY_CHECK(binary.size() == Size() || binary.size() == 0)                 \
-          << "expected size is " << Size() << ", but got " << binary.size();   \
-      std::memcpy(&id_, binary.data(), binary.size());                         \
-    }                                                                          \
+#define DEFINE_UNIQUE_ID(type)                                                           \
+  class RAY_EXPORT type : public UniqueID {                                              \
+   public:                                                                               \
+    explicit type(const UniqueID &from) {                                                \
+      std::memcpy(&id_, from.Data(), kUniqueIDSize);                                     \
+    }                                                                                    \
+    type() : UniqueID() {}                                                               \
+    static type FromRandom() { return type(UniqueID::FromRandom()); }                    \
+    static type FromBinary(const std::string &binary) { return type(binary); }           \
+    static type Nil() { return type(UniqueID::Nil()); }                                  \
+    static size_t Size() { return kUniqueIDSize; }                                       \
+                                                                                         \
+   private:                                                                              \
+    explicit type(const std::string &binary) {                                           \
+      RAY_CHECK(binary.size() == Size() || binary.size() == 0)                           \
+          << "expected size is " << Size() << ", but got data " << binary << " of size " \
+          << binary.size();                                                              \
+      std::memcpy(&id_, binary.data(), binary.size());                                   \
+    }                                                                                    \
   };
 
 #include "ray/common/id_def.h"
@@ -385,7 +387,8 @@ T BaseID<T>::FromRandom() {
 template <typename T>
 T BaseID<T>::FromBinary(const std::string &binary) {
   RAY_CHECK(binary.size() == T::Size() || binary.size() == 0)
-      << "expected size is " << T::Size() << ", but got " << binary.size();
+      << "expected size is " << T::Size() << ", but got data " << binary << " of size "
+      << binary.size();
   T t;
   std::memcpy(t.MutableData(), binary.data(), binary.size());
   return t;

src/ray/common/task/task_spec.cc

@@ -120,7 +120,7 @@ ObjectID TaskSpecification::ReturnId(size_t return_index) const {
 }
 
 bool TaskSpecification::ArgByRef(size_t arg_index) const {
-  return message_->args(arg_index).object_ref().object_id() != "";
+  return message_->args(arg_index).has_object_ref();
 }
 
 ObjectID TaskSpecification::ArgId(size_t arg_index) const {

src/ray/common/task/task_util.h

@@ -41,7 +41,7 @@ class TaskArgByValue : public TaskArg {
   ///
   /// \param[in] value Value of the argument.
   /// \return The task argument.
-  TaskArgByValue(const std::shared_ptr<RayObject> &value) : value_(value) {
+  explicit TaskArgByValue(const std::shared_ptr<RayObject> &value) : value_(value) {
     RAY_CHECK(value) << "Value can't be null.";
   }
 

src/ray/core_worker/transport/dependency_resolver.cc

@@ -46,7 +46,7 @@ void InlineDependencies(
         if (!it->second->IsInPlasmaError()) {
           // The object has not been promoted to plasma. Inline the object by
           // clearing the reference and replacing it with the raw value.
-          mutable_arg->mutable_object_ref()->Clear();
+          mutable_arg->clear_object_ref();
           if (it->second->HasData()) {
             const auto &data = it->second->GetData();
             mutable_arg->set_data(data->Data(), data->Size());

src/ray/exported/exported.cc

@@ -0,0 +1,34 @@
+#include "ray/core_worker/core_worker.h"
+#include "ray/exported/test.h"
+
+namespace ray {
+namespace exported_ {
+
+void SubmitActorTask(const ActorID &peer_actor_id,
+                     std::shared_ptr<LocalMemoryBuffer> buffer, RayFunction &function,
+                     int return_num, std::vector<ObjectID> &return_ids) {
+  std::unordered_map<std::string, double> resources;
+  std::string name = function.GetFunctionDescriptor()->DefaultTaskName();
+  TaskOptions options{name, return_num, resources};
+
+  char meta_data[3] = {'R', 'A', 'W'};
+  std::shared_ptr<LocalMemoryBuffer> meta =
+      std::make_shared<LocalMemoryBuffer>((uint8_t *)meta_data, 3, true);
+
+  std::vector<std::unique_ptr<TaskArg>> args;
+  if (function.GetLanguage() == Language::PYTHON) {
+    auto dummy = "__RAY_DUMMY__";
+    std::shared_ptr<LocalMemoryBuffer> dummyBuffer =
+        std::make_shared<LocalMemoryBuffer>((uint8_t *)dummy, 13, true);
+    args.emplace_back(new TaskArgByValue(std::make_shared<RayObject>(
+        std::move(dummyBuffer), meta, std::vector<ObjectID>(), true)));
+  }
+  args.emplace_back(new TaskArgByValue(std::make_shared<RayObject>(
+      std::move(buffer), meta, std::vector<ObjectID>(), true)));
+
+  std::vector<std::shared_ptr<RayObject>> results;
+  CoreWorkerProcess::GetCoreWorker().SubmitActorTask(peer_actor_id, function, args,
+                                                     options, &return_ids);
+}
+}  // namespace exported_
+}  // namespace ray

src/ray/exported/exported.h

@@ -0,0 +1,14 @@
+#pragma once
+#include "ray/common/buffer.h"
+#include "ray/common/id.h"
+#include "ray/core_worker/common.h"
+
+namespace ray {
+namespace exported_ {
+
+void SubmitActorTask(const ActorID &peer_actor_id,
+                     std::shared_ptr<LocalMemoryBuffer> buffer, RayFunction &function,
+                     int return_num, std::vector<ObjectID> &return_ids);
+
+}  // namespace exported_
+}  // namespace ray

src/ray/ray_exported_symbols.lds

@@ -21,9 +21,11 @@
 *ray*TaskID*
 *ray*ActorID*
 *ray*ObjectID*
+*ray*ObjectReference*
 # Others
 *ray*CoreWorker*
 *PyInit*
 *init_raylet*
 *Java*
 *JNI_*
+*ray*exported_*

src/ray/ray_version_script.lds

@@ -4,6 +4,7 @@
 # Ray ABI is not finalized, the exact set of exported (C/C++) APIs is subject to change.
 VERSION_1.0 {
     global:
+        *ray*exported_*
         # common
         *ray*Language*;
         *ray*RayObject*;
@@ -23,6 +24,7 @@ VERSION_1.0 {
         *ray*TaskID*;
         *ray*ActorID*;
         *ray*ObjectID*;
+        *ray*ObjectReference*;
         # Others
         *ray*CoreWorker*;
         *PyInit*;

streaming/BUILD.bazel

@@ -61,6 +61,14 @@ cc_binary(
     deps = ["//:core_worker_lib"],
 )
 
+cc_binary(
+    name = "exported_lib.so",
+    copts = COPTS,
+    linkshared = 1,
+    deps = ["//:exported_lib"],
+)
+
+
 cc_library(
     name = "streaming_util",
     srcs = glob([
@@ -139,6 +147,7 @@ cc_library(
         ],
         "//conditions:default": [
             "core_worker_lib.so",
+            "exported_lib.so",
         ],
     }),
 )
@@ -283,7 +292,9 @@ genrule(
         GENERATED_DIR="streaming/python/generated"
         mkdir -p "$$GENERATED_DIR"
         touch "$$GENERATED_DIR/__init__.py"
+        # Use this `sed` command to change the import path in the generated file.
         sed -i -E 's/from streaming.src.protobuf/from ./' "$$GENERATED_DIR/remote_call_pb2.py"
+        sed -i -E 's/from protobuf/from ./' "$$GENERATED_DIR/remote_call_pb2.py"
         date > $@
     """,
     local = 1,

streaming/python/tests/test_hybrid_stream.py

@@ -1,8 +1,10 @@
 import json
+import os
+import subprocess
+
 import ray
 from ray.streaming import StreamingContext
-import subprocess
-import os
+from ray.test_utils import wait_for_condition
 
 
 def map_func1(x):
@@ -45,6 +47,7 @@ def sink_func(x):
         print("HybridStreamTest", x)
         with open(sink_file, "a") as f:
             f.write(str(x))
+            f.flush()
 
     ctx = StreamingContext.Builder().build()
     ctx.from_values("a", "b", "c") \
@@ -54,14 +57,23 @@ def sink_func(x):
         .as_python_stream() \
         .sink(sink_func)
     ctx.submit("HybridStreamTest")
-    import time
-    time.sleep(3)
+
+    def check_succeed():
+        if os.path.exists(sink_file):
+            import time
+            time.sleep(3)  # Wait all data be written
+            with open(sink_file, "r") as f:
+                result = f.read()
+                assert "a" in result
+                assert "b" not in result
+                assert "c" in result
+            print("Execution succeed")
+            return True
+        return False
+
+    wait_for_condition(check_succeed, timeout=60, retry_interval_ms=1000)
+    print("Execution succeed")
     ray.shutdown()
-    with open(sink_file, "r") as f:
-        result = f.read()
-        assert "a" in result
-        assert "b" not in result
-        assert "c" in result
 
 
 if __name__ == "__main__":

streaming/src/queue/transport.cc

@@ -1,50 +1,24 @@
 #include "queue/transport.h"
 
 #include "queue/utils.h"
+#include "ray/exported/test.h"
 
 namespace ray {
 namespace streaming {
 
 static constexpr int TASK_OPTION_RETURN_NUM_0 = 0;
 static constexpr int TASK_OPTION_RETURN_NUM_1 = 1;
 
-void Transport::SendInternal(std::shared_ptr<LocalMemoryBuffer> buffer,
-                             RayFunction &function, int return_num,
-                             std::vector<ObjectID> &return_ids) {
-  std::unordered_map<std::string, double> resources;
-  std::string name = function.GetFunctionDescriptor()->DefaultTaskName();
-  TaskOptions options{name, return_num, resources};
-
-  char meta_data[3] = {'R', 'A', 'W'};
-  std::shared_ptr<LocalMemoryBuffer> meta =
-      std::make_shared<LocalMemoryBuffer>((uint8_t *)meta_data, 3, true);
-
-  std::vector<std::unique_ptr<TaskArg>> args;
-  if (function.GetLanguage() == Language::PYTHON) {
-    auto dummy = "__RAY_DUMMY__";
-    std::shared_ptr<LocalMemoryBuffer> dummyBuffer =
-        std::make_shared<LocalMemoryBuffer>((uint8_t *)dummy, 13, true);
-    args.emplace_back(new TaskArgByValue(std::make_shared<RayObject>(
-        std::move(dummyBuffer), meta, std::vector<ObjectID>(), true)));
-  }
-  args.emplace_back(new TaskArgByValue(std::make_shared<RayObject>(
-      std::move(buffer), meta, std::vector<ObjectID>(), true)));
-
-  std::vector<std::shared_ptr<RayObject>> results;
-  CoreWorkerProcess::GetCoreWorker().SubmitActorTask(peer_actor_id_, function, args,
-                                                     options, &return_ids);
-}
-
 void Transport::Send(std::shared_ptr<LocalMemoryBuffer> buffer) {
   STREAMING_LOG(INFO) << "Transport::Send buffer size: " << buffer->Size();
   std::vector<ObjectID> return_ids;
-  SendInternal(std::move(buffer), async_func_, TASK_OPTION_RETURN_NUM_0, return_ids);
+  ray::exported_::SendInternal(peer_actor_id_, std::move(buffer), async_func_, TASK_OPTION_RETURN_NUM_0, return_ids);
 }
 
 std::shared_ptr<LocalMemoryBuffer> Transport::SendForResult(
     std::shared_ptr<LocalMemoryBuffer> buffer, int64_t timeout_ms) {
   std::vector<ObjectID> return_ids;
-  SendInternal(buffer, sync_func_, TASK_OPTION_RETURN_NUM_1, return_ids);
+  ray_exported::ray::SendInternal(peer_actor_id_, buffer, sync_func_, TASK_OPTION_RETURN_NUM_1, return_ids);
 
   std::vector<std::shared_ptr<RayObject>> results;
   Status get_st =

streaming/src/queue/transport.h

@@ -47,16 +47,6 @@ class Transport {
   std::shared_ptr<LocalMemoryBuffer> SendForResultWithRetry(
       std::shared_ptr<LocalMemoryBuffer> buffer, int retry_cnt, int64_t timeout_ms);
 
- private:
-  /// Send buffer internal
-  /// \param[in] buffer buffer to be sent.
-  /// \param[in] function the function descriptor of peer's function.
-  /// \param[in] return_num return value number of the call.
-  /// \param[out] return_ids return ids from SubmitActorTask.
-  virtual void SendInternal(std::shared_ptr<LocalMemoryBuffer> buffer,
-                            RayFunction &function, int return_num,
-                            std::vector<ObjectID> &return_ids);
-
  private:
   WorkerID worker_id_;
   ActorID peer_actor_id_;