[xlang] Cross language serialization for ActorHandle (#10335)

java/runtime/src/main/java/io/ray/runtime/object/ObjectSerializer.java

@@ -2,6 +2,7 @@ package io.ray.runtime.object;
 
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.ray.api.id.ObjectId;
+import io.ray.runtime.actor.NativeActorHandle;
 import io.ray.runtime.exception.RayActorException;
 import io.ray.runtime.exception.RayTaskException;
 import io.ray.runtime.exception.RayWorkerException;
@@ -35,6 +36,10 @@ public class ObjectSerializer {
   public static final byte[] OBJECT_METADATA_TYPE_JAVA = "JAVA".getBytes();
   public static final byte[] OBJECT_METADATA_TYPE_PYTHON = "PYTHON".getBytes();
   public static final byte[] OBJECT_METADATA_TYPE_RAW = "RAW".getBytes();
+  // A constant used as object metadata to indicate the object is an actor handle.
+  // This value should be synchronized with the Python definition in ray_constants.py
+  // TODO(fyrestone): Serialize the ActorHandle via the custom type feature of XLANG.
+  public static final byte[] OBJECT_METADATA_TYPE_ACTOR_HANDLE = "ACTOR_HANDLE".getBytes();
 
   // When an outer object is being serialized, the nested ObjectRefs are all
   // serialized and the writeExternal method of the nested ObjectRefs are
@@ -86,6 +91,9 @@ public class ObjectSerializer {
               "Can't deserialize RayTaskException object: " + objectId
                   .toString());
         }
+      } else if (Arrays.equals(meta, OBJECT_METADATA_TYPE_ACTOR_HANDLE)) {
+        byte[] serialized = Serializer.decode(data, byte[].class);
+        return NativeActorHandle.fromBytes(serialized);
       } else if (Arrays.equals(meta, OBJECT_METADATA_TYPE_PYTHON)) {
         throw new IllegalArgumentException("Can't deserialize Python object: " + objectId
             .toString());
@@ -129,6 +137,13 @@ public class ObjectSerializer {
       // Only OBJECT_METADATA_TYPE_RAW is raw bytes,
       // any other type should be the MessagePack serialized bytes.
       return new NativeRayObject(serializedBytes, TASK_EXECUTION_EXCEPTION_META);
+    } else if (object instanceof NativeActorHandle) {
+      NativeActorHandle actorHandle = (NativeActorHandle)object;
+      byte[] serializedBytes = Serializer.encode(actorHandle.toBytes()).getLeft();
+      // serializedBytes is MessagePack serialized bytes
+      // Only OBJECT_METADATA_TYPE_RAW is raw bytes,
+      // any other type should be the MessagePack serialized bytes.
+      return new NativeRayObject(serializedBytes, OBJECT_METADATA_TYPE_ACTOR_HANDLE);
     } else {
       try {
         Pair<byte[], Boolean> serialized = Serializer.encode(object);

java/runtime/src/main/java/io/ray/runtime/task/ArgumentsBuilder.java

@@ -48,7 +48,8 @@ public class ArgumentsBuilder {
         if (language != Language.JAVA) {
           boolean isCrossData =
               Arrays.equals(value.metadata, ObjectSerializer.OBJECT_METADATA_TYPE_CROSS_LANGUAGE) ||
-                  Arrays.equals(value.metadata, ObjectSerializer.OBJECT_METADATA_TYPE_RAW);
+                  Arrays.equals(value.metadata, ObjectSerializer.OBJECT_METADATA_TYPE_RAW) ||
+                  Arrays.equals(value.metadata, ObjectSerializer.OBJECT_METADATA_TYPE_ACTOR_HANDLE);
           if (!isCrossData) {
             throw new IllegalArgumentException(String.format("Can't transfer %s data to %s",
                 Arrays.toString(value.metadata), language.getValueDescriptor().getName()));

java/test/src/main/java/io/ray/test/CrossLanguageInvocationTest.java

@@ -167,23 +167,21 @@ public class CrossLanguageInvocationTest extends BaseMultiLanguageTest {
     // Create a java actor, and pass actor handle to python.
     ActorHandle<TestActor> javaActor = Ray.actor(TestActor::new, "1".getBytes()).remote();
     Preconditions.checkState(javaActor instanceof NativeActorHandle);
-    byte[] actorHandleBytes = ((NativeActorHandle) javaActor).toBytes();
     ObjectRef<byte[]> res = Ray.task(
         PyFunction.of(PYTHON_MODULE,
             "py_func_call_java_actor_from_handle",
             byte[].class),
-        actorHandleBytes).remote();
+        javaActor).remote();
     Assert.assertEquals(res.get(), "12".getBytes());
     // Create a python actor, and pass actor handle to python.
     PyActorHandle pyActor = Ray.actor(
         PyActorClass.of(PYTHON_MODULE, "Counter"), "1".getBytes()).remote();
     Preconditions.checkState(pyActor instanceof NativeActorHandle);
-    actorHandleBytes = ((NativeActorHandle) pyActor).toBytes();
     res = Ray.task(
         PyFunction.of(PYTHON_MODULE,
             "py_func_call_python_actor_from_handle",
             byte[].class),
-        actorHandleBytes).remote();
+        pyActor).remote();
     Assert.assertEquals(res.get(), "3".getBytes());
   }
 
@@ -301,9 +299,8 @@ public class CrossLanguageInvocationTest extends BaseMultiLanguageTest {
     return l;
   }
 
-  public static byte[] callPythonActorHandle(byte[] value) {
+  public static byte[] callPythonActorHandle(PyActorHandle actor) {
     // This function will be called from test_cross_language_invocation.py
-    NativePyActorHandle actor = (NativePyActorHandle) NativeActorHandle.fromBytes(value);
     ObjectRef<byte[]> res = actor.task(
         PyActorMethod.of("increase", byte[].class),
         "1".getBytes()).remote();

java/test/src/main/resources/test_cross_language_invocation.py

@@ -59,17 +59,13 @@ def py_func_call_java_actor(value):
 
 
 @ray.remote
-def py_func_call_java_actor_from_handle(value):
-    assert isinstance(value, bytes)
-    actor_handle = ray.actor.ActorHandle._deserialization_helper(value)
+def py_func_call_java_actor_from_handle(actor_handle):
     r = actor_handle.concat.remote(b"2")
     return ray.get(r)
 
 
 @ray.remote
-def py_func_call_python_actor_from_handle(value):
-    assert isinstance(value, bytes)
-    actor_handle = ray.actor.ActorHandle._deserialization_helper(value)
+def py_func_call_python_actor_from_handle(actor_handle):
     r = actor_handle.increase.remote(2)
     return ray.get(r)
 
@@ -79,7 +75,7 @@ def py_func_pass_python_actor_handle():
     counter = Counter.remote(2)
     f = ray.java_function("io.ray.test.CrossLanguageInvocationTest",
                           "callPythonActorHandle")
-    r = f.remote(counter._serialization_helper()[0])
+    r = f.remote(counter)
     return ray.get(r)