Set RAY_FORCE_DIRECT=1 for run_rllib_tests, test_basic (#6171)

src/ray/core_worker/core_worker.cc

@@ -28,11 +28,6 @@ void BuildCommonTaskSpec(
   // Set task arguments.
   for (const auto &arg : args) {
     if (arg.IsPassedByReference()) {
-      // TODO(ekl) remove this check once we deprecate TaskTransportType::RAYLET
-      if (transport_type == ray::TaskTransportType::RAYLET) {
-        RAY_CHECK(!arg.GetReference().IsDirectCallType())
-            << "Passing direct call objects to non-direct tasks is not allowed.";
-      }
       builder.AddByRefArg(arg.GetReference());
     } else {
       builder.AddByValueArg(arg.GetValue());
@@ -169,6 +164,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
         RAY_CHECK_OK(plasma_store_provider_->Put(obj, obj_id));
       },
       ref_counting_enabled ? reference_counter_ : nullptr, raylet_client_));
+  resolver_.reset(new LocalDependencyResolver(memory_store_));
 
   // Create an entry for the driver task in the task table. This task is
   // added immediately with status RUNNING. This allows us to push errors
@@ -621,7 +617,12 @@ Status CoreWorker::CreateActor(const RayFunction &function,
   *return_actor_id = actor_id;
   TaskSpecification task_spec = builder.Build();
   PinObjectReferences(task_spec, TaskTransportType::RAYLET);
-  return raylet_client_->SubmitTask(task_spec);
+  // TODO(ekl) if we moved actor creation to use direct call tasks, then we won't
+  // need to manually resolve direct call args here.
+  resolver_->ResolveDependencies(task_spec, [this, task_spec]() {
+    RAY_CHECK_OK(raylet_client_->SubmitTask(task_spec));
+  });
+  return Status::OK();
 }
 
 Status CoreWorker::SubmitActorTask(const ActorID &actor_id, const RayFunction &function,
@@ -874,7 +875,7 @@ Status CoreWorker::BuildArgsForExecutor(const TaskSpecification &task,
         metadata = std::make_shared<LocalMemoryBuffer>(
             const_cast<uint8_t *>(task.ArgMetadata(i)), task.ArgMetadataSize(i));
       }
-      args->at(i) = std::make_shared<RayObject>(data, metadata);
+      args->at(i) = std::make_shared<RayObject>(data, metadata, /*copy_data*/ true);
       arg_reference_ids->at(i) = ObjectID::Nil();
     }
   }

src/ray/core_worker/core_worker.h

@@ -522,6 +522,9 @@ class CoreWorker {
   /// Map from actor ID to a handle to that actor.
   absl::flat_hash_map<ActorID, std::unique_ptr<ActorHandle>> actor_handles_;
 
+  /// Resolve local and remote dependencies for actor creation.
+  std::unique_ptr<LocalDependencyResolver> resolver_;
+
   ///
   /// Fields related to task execution.
   ///

src/ray/core_worker/transport/dependency_resolver.cc

@@ -42,7 +42,7 @@ void DoInlineObjectValue(const ObjectID &obj_id, std::shared_ptr<RayObject> valu
   RAY_CHECK(found) << "obj id " << obj_id << " not found";
 }
 
-void LocalDependencyResolver::ResolveDependencies(const TaskSpecification &task,
+void LocalDependencyResolver::ResolveDependencies(TaskSpecification &task,
                                                   std::function<void()> on_complete) {
   absl::flat_hash_set<ObjectID> local_dependencies;
   for (size_t i = 0; i < task.NumArgs(); i++) {

src/ray/core_worker/transport/dependency_resolver.h

@@ -22,8 +22,7 @@ class LocalDependencyResolver {
   /// Note: This method **will mutate** the given TaskSpecification.
   ///
   /// Postcondition: all direct call ids in arguments are converted to values.
-  void ResolveDependencies(const TaskSpecification &task,
-                           std::function<void()> on_complete);
+  void ResolveDependencies(TaskSpecification &task, std::function<void()> on_complete);
 
   /// Return the number of tasks pending dependency resolution.
   /// TODO(ekl) this should be exposed in worker stats.

src/ray/core_worker/transport/raylet_transport.cc

@@ -17,7 +17,8 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
     rpc::SendReplyCallback send_reply_callback) {
   const Task task(request.task());
   const auto &task_spec = task.GetTaskSpecification();
-  RAY_LOG(DEBUG) << "Received task " << task_spec.TaskId();
+  RAY_LOG(DEBUG) << "Received task " << task_spec.TaskId() << " is create "
+                 << task_spec.IsActorCreationTask();
 
   // Set the resource IDs for this task.
   // TODO: convert the resource map to protobuf and change this.

src/ray/raylet/raylet_client.cc

@@ -230,6 +230,14 @@ RayletClient::RayletClient(std::shared_ptr<ray::rpc::NodeManagerWorkerClient> gr
 
 ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) {
   ray::rpc::SubmitTaskRequest request;
+  for (size_t i = 0; i < task_spec.NumArgs(); i++) {
+    if (task_spec.ArgByRef(i)) {
+      for (size_t j = 0; j < task_spec.ArgIdCount(i); j++) {
+        RAY_CHECK(!task_spec.ArgId(i, j).IsDirectCallType())
+            << "Passing direct call objects to non-direct tasks is not allowed.";
+      }
+    }
+  }
   request.mutable_task_spec()->CopyFrom(task_spec.GetMessage());
   return grpc_client_->SubmitTask(request, /*callback=*/nullptr);
 }

src/ray/rpc/client_call.h

@@ -198,8 +198,12 @@ class ClientCallManager {
       auto status = cqs_[index].AsyncNext(&got_tag, &ok, deadline);
       if (status == grpc::CompletionQueue::SHUTDOWN) {
         break;
-      }
-      if (status != grpc::CompletionQueue::TIMEOUT) {
+      } else if (status == grpc::CompletionQueue::TIMEOUT && shutdown_) {
+        // If we timed out and shutdown, then exit immediately. This should not
+        // be needed, but gRPC seems to not return SHUTDOWN correctly in these
+        // cases (e.g., test_wait will hang on shutdown without this check).
+        break;
+      } else if (status != grpc::CompletionQueue::TIMEOUT) {
         auto tag = reinterpret_cast<ClientCallTag *>(got_tag);
         if (ok && !main_service_.stopped() && !shutdown_) {
           // Post the callback to the main event loop.