[core worker] Python core worker normal task submission (#5566)

src/ray/core_worker/common.h

@@ -13,11 +13,19 @@ namespace ray {
 using WorkerType = rpc::WorkerType;
 
 /// Information about a remote function.
-struct RayFunction {
-  /// Language of the remote function.
-  const Language language;
-  /// Function descriptor of the remote function.
-  const std::vector<std::string> function_descriptor;
+class RayFunction {
+ public:
+  RayFunction() {}
+  RayFunction(Language language, const std::vector<std::string> &function_descriptor)
+      : language_(language), function_descriptor_(function_descriptor) {}
+
+  Language GetLanguage() const { return language_; }
+
+  std::vector<std::string> GetFunctionDescriptor() const { return function_descriptor_; }
+
+ private:
+  Language language_;
+  std::vector<std::string> function_descriptor_;
 };
 
 /// Argument of a task.

src/ray/core_worker/context.cc

@@ -83,7 +83,7 @@ void WorkerContext::SetCurrentTaskId(const TaskID &task_id) {
 }
 
 void WorkerContext::SetCurrentTask(const TaskSpecification &task_spec) {
-  current_job_id_ = task_spec.JobId();
+  SetCurrentJobId(task_spec.JobId());
   GetThreadContext().SetCurrentTask(task_spec);
   if (task_spec.IsActorCreationTask()) {
     RAY_CHECK(current_actor_id_.IsNil());

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

@@ -46,7 +46,7 @@ JNIEXPORT jlong JNICALL Java_org_ray_runtime_RayNativeRuntime_nativeInitCoreWork
     RAY_CHECK(local_java_task_executor);
     // convert RayFunction
     jobject ray_function_array_list =
-        NativeStringVectorToJavaStringList(env, ray_function.function_descriptor);
+        NativeStringVectorToJavaStringList(env, ray_function.GetFunctionDescriptor());
     // convert args
     // TODO (kfstorm): Avoid copying binary data from Java to C++
     jobject args_array_list = NativeVectorToJavaList<std::shared_ptr<ray::RayObject>>(

src/ray/core_worker/store_provider/memory_store/memory_store.h

@@ -26,26 +26,6 @@ class CoreWorkerMemoryStore {
   /// \return Status.
   Status Put(const ObjectID &object_id, const RayObject &object);
 
-  /// Create and return a buffer in the object store that can be directly written
-  /// into. After writing to the buffer, the caller must call `Seal()` to finalize
-  /// the object. The `Create()` and `Seal()` combination is an alternative interface
-  /// to `Put()` that allows frontends to avoid an extra copy when possible.
-  ///
-  /// \param[in] metadata Metadata of the object to be written.
-  /// \param[in] data_size Size of the object to be written.
-  /// \param[in] object_id Object ID specified by the user.
-  /// \param[out] data Buffer for the user to write the object into.
-  /// \return Status.
-  Status Create(const std::shared_ptr<Buffer> &metadata, const size_t data_size,
-                const ObjectID &object_id, std::shared_ptr<Buffer> *data);
-
-  /// Finalize placing an object into the object store. This should be called after
-  /// a corresponding `Create()` call and then writing into the returned buffer.
-  ///
-  /// \param[in] object_id Object ID corresponding to the object.
-  /// \return Status.
-  Status Seal(const ObjectID &object_id);
-
   /// Get a list of objects from the object store.
   ///
   /// \param[in] object_ids IDs of the objects to get. Duplicates are not allowed.

src/ray/core_worker/task_interface.cc

@@ -121,10 +121,10 @@ void CoreWorkerTaskInterface::BuildCommonTaskSpec(
     const std::unordered_map<std::string, double> &required_placement_resources,
     TaskTransportType transport_type, std::vector<ObjectID> *return_ids) {
   // Build common task spec.
-  builder.SetCommonTaskSpec(task_id, function.language, function.function_descriptor,
-                            worker_context_.GetCurrentJobID(),
-                            worker_context_.GetCurrentTaskID(), task_index, num_returns,
-                            required_resources, required_placement_resources);
+  builder.SetCommonTaskSpec(
+      task_id, function.GetLanguage(), function.GetFunctionDescriptor(),
+      worker_context_.GetCurrentJobID(), worker_context_.GetCurrentTaskID(), task_index,
+      num_returns, required_resources, required_placement_resources);
   // Set task arguments.
   for (const auto &arg : args) {
     if (arg.IsPassedByReference()) {
@@ -177,8 +177,8 @@ Status CoreWorkerTaskInterface::CreateActor(
                                    actor_creation_options.is_direct_call);
 
   *actor_handle = std::unique_ptr<ActorHandle>(new ActorHandle(
-      actor_id, ActorHandleID::Nil(), function.language,
-      actor_creation_options.is_direct_call, function.function_descriptor));
+      actor_id, ActorHandleID::Nil(), function.GetLanguage(),
+      actor_creation_options.is_direct_call, function.GetFunctionDescriptor()));
   (*actor_handle)->IncreaseTaskCounter();
   (*actor_handle)->SetActorCursor(return_ids[0]);
 

src/ray/core_worker/task_interface.h

@@ -22,13 +22,13 @@ class CoreWorker;
 /// Options of a non-actor-creation task.
 struct TaskOptions {
   TaskOptions() {}
-  TaskOptions(int num_returns, const std::unordered_map<std::string, double> &resources)
+  TaskOptions(int num_returns, std::unordered_map<std::string, double> &resources)
       : num_returns(num_returns), resources(resources) {}
 
   /// Number of returns of this task.
-  const int num_returns = 1;
+  int num_returns = 1;
   /// Resources required by this task.
-  const std::unordered_map<std::string, double> resources;
+  std::unordered_map<std::string, double> resources;
 };
 
 /// Options of an actor creation task.

src/ray/core_worker/test/core_worker_test.cc

@@ -49,7 +49,7 @@ std::shared_ptr<Buffer> GenerateRandomBuffer() {
 }
 
 std::unique_ptr<ActorHandle> CreateActorHelper(
-    CoreWorker &worker, const std::unordered_map<std::string, double> &resources,
+    CoreWorker &worker, std::unordered_map<std::string, double> &resources,
     bool is_direct_call, uint64_t max_reconstructions) {
   std::unique_ptr<ActorHandle> actor_handle;
 
@@ -57,7 +57,7 @@ std::unique_ptr<ActorHandle> CreateActorHelper(
   uint8_t array[] = {1, 2, 3};
   auto buffer = std::make_shared<LocalMemoryBuffer>(array, sizeof(array));
 
-  RayFunction func{ray::Language::PYTHON, {"actor creation task"}};
+  RayFunction func(ray::Language::PYTHON, {"actor creation task"});
   std::vector<TaskArg> args;
   args.emplace_back(TaskArg::PassByValue(std::make_shared<RayObject>(buffer, nullptr)));
 
@@ -171,22 +171,22 @@ class CoreWorkerTest : public ::testing::Test {
   void TestStoreProvider(StoreProviderType type);
 
   // Test normal tasks.
-  void TestNormalTask(const std::unordered_map<std::string, double> &resources);
+  void TestNormalTask(std::unordered_map<std::string, double> &resources);
 
   // Test actor tasks.
-  void TestActorTask(const std::unordered_map<std::string, double> &resources,
+  void TestActorTask(std::unordered_map<std::string, double> &resources,
                      bool is_direct_call);
 
   // Test actor failure case, verify that the tasks would either succeed or
   // fail with exceptions, in that case the return objects fetched from `Get`
   // contain errors.
-  void TestActorFailure(const std::unordered_map<std::string, double> &resources,
+  void TestActorFailure(std::unordered_map<std::string, double> &resources,
                         bool is_direct_call);
 
   // Test actor failover case. Verify that actor can be reconstructed successfully,
   // and as long as we wait for actor reconstruction before submitting new tasks,
   // it is guaranteed that all tasks are successfully completed.
-  void TestActorReconstruction(const std::unordered_map<std::string, double> &resources,
+  void TestActorReconstruction(std::unordered_map<std::string, double> &resources,
                                bool is_direct_call);
 
  protected:
@@ -216,8 +216,7 @@ bool CoreWorkerTest::WaitForDirectCallActorState(CoreWorker &worker,
   return WaitForCondition(condition_func, timeout_ms);
 }
 
-void CoreWorkerTest::TestNormalTask(
-    const std::unordered_map<std::string, double> &resources) {
+void CoreWorkerTest::TestNormalTask(std::unordered_map<std::string, double> &resources) {
   CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
                     raylet_socket_names_[0], NextJobId(), gcs_options_, "", nullptr);
 
@@ -236,7 +235,7 @@ void CoreWorkerTest::TestNormalTask(
           TaskArg::PassByValue(std::make_shared<RayObject>(buffer1, nullptr)));
       args.emplace_back(TaskArg::PassByReference(object_id));
 
-      RayFunction func{ray::Language::PYTHON, {}};
+      RayFunction func(ray::Language::PYTHON, {});
       TaskOptions options;
 
       std::vector<ObjectID> return_ids;
@@ -258,8 +257,8 @@ void CoreWorkerTest::TestNormalTask(
   }
 }
 
-void CoreWorkerTest::TestActorTask(
-    const std::unordered_map<std::string, double> &resources, bool is_direct_call) {
+void CoreWorkerTest::TestActorTask(std::unordered_map<std::string, double> &resources,
+                                   bool is_direct_call) {
   CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
                     raylet_socket_names_[0], NextJobId(), gcs_options_, "", nullptr);
 
@@ -281,7 +280,7 @@ void CoreWorkerTest::TestActorTask(
 
       TaskOptions options{1, resources};
       std::vector<ObjectID> return_ids;
-      RayFunction func{ray::Language::PYTHON, {}};
+      RayFunction func(ray::Language::PYTHON, {});
 
       RAY_CHECK_OK(driver.Tasks().SubmitActorTask(*actor_handle, func, args, options,
                                                   &return_ids));
@@ -323,7 +322,7 @@ void CoreWorkerTest::TestActorTask(
 
     TaskOptions options{1, resources};
     std::vector<ObjectID> return_ids;
-    RayFunction func{ray::Language::PYTHON, {}};
+    RayFunction func(ray::Language::PYTHON, {});
     auto status =
         driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids);
     if (is_direct_call) {
@@ -348,7 +347,7 @@ void CoreWorkerTest::TestActorTask(
 }
 
 void CoreWorkerTest::TestActorReconstruction(
-    const std::unordered_map<std::string, double> &resources, bool is_direct_call) {
+    std::unordered_map<std::string, double> &resources, bool is_direct_call) {
   CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
                     raylet_socket_names_[0], NextJobId(), gcs_options_, "", nullptr);
 
@@ -389,7 +388,7 @@ void CoreWorkerTest::TestActorReconstruction(
 
       TaskOptions options{1, resources};
       std::vector<ObjectID> return_ids;
-      RayFunction func{ray::Language::PYTHON, {}};
+      RayFunction func(ray::Language::PYTHON, {});
 
       auto status =
           driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids);
@@ -405,8 +404,8 @@ void CoreWorkerTest::TestActorReconstruction(
   }
 }
 
-void CoreWorkerTest::TestActorFailure(
-    const std::unordered_map<std::string, double> &resources, bool is_direct_call) {
+void CoreWorkerTest::TestActorFailure(std::unordered_map<std::string, double> &resources,
+                                      bool is_direct_call) {
   CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0],
                     raylet_socket_names_[0], NextJobId(), gcs_options_, "", nullptr);
 
@@ -435,7 +434,7 @@ void CoreWorkerTest::TestActorFailure(
 
       TaskOptions options{1, resources};
       std::vector<ObjectID> return_ids;
-      RayFunction func{ray::Language::PYTHON, {}};
+      RayFunction func(ray::Language::PYTHON, {});
 
       auto status =
           driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids);
@@ -641,7 +640,7 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) {
   // to benchmark performance.
   uint8_t array[] = {1, 2, 3};
   auto buffer = std::make_shared<LocalMemoryBuffer>(array, sizeof(array));
-  RayFunction function{ray::Language::PYTHON, {}};
+  RayFunction function(ray::Language::PYTHON, {});
   std::vector<TaskArg> args;
   args.emplace_back(TaskArg::PassByValue(std::make_shared<RayObject>(buffer, nullptr)));
 
@@ -649,8 +648,8 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) {
   ActorCreationOptions actor_options{0, /*is_direct_call*/ true, resources, {}};
   const auto job_id = NextJobId();
   ActorHandle actor_handle(ActorID::Of(job_id, TaskID::ForDriverTask(job_id), 1),
-                           ActorHandleID::Nil(), function.language, true,
-                           function.function_descriptor);
+                           ActorHandleID::Nil(), function.GetLanguage(), true,
+                           function.GetFunctionDescriptor());
 
   // Manually create `num_tasks` task specs, and for each of them create a
   // `PushTaskRequest`, this is to batch performance of TaskSpec
@@ -664,8 +663,8 @@ TEST_F(ZeroNodeTest, TestTaskSpecPerf) {
     auto num_returns = options.num_returns;
 
     TaskSpecBuilder builder;
-    builder.SetCommonTaskSpec(RandomTaskId(), function.language,
-                              function.function_descriptor, job_id, RandomTaskId(), 0,
+    builder.SetCommonTaskSpec(RandomTaskId(), function.GetLanguage(),
+                              function.GetFunctionDescriptor(), job_id, RandomTaskId(), 0,
                               num_returns, resources, resources);
     // Set task arguments.
     for (const auto &arg : args) {
@@ -703,7 +702,7 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) {
   // Test creating actor.
   uint8_t array[] = {1, 2, 3};
   auto buffer = std::make_shared<LocalMemoryBuffer>(array, sizeof(array));
-  RayFunction func{ray::Language::PYTHON, {}};
+  RayFunction func(ray::Language::PYTHON, {});
   std::vector<TaskArg> args;
   args.emplace_back(TaskArg::PassByValue(std::make_shared<RayObject>(buffer, nullptr)));
 
@@ -725,7 +724,7 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) {
 
     TaskOptions options{1, resources};
     std::vector<ObjectID> return_ids;
-    RayFunction func{ray::Language::PYTHON, {}};
+    RayFunction func(ray::Language::PYTHON, {});
 
     RAY_CHECK_OK(
         driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids));