[xlang] Cross language serialize ActorHandle (#7134)

src/ray/common/function_descriptor.cc

@@ -39,8 +39,6 @@ FunctionDescriptor FunctionDescriptorBuilder::FromProto(rpc::FunctionDescriptor
   default:
     break;
   }
-  RAY_LOG(DEBUG) << "Unknown function descriptor case: "
-                 << message.function_descriptor_case();
   // When TaskSpecification() constructed without function_descriptor set,
   // we should return a valid ray::FunctionDescriptor instance.
   return FunctionDescriptorBuilder::Empty();

src/ray/core_worker/actor_handle.cc

@@ -7,7 +7,8 @@ namespace {
 ray::rpc::ActorHandle CreateInnerActorHandle(
     const class ActorID &actor_id, const class JobID &job_id,
     const ObjectID &initial_cursor, const Language actor_language, bool is_direct_call,
-    const ray::FunctionDescriptor &actor_creation_task_function_descriptor) {
+    const ray::FunctionDescriptor &actor_creation_task_function_descriptor,
+    const std::string &extension_data) {
   ray::rpc::ActorHandle inner;
   inner.set_actor_id(actor_id.Data(), actor_id.Size());
   inner.set_creation_job_id(job_id.Data(), job_id.Size());
@@ -16,6 +17,7 @@ ray::rpc::ActorHandle CreateInnerActorHandle(
       actor_creation_task_function_descriptor->GetMessage();
   inner.set_actor_cursor(initial_cursor.Binary());
   inner.set_is_direct_call(is_direct_call);
+  inner.set_extension_data(extension_data);
   return inner;
 }
 
@@ -32,10 +34,11 @@ namespace ray {
 ActorHandle::ActorHandle(
     const class ActorID &actor_id, const class JobID &job_id,
     const ObjectID &initial_cursor, const Language actor_language, bool is_direct_call,
-    const ray::FunctionDescriptor &actor_creation_task_function_descriptor)
-    : ActorHandle(CreateInnerActorHandle(actor_id, job_id, initial_cursor, actor_language,
-                                         is_direct_call,
-                                         actor_creation_task_function_descriptor)) {}
+    const ray::FunctionDescriptor &actor_creation_task_function_descriptor,
+    const std::string &extension_data)
+    : ActorHandle(CreateInnerActorHandle(
+          actor_id, job_id, initial_cursor, actor_language, is_direct_call,
+          actor_creation_task_function_descriptor, extension_data)) {}
 
 ActorHandle::ActorHandle(const std::string &serialized)
     : ActorHandle(CreateInnerActorHandleFromString(serialized)) {}

src/ray/core_worker/actor_handle.h

@@ -21,7 +21,8 @@ class ActorHandle {
   ActorHandle(const ActorID &actor_id, const JobID &job_id,
               const ObjectID &initial_cursor, const Language actor_language,
               bool is_direct_call,
-              const ray::FunctionDescriptor &actor_creation_task_function_descriptor);
+              const ray::FunctionDescriptor &actor_creation_task_function_descriptor,
+              const std::string &extension_data);
 
   /// Constructs an ActorHandle from a serialized string.
   ActorHandle(const std::string &serialized);
@@ -39,6 +40,8 @@ class ActorHandle {
         inner_.actor_creation_task_function_descriptor());
   };
 
+  std::string ExtensionData() const { return inner_.extension_data(); }
+
   bool IsDirectCallActor() const { return inner_.is_direct_call(); }
 
   void SetActorTaskSpec(TaskSpecBuilder &builder, const TaskTransportType transport_type,

src/ray/core_worker/core_worker.cc

@@ -698,6 +698,7 @@ Status CoreWorker::SubmitTask(const RayFunction &function,
 Status CoreWorker::CreateActor(const RayFunction &function,
                                const std::vector<TaskArg> &args,
                                const ActorCreationOptions &actor_creation_options,
+                               const std::string &extension_data,
                                ActorID *return_actor_id) {
   const int next_task_index = worker_context_.GetNextTaskIndex();
   const ActorID actor_id =
@@ -720,9 +721,10 @@ Status CoreWorker::CreateActor(const RayFunction &function,
       actor_creation_options.is_direct_call, actor_creation_options.max_concurrency,
       actor_creation_options.is_detached, actor_creation_options.is_asyncio);
 
-  std::unique_ptr<ActorHandle> actor_handle(new ActorHandle(
-      actor_id, job_id, /*actor_cursor=*/return_ids[0], function.GetLanguage(),
-      actor_creation_options.is_direct_call, function.GetFunctionDescriptor()));
+  std::unique_ptr<ActorHandle> actor_handle(
+      new ActorHandle(actor_id, job_id, /*actor_cursor=*/return_ids[0],
+                      function.GetLanguage(), actor_creation_options.is_direct_call,
+                      function.GetFunctionDescriptor(), extension_data));
   RAY_CHECK(AddActorHandle(std::move(actor_handle)))
       << "Actor " << actor_id << " already exists";
 

src/ray/core_worker/core_worker.h

@@ -348,13 +348,14 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   /// \param[in] function The remote function that generates the actor object.
   /// \param[in] args Arguments of this task.
   /// \param[in] actor_creation_options Options for this actor creation task.
-  /// \param[out] actor_handle Handle to the actor.
+  /// \param[in] extension_data Extension data of the actor handle,
+  /// see `ActorHandle` in `core_worker.proto`.
   /// \param[out] actor_id ID of the created actor. This can be used to submit
   /// tasks on the actor.
   /// \return Status error if actor creation fails, likely due to raylet failure.
   Status CreateActor(const RayFunction &function, const std::vector<TaskArg> &args,
                      const ActorCreationOptions &actor_creation_options,
-                     ActorID *actor_id);
+                     const std::string &extension_data, ActorID *actor_id);
 
   /// Submit an actor task.
   ///

src/ray/core_worker/lib/java/org_ray_runtime_RayNativeRuntime.h

@@ -21,8 +21,8 @@ JNIEXPORT jlong JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeInitCoreWork
  * Method:    nativeRunTaskExecutor
  * Signature: (J)V
  */
-JNIEXPORT void JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeRunTaskExecutor(
-    JNIEnv *, jclass, jlong);
+JNIEXPORT void JNICALL
+Java_org_ray_runtime_RayNativeRuntime_nativeRunTaskExecutor(JNIEnv *, jclass, jlong);
 
 /*
  * Class:     org_ray_runtime_RayNativeRuntime

src/ray/core_worker/lib/java/org_ray_runtime_actor_NativeRayActor.cc

@@ -13,6 +13,16 @@ inline ray::CoreWorker &GetCoreWorker(jlong nativeCoreWorkerPointer) {
 extern "C" {
 #endif
 
+JNIEXPORT jint JNICALL Java_org_ray_runtime_actor_NativeRayActor_nativeGetLanguage(
+    JNIEnv *env, jclass o, jlong nativeCoreWorkerPointer, jbyteArray actorId) {
+  auto actor_id = JavaByteArrayToId<ray::ActorID>(env, actorId);
+  ray::ActorHandle *native_actor_handle = nullptr;
+  auto status = GetCoreWorker(nativeCoreWorkerPointer)
+                    .GetActorHandle(actor_id, &native_actor_handle);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, false);
+  return native_actor_handle->ActorLanguage();
+}
+
 JNIEXPORT jboolean JNICALL
 Java_org_ray_runtime_actor_NativeRayActor_nativeIsDirectCallActor(
     JNIEnv *env, jclass o, jlong nativeCoreWorkerPointer, jbyteArray actorId) {

src/ray/core_worker/lib/java/org_ray_runtime_actor_NativeRayActor.h

@@ -7,6 +7,14 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+/*
+ * Class:     org_ray_runtime_actor_NativeRayActor
+ * Method:    nativeGetLanguage
+ * Signature: (J[B)I
+ */
+JNIEXPORT jint JNICALL Java_org_ray_runtime_actor_NativeRayActor_nativeGetLanguage(
+    JNIEnv *, jclass, jlong, jbyteArray);
+
 /*
  * Class:     org_ray_runtime_actor_NativeRayActor
  * Method:    nativeIsDirectCallActor

src/ray/core_worker/lib/java/org_ray_runtime_task_NativeTaskSubmitter.cc

@@ -157,9 +157,9 @@ Java_org_ray_runtime_task_NativeTaskSubmitter_nativeCreateActor(
   auto actor_creation_options = ToActorCreationOptions(env, actorCreationOptions);
 
   ray::ActorID actor_id;
-  auto status =
-      GetCoreWorker(nativeCoreWorkerPointer)
-          .CreateActor(ray_function, task_args, actor_creation_options, &actor_id);
+  auto status = GetCoreWorker(nativeCoreWorkerPointer)
+                    .CreateActor(ray_function, task_args, actor_creation_options,
+                                 /*extension_data*/ "", &actor_id);
 
   THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
   return IdToJavaByteArray<ray::ActorID>(env, actor_id);

src/ray/core_worker/test/core_worker_test.cc

@@ -63,7 +63,8 @@ ActorID CreateActorHelper(CoreWorker &worker,
 
   // Create an actor.
   ActorID actor_id;
-  RAY_CHECK_OK(worker.CreateActor(func, args, actor_options, &actor_id));
+  RAY_CHECK_OK(
+      worker.CreateActor(func, args, actor_options, /*extension_data*/ "", &actor_id));
   return actor_id;
 }
 
@@ -614,7 +615,7 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) {
   const auto job_id = NextJobId();
   ActorHandle actor_handle(ActorID::Of(job_id, TaskID::ForDriverTask(job_id), 1), job_id,
                            ObjectID::FromRandom(), function.GetLanguage(), true,
-                           function.GetFunctionDescriptor());
+                           function.GetFunctionDescriptor(), "");
 
   // Manually create `num_tasks` task specs, and for each of them create a
   // `PushTaskRequest`, this is to batch performance of TaskSpec
@@ -729,7 +730,7 @@ TEST_F(ZeroNodeTest, TestActorHandle) {
   JobID job_id = NextJobId();
   ActorHandle original(ActorID::Of(job_id, TaskID::ForDriverTask(job_id), 0), job_id,
                        ObjectID::FromRandom(), Language::PYTHON, /*is_direct_call=*/false,
-                       ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""));
+                       ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""), "");
   std::string output;
   original.Serialize(&output);
   ActorHandle deserialized(output);

src/ray/protobuf/core_worker.proto

@@ -30,6 +30,9 @@ message ActorHandle {
 
   // Whether direct actor call is used.
   bool is_direct_call = 7;
+
+  // An extension field that is used for storing app-language-specific data.
+  bytes extension_data = 8;
 }
 
 message AssignTaskRequest {