Support pass by ref args in for direct actor calls (#6040)

BUILD.bazel

@@ -213,6 +213,7 @@ cc_library(
         ":ray_common",
         ":worker_cc_grpc",
         "@boost//:asio",
+        "@boost//:thread",
         "@com_github_grpc_grpc//:grpc++",
     ],
 )

python/ray/ray_perf.py

@@ -1,27 +1,46 @@
 """This is the script for `ray microbenchmark`."""
 
+import os
 import time
 import numpy as np
 import multiprocessing
 import ray
 
+# Only run tests matching this filter pattern.
+filter_pattern = os.environ.get("TESTS_TO_RUN", "")
+
 
 @ray.remote
 class Actor(object):
     def small_value(self):
         return 0
 
+    def small_value_arg(self, x):
+        return 0
+
     def small_value_batch(self, n):
         ray.get([small_value.remote() for _ in range(n)])
 
 
 @ray.remote
 class Client(object):
-    def __init__(self, server):
-        self.server = server
+    def __init__(self, servers):
+        if not isinstance(servers, list):
+            servers = [servers]
+        self.servers = servers
 
     def small_value_batch(self, n):
-        ray.get([self.server.small_value.remote() for _ in range(n)])
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value.remote() for _ in range(n)])
+        ray.get(results)
+
+    def small_value_batch_arg(self, n):
+        x = ray.put(0)
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value_arg.remote(x) for _ in range(n)])
+        ray.get(results)
 
 
 @ray.remote
@@ -37,6 +56,8 @@ def small_value_batch(n):
 
 
 def timeit(name, fn, multiplier=1):
+    if filter_pattern not in name:
+        return
     # warmup
     start = time.time()
     while time.time() - start < 1:
@@ -56,6 +77,7 @@ def timeit(name, fn, multiplier=1):
 
 
 def main():
+    print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks")
     ray.init()
     value = ray.put(0)
     arr = np.zeros(100 * 1024 * 1024, dtype=np.int64)
@@ -141,23 +163,17 @@ def main():
 
     timeit("multi client actor calls async", actor_multi2, m * n)
 
-    a = Actor._remote(is_direct_call=True)
-
-    def actor_sync_direct():
-        ray.get(a.small_value.remote())
-
-    timeit("single client direct actor calls sync", actor_sync_direct)
-
-    a = Actor._remote(is_direct_call=True)
-
-    def actor_async_direct():
-        ray.get([a.small_value.remote() for _ in range(1000)])
-
-    timeit("single client direct actor calls async", actor_async_direct, 1000)
-
     n = 5000
     n_cpu = multiprocessing.cpu_count() // 2
     actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)]
+    client = Client.remote(actors)
+
+    def actor_async_direct():
+        ray.get(client.small_value_batch.remote(n))
+
+    timeit("single client direct actor calls async", actor_async_direct,
+           n * len(actors))
+
     clients = [Client.remote(a) for a in actors]
 
     def actor_multi2_direct():
@@ -166,6 +182,16 @@ def main():
     timeit("multi client direct actor calls async", actor_multi2_direct,
            n * len(clients))
 
+    n = 1000
+    actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)]
+    clients = [Client.remote(a) for a in actors]
+
+    def actor_multi2_direct_arg():
+        ray.get([c.small_value_batch_arg.remote(n) for c in clients])
+
+    timeit("multi client direct actor calls with arg async",
+           actor_multi2_direct_arg, n * len(clients))
+
 
 if __name__ == "__main__":
     main()

python/ray/tests/test_basic.py

@@ -1223,9 +1223,70 @@ def test_direct_actor_errors(ray_start_regular):
     with pytest.raises(Exception):
         ray.get(f.remote([a.f.remote(2)]))
 
-    # by ref args not implemented
-    with pytest.raises(ray.exceptions.RayletError):
-        a.f.remote(f.remote(2))
+
+def test_direct_actor_pass_by_ref(ray_start_regular):
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            return x * 2
+
+    @ray.remote
+    def f(x):
+        return x
+
+    @ray.remote
+    def error():
+        sys.exit(0)
+
+    a = Actor._remote(is_direct_call=True)
+    assert ray.get(a.f.remote(f.remote(1))) == 2
+
+    fut = [a.f.remote(f.remote(i)) for i in range(100)]
+    assert ray.get(fut) == [i * 2 for i in range(100)]
+
+    # propagates errors for pass by ref
+    with pytest.raises(Exception):
+        ray.get(a.f.remote(error.remote()))
+
+
+def test_direct_actor_pass_by_ref_order_optimization(shutdown_only):
+    ray.init(num_cpus=4)
+
+    @ray.remote
+    class Actor(object):
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            pass
+
+    a = Actor._remote(is_direct_call=True)
+
+    @ray.remote
+    def fast_value():
+        print("fast value")
+        pass
+
+    @ray.remote
+    def slow_value():
+        print("start sleep")
+        time.sleep(30)
+
+    @ray.remote
+    def runner(f):
+        print("runner", a, f)
+        return ray.get(a.f.remote(f.remote()))
+
+    runner.remote(slow_value)
+    time.sleep(1)
+    x2 = runner.remote(fast_value)
+    start = time.time()
+    ray.get(x2)
+    delta = time.time() - start
+    assert delta < 10, "did not skip slow value"
 
 
 def test_direct_actor_recursive(ray_start_regular):

src/ray/core_worker/core_worker.cc

@@ -131,6 +131,10 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
       raylet_socket, WorkerID::FromBinary(worker_context_.GetWorkerID().Binary()),
       (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(),
       language_, worker_server_.GetPort()));
+  // Unfortunately the raylet client has to be constructed after the receivers.
+  if (direct_actor_task_receiver_ != nullptr) {
+    direct_actor_task_receiver_->Init(*raylet_client_);
+  }
 
   // Set timer to periodically send heartbeats containing active object IDs to the raylet.
   // If the heartbeat timeout is < 0, the heartbeats are disabled.

src/ray/core_worker/test/core_worker_test.cc

@@ -316,12 +316,6 @@ void CoreWorkerTest::TestActorTask(std::unordered_map<std::string, double> &reso
     std::vector<ObjectID> return_ids;
     RayFunction func(ray::Language::PYTHON, {});
     auto status = driver.SubmitActorTask(actor_id, func, args, options, &return_ids);
-    if (is_direct_call) {
-      // For direct actor call, submitting a task with by-reference arguments
-      // would fail.
-      ASSERT_TRUE(!status.ok());
-      return;
-    }
     ASSERT_TRUE(status.ok());
 
     ASSERT_EQ(return_ids.size(), 1);

src/ray/core_worker/test/scheduling_queue_test.cc

@@ -5,9 +5,25 @@
 
 namespace ray {
 
+class MockWaiter : public DependencyWaiter {
+ public:
+  MockWaiter() {}
+
+  void Wait(const std::vector<ObjectID> &dependencies,
+            std::function<void()> on_dependencies_available) override {
+    callbacks_.push_back([on_dependencies_available]() { on_dependencies_available(); });
+  }
+
+  void Complete(int index) { callbacks_[index](); }
+
+ private:
+  std::vector<std::function<void()>> callbacks_;
+};
+
 TEST(SchedulingQueueTest, TestInOrder) {
   boost::asio::io_service io_service;
-  SchedulingQueue queue(io_service, 0);
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -21,9 +37,55 @@ TEST(SchedulingQueueTest, TestInOrder) {
   ASSERT_EQ(n_rej, 0);
 }
 
+TEST(SchedulingQueueTest, TestWaitForObjects) {
+  ObjectID obj1 = ObjectID::FromRandom();
+  ObjectID obj2 = ObjectID::FromRandom();
+  ObjectID obj3 = ObjectID::FromRandom();
+  boost::asio::io_service io_service;
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(0, -1, fn_ok, fn_rej);
+  queue.Add(1, -1, fn_ok, fn_rej, {obj1});
+  queue.Add(2, -1, fn_ok, fn_rej, {obj2});
+  queue.Add(3, -1, fn_ok, fn_rej, {obj3});
+  ASSERT_EQ(n_ok, 1);
+
+  waiter.Complete(0);
+  ASSERT_EQ(n_ok, 2);
+
+  waiter.Complete(2);
+  ASSERT_EQ(n_ok, 2);
+
+  waiter.Complete(1);
+  ASSERT_EQ(n_ok, 4);
+}
+
+TEST(SchedulingQueueTest, TestWaitForObjectsNotSubjectToSeqTimeout) {
+  ObjectID obj1 = ObjectID::FromRandom();
+  boost::asio::io_service io_service;
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
+  int n_ok = 0;
+  int n_rej = 0;
+  auto fn_ok = [&n_ok]() { n_ok++; };
+  auto fn_rej = [&n_rej]() { n_rej++; };
+  queue.Add(0, -1, fn_ok, fn_rej);
+  queue.Add(1, -1, fn_ok, fn_rej, {obj1});
+  ASSERT_EQ(n_ok, 1);
+  io_service.run();
+  ASSERT_EQ(n_rej, 0);
+  waiter.Complete(0);
+  ASSERT_EQ(n_ok, 2);
+}
+
 TEST(SchedulingQueueTest, TestOutOfOrder) {
   boost::asio::io_service io_service;
-  SchedulingQueue queue(io_service, 0);
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -37,9 +99,10 @@ TEST(SchedulingQueueTest, TestOutOfOrder) {
   ASSERT_EQ(n_rej, 0);
 }
 
-TEST(SchedulingQueueTest, TestDepWaitTimeout) {
+TEST(SchedulingQueueTest, TestSeqWaitTimeout) {
   boost::asio::io_service io_service;
-  SchedulingQueue queue(io_service, 0);
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };
@@ -60,7 +123,8 @@ TEST(SchedulingQueueTest, TestDepWaitTimeout) {
 
 TEST(SchedulingQueueTest, TestSkipAlreadyProcessedByClient) {
   boost::asio::io_service io_service;
-  SchedulingQueue queue(io_service, 0);
+  MockWaiter waiter;
+  SchedulingQueue queue(io_service, waiter, 0);
   int n_ok = 0;
   int n_rej = 0;
   auto fn_ok = [&n_ok]() { n_ok++; };

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -24,9 +24,6 @@ CoreWorkerDirectActorTaskSubmitter::CoreWorkerDirectActorTaskSubmitter(
 Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(
     const TaskSpecification &task_spec) {
   RAY_LOG(DEBUG) << "Submitting task " << task_spec.TaskId();
-  if (HasByReferenceArgs(task_spec)) {
-    return Status::Invalid("Direct actor call only supports by-value arguments");
-  }
 
   RAY_CHECK(task_spec.IsActorTask());
   const auto &actor_id = task_spec.ActorId();
@@ -84,7 +81,7 @@ void CoreWorkerDirectActorTaskSubmitter::HandleActorUpdate(
     // Check if this actor is the one that we're interested, if we already have
     // a connection to the actor, or have pending requests for it, we should
     // create a new connection.
-    if (pending_requests_.count(actor_id) > 0) {
+    if (pending_requests_.count(actor_id) > 0 && rpc_clients_.count(actor_id) == 0) {
       ConnectAndSendPendingTasks(actor_id, actor_data.ip_address(), actor_data.port());
     }
   } else {
@@ -147,6 +144,9 @@ void CoreWorkerDirectActorTaskSubmitter::PushTask(
           waiting_reply_tasks_[actor_id].erase(task_id);
         }
         if (!status.ok()) {
+          // Note that this might be the __ray_terminate__ task, so we don't log
+          // loudly with ERROR here.
+          RAY_LOG(DEBUG) << "Task failed with error: " << status;
           TreatTaskAsFailed(task_id, num_returns, rpc::ErrorType::ACTOR_DIED);
           return;
         }
@@ -208,28 +208,38 @@ CoreWorkerDirectActorTaskReceiver::CoreWorkerDirectActorTaskReceiver(
   server.RegisterService(task_service_);
 }
 
+void CoreWorkerDirectActorTaskReceiver::Init(RayletClient &raylet_client) {
+  waiter_.reset(new DependencyWaiterImpl(raylet_client));
+}
+
 void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
     const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply,
     rpc::SendReplyCallback send_reply_callback) {
+  RAY_CHECK(waiter_ != nullptr) << "Must call init() prior to use";
   const TaskSpecification task_spec(request.task_spec());
   RAY_LOG(DEBUG) << "Received task " << task_spec.TaskId();
-  if (HasByReferenceArgs(task_spec)) {
-    send_reply_callback(
-        Status::Invalid("Direct actor call only supports by value arguments"), nullptr,
-        nullptr);
-    return;
-  }
   if (task_spec.IsActorTask() && !worker_context_.CurrentActorUseDirectCall()) {
     send_reply_callback(Status::Invalid("This actor doesn't accept direct calls."),
                         nullptr, nullptr);
     return;
   }
 
+  // TODO(ekl) resolving object dependencies is expensive and requires an IPC to
+  // the raylet, which is a central bottleneck. In the future, we should inline
+  // dependencies that are small and already known to be local to the client.
+  std::vector<ObjectID> dependencies;
+  for (size_t i = 0; i < task_spec.NumArgs(); ++i) {
+    int count = task_spec.ArgIdCount(i);
+    for (int j = 0; j < count; j++) {
+      dependencies.push_back(task_spec.ArgId(i, j));
+    }
+  }
+
   auto it = scheduling_queue_.find(task_spec.CallerId());
   if (it == scheduling_queue_.end()) {
     auto result = scheduling_queue_.emplace(
-        task_spec.CallerId(),
-        std::unique_ptr<SchedulingQueue>(new SchedulingQueue(task_main_io_service_)));
+        task_spec.CallerId(), std::unique_ptr<SchedulingQueue>(
+                                  new SchedulingQueue(task_main_io_service_, *waiter_)));
     it = result.first;
   }
   it->second->Add(
@@ -268,7 +278,17 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
       },
       [send_reply_callback]() {
         send_reply_callback(Status::Invalid("client cancelled rpc"), nullptr, nullptr);
-      });
+      },
+      dependencies);
+}
+
+void CoreWorkerDirectActorTaskReceiver::HandleDirectActorCallArgWaitComplete(
+    const rpc::DirectActorCallArgWaitCompleteRequest &request,
+    rpc::DirectActorCallArgWaitCompleteReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
+  RAY_LOG(DEBUG) << "Arg wait complete for tag " << request.tag();
+  waiter_->OnWaitComplete(request.tag());
+  send_reply_callback(Status::OK(), nullptr, nullptr);
 }
 
 }  // namespace ray

src/ray/core_worker/transport/direct_actor_transport.h

@@ -1,6 +1,7 @@
 #ifndef RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H
 #define RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H
 
+#include <boost/thread.hpp>
 #include <list>
 #include <set>
 #include <utility>
@@ -125,60 +126,139 @@ class CoreWorkerDirectActorTaskSubmitter {
   friend class CoreWorkerTest;
 };
 
+/// Object dependency and RPC state of an inbound request.
+class InboundRequest {
+ public:
+  InboundRequest(){};
+  InboundRequest(std::function<void()> accept_callback,
+                 std::function<void()> reject_callback, bool has_dependencies)
+      : accept_callback_(accept_callback),
+        reject_callback_(reject_callback),
+        has_pending_dependencies_(has_dependencies) {}
+
+  void Accept() { accept_callback_(); }
+  void Cancel() { reject_callback_(); }
+  bool CanExecute() const { return !has_pending_dependencies_; }
+  void MarkDependenciesSatisfied() { has_pending_dependencies_ = false; }
+
+ private:
+  std::function<void()> accept_callback_;
+  std::function<void()> reject_callback_;
+  bool has_pending_dependencies_;
+};
+
+/// Waits for an object dependency to become available. Abstract for testing.
+class DependencyWaiter {
+ public:
+  /// Calls `callback` once the specified objects become available.
+  virtual void Wait(const std::vector<ObjectID> &dependencies,
+                    std::function<void()> on_dependencies_available) = 0;
+};
+
+class DependencyWaiterImpl : public DependencyWaiter {
+ public:
+  DependencyWaiterImpl(RayletClient &raylet_client) : raylet_client_(raylet_client) {}
+
+  void Wait(const std::vector<ObjectID> &dependencies,
+            std::function<void()> on_dependencies_available) override {
+    auto tag = next_request_id_++;
+    requests_[tag] = on_dependencies_available;
+    raylet_client_.WaitForDirectActorCallArgs(dependencies, tag);
+  }
+
+  /// Fulfills the callback stored by Wait().
+  void OnWaitComplete(int64_t tag) {
+    auto it = requests_.find(tag);
+    RAY_CHECK(it != requests_.end());
+    it->second();
+    requests_.erase(it);
+  }
+
+ private:
+  int64_t next_request_id_ = 0;
+  std::unordered_map<int64_t, std::function<void()>> requests_;
+  RayletClient &raylet_client_;
+};
+
 /// Used to ensure serial order of task execution per actor handle.
 /// See direct_actor.proto for a description of the ordering protocol.
 class SchedulingQueue {
  public:
-  SchedulingQueue(boost::asio::io_service &io_service,
+  SchedulingQueue(boost::asio::io_service &main_io_service, DependencyWaiter &waiter,
                   int64_t reorder_wait_seconds = kMaxReorderWaitSeconds)
-      : wait_timer_(io_service), reorder_wait_seconds_(reorder_wait_seconds) {}
+      : wait_timer_(main_io_service),
+        waiter_(waiter),
+        reorder_wait_seconds_(reorder_wait_seconds),
+        main_thread_id_(boost::this_thread::get_id()) {}
 
   void Add(int64_t seq_no, int64_t client_processed_up_to,
-           std::function<void()> accept_request, std::function<void()> reject_request) {
+           std::function<void()> accept_request, std::function<void()> reject_request,
+           const std::vector<ObjectID> &dependencies = {}) {
+    RAY_CHECK(boost::this_thread::get_id() == main_thread_id_);
     if (client_processed_up_to >= next_seq_no_) {
-      RAY_LOG(DEBUG) << "client skipping requests " << next_seq_no_ << " to "
+      RAY_LOG(ERROR) << "client skipping requests " << next_seq_no_ << " to "
                      << client_processed_up_to;
       next_seq_no_ = client_processed_up_to + 1;
     }
-    pending_tasks_[seq_no] = make_pair(accept_request, reject_request);
+    pending_tasks_[seq_no] =
+        InboundRequest(accept_request, reject_request, dependencies.size() > 0);
+    if (dependencies.size() > 0) {
+      waiter_.Wait(dependencies, [seq_no, this]() {
+        RAY_CHECK(boost::this_thread::get_id() == main_thread_id_);
+        auto it = pending_tasks_.find(seq_no);
+        if (it != pending_tasks_.end()) {
+          it->second.MarkDependenciesSatisfied();
+          ScheduleRequests();
+        }
+      });
+    }
+    ScheduleRequests();
+  }
 
-    // Reject any stale requests that the client doesn't need any longer.
+ private:
+  /// Schedules as many requests as possible in sequence.
+  void ScheduleRequests() {
+    // Cancel any stale requests that the client doesn't need any longer.
     while (!pending_tasks_.empty() && pending_tasks_.begin()->first < next_seq_no_) {
       auto head = pending_tasks_.begin();
-      head->second.second();  // reject_request
+      head->second.Cancel();
       pending_tasks_.erase(head);
     }
 
     // Process as many in-order requests as we can.
-    while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_) {
+    while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_ &&
+           pending_tasks_.begin()->second.CanExecute()) {
       auto head = pending_tasks_.begin();
-      head->second.first();  // accept_request
+      head->second.Accept();
       pending_tasks_.erase(head);
       next_seq_no_++;
     }
 
-    // Set a timeout on the queued tasks to avoid an infinite wait on failure.
-    wait_timer_.expires_from_now(boost::posix_time::seconds(reorder_wait_seconds_));
-    if (!pending_tasks_.empty()) {
+    if (pending_tasks_.empty() || !pending_tasks_.begin()->second.CanExecute()) {
+      // No timeout for object dependency waits.
+      wait_timer_.cancel();
+    } else {
+      // Set a timeout on the queued tasks to avoid an infinite wait on failure.
+      wait_timer_.expires_from_now(boost::posix_time::seconds(reorder_wait_seconds_));
       RAY_LOG(DEBUG) << "waiting for " << next_seq_no_ << " queue size "
                      << pending_tasks_.size();
       wait_timer_.async_wait([this](const boost::system::error_code &error) {
         if (error == boost::asio::error::operation_aborted) {
           return;  // time deadline was adjusted
         }
-        OnDependencyWaitTimeout();
+        OnSequencingWaitTimeout();
       });
     }
   }
 
- private:
-  /// Called when we time out waiting for a task dependency to show up.
-  void OnDependencyWaitTimeout() {
+  /// Called when we time out waiting for an earlier task to show up.
+  void OnSequencingWaitTimeout() {
+    RAY_CHECK(boost::this_thread::get_id() == main_thread_id_);
     RAY_LOG(ERROR) << "timed out waiting for " << next_seq_no_
                    << ", cancelling all queued tasks";
     while (!pending_tasks_.empty()) {
       auto head = pending_tasks_.begin();
-      head->second.second();  // reject_request
+      head->second.Cancel();
       pending_tasks_.erase(head);
       next_seq_no_ = std::max(next_seq_no_, head->first + 1);
     }
@@ -187,12 +267,15 @@ class SchedulingQueue {
   /// Max time in seconds to wait for dependencies to show up.
   const int64_t reorder_wait_seconds_ = 0;
   /// Sorted map of (accept, rej) task callbacks keyed by their sequence number.
-  std::map<int64_t, std::pair<std::function<void()>, std::function<void()>>>
-      pending_tasks_;
+  std::map<int64_t, InboundRequest> pending_tasks_;
   /// The next sequence number we are waiting for to arrive.
   int64_t next_seq_no_ = 0;
   /// Timer for waiting on dependencies.
   boost::asio::deadline_timer wait_timer_;
+  /// The id of the thread that constructed this scheduling queue.
+  boost::thread::id main_thread_id_;
+  /// Reference to the waiter owned by the task receiver.
+  DependencyWaiter &waiter_;
 
   friend class SchedulingQueueTest;
 };
@@ -208,16 +291,27 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler {
                                     rpc::GrpcServer &server,
                                     const TaskHandler &task_handler);
 
+  /// Initialize this receiver. This must be called prior to use.
+  void Init(RayletClient &client);
+
   /// Handle a `PushTask` request.
-  /// The implementation can handle this request asynchronously. When hanling is done, the
-  /// `done_callback` should be called.
   ///
   /// \param[in] request The request message.
   /// \param[out] reply The reply message.
-  /// \param[in] done_callback The callback to be called when the request is done.
+  /// \param[in] send_reply_callback The callback to be called when the request is done.
   void HandlePushTask(const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply,
                       rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle a `DirectActorCallArgWaitComplete` request.
+  ///
+  /// \param[in] request The request message.
+  /// \param[out] reply The reply message.
+  /// \param[in] send_reply_callback The callback to be called when the request is done.
+  void HandleDirectActorCallArgWaitComplete(
+      const rpc::DirectActorCallArgWaitCompleteRequest &request,
+      rpc::DirectActorCallArgWaitCompleteReply *reply,
+      rpc::SendReplyCallback send_reply_callback) override;
+
  private:
   // Worker context.
   WorkerContext &worker_context_;
@@ -227,6 +321,8 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler {
   TaskHandler task_handler_;
   /// The IO event loop for running tasks on.
   boost::asio::io_service &task_main_io_service_;
+  /// Shared waiter for dependencies required by incoming tasks.
+  std::unique_ptr<DependencyWaiterImpl> waiter_;
   /// Queue of pending requests per actor handle.
   /// TODO(ekl) GC these queues once the handle is no longer active.
   std::unordered_map<TaskID, std::unique_ptr<SchedulingQueue>> scheduling_queue_;

src/ray/core_worker/transport/raylet_transport.h

@@ -22,8 +22,8 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler {
                                rpc::GrpcServer &server, const TaskHandler &task_handler);
 
   /// Handle a `AssignTask` request.
-  /// The implementation can handle this request asynchronously. When hanling is done, the
-  /// `send_reply_callback` should be called.
+  /// The implementation can handle this request asynchronously. When handling is done,
+  /// the `send_reply_callback` should be called.
   ///
   /// \param[in] request The request message.
   /// \param[out] reply The reply message.
@@ -41,6 +41,8 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler {
   rpc::WorkerTaskGrpcService task_service_;
   /// The callback function to process a task.
   TaskHandler task_handler_;
+  /// The callback to process arg wait complete.
+  std::function<void(int64_t)> on_wait_complete_;
 };
 
 }  // namespace ray

src/ray/protobuf/direct_actor.proto

@@ -33,8 +33,20 @@ message PushTaskReply {
   repeated ReturnObject return_objects = 1;
 }
 
+message DirectActorCallArgWaitCompleteRequest {
+  // Id used to uniquely identify this request. This is sent back to the core
+  // worker to notify the wait has completed.
+  int64 tag = 1;
+}
+
+message DirectActorCallArgWaitCompleteReply {
+}
+
 // Service for direct actor.
 service DirectActorService {
   // Push a task to a worker.
   rpc PushTask(PushTaskRequest) returns (PushTaskReply);
+  // Notify wait for direct actor call args has completed
+  rpc DirectActorCallArgWaitComplete(DirectActorCallArgWaitCompleteRequest)
+      returns (DirectActorCallArgWaitCompleteReply);
 }

src/ray/raylet/format/node_manager.fbs

@@ -55,6 +55,8 @@ enum MessageType:int {
   // The response message to WaitRequest; replies with the objects found and objects
   // remaining.
   WaitReply,
+  // Wait for objects asynchronously. The reply will be sent back via gRPC push.
+  WaitForDirectActorCallArgsRequest,
   // Push an error to the relevant driver. This is sent from a worker to the
   // node manager.
   PushErrorRequest,
@@ -187,6 +189,14 @@ table WaitReply {
   remaining: [string];
 }
 
+table WaitForDirectActorCallArgsRequest {
+  // List of object ids we'll be waiting on.
+  object_ids: [string];
+  // Id used to uniquely identify this request. This is sent back to the core
+  // worker to notify the wait has completed.
+  tag: int;
+}
+
 // This struct is the same as ErrorTableData.
 table PushErrorRequest {
   // The ID of the job that the error is for.

src/ray/raylet/node_manager.cc

@@ -898,6 +898,9 @@ void NodeManager::ProcessClientMessage(
   case protocol::MessageType::WaitRequest: {
     ProcessWaitRequestMessage(client, message_data);
   } break;
+  case protocol::MessageType::WaitForDirectActorCallArgsRequest: {
+    ProcessWaitForDirectActorCallArgsRequestMessage(client, message_data);
+  } break;
   case protocol::MessageType::PushErrorRequest: {
     ProcessPushErrorRequestMessage(message_data);
   } break;
@@ -1253,6 +1256,37 @@ void NodeManager::ProcessWaitRequestMessage(
   RAY_CHECK_OK(status);
 }
 
+void NodeManager::ProcessWaitForDirectActorCallArgsRequestMessage(
+    const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
+  // Read the data.
+  auto message =
+      flatbuffers::GetRoot<protocol::WaitForDirectActorCallArgsRequest>(message_data);
+  int64_t tag = message->tag();
+  std::vector<ObjectID> object_ids = from_flatbuf<ObjectID>(*message->object_ids());
+  std::vector<ObjectID> required_object_ids;
+  for (auto const &object_id : object_ids) {
+    if (!task_dependency_manager_.CheckObjectLocal(object_id)) {
+      // Add any missing objects to the list to subscribe to in the task
+      // dependency manager. These objects will be pulled from remote node
+      // managers and reconstructed if necessary.
+      required_object_ids.push_back(object_id);
+    }
+  }
+
+  ray::Status status = object_manager_.Wait(
+      object_ids, -1, object_ids.size(), false,
+      [this, client, tag](std::vector<ObjectID> found, std::vector<ObjectID> remaining) {
+        RAY_CHECK(remaining.empty());
+        std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
+        if (worker == nullptr) {
+          RAY_LOG(ERROR) << "Lost worker for wait request " << client;
+        } else {
+          worker->DirectActorCallArgWaitComplete(tag);
+        }
+      });
+  RAY_CHECK_OK(status);
+}
+
 void NodeManager::ProcessPushErrorRequestMessage(const uint8_t *message_data) {
   auto message = flatbuffers::GetRoot<protocol::PushErrorRequest>(message_data);
 

src/ray/raylet/node_manager.h

@@ -428,6 +428,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
   void ProcessWaitRequestMessage(const std::shared_ptr<LocalClientConnection> &client,
                                  const uint8_t *message_data);
 
+  /// Process client message of WaitForDirectActorCallArgsRequest
+  ///
+  /// \param client The client that sent the message.
+  /// \param message_data A pointer to the message data.
+  /// \return Void.
+  void ProcessWaitForDirectActorCallArgsRequestMessage(
+      const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data);
+
   /// Process client message of PushErrorRequest
   ///
   /// \param message_data A pointer to the message data.

src/ray/raylet/raylet_client.cc

@@ -279,6 +279,15 @@ ray::Status RayletClient::Wait(const std::vector<ObjectID> &object_ids, int num_
   return ray::Status::OK();
 }
 
+ray::Status RayletClient::WaitForDirectActorCallArgs(
+    const std::vector<ObjectID> &object_ids, int64_t tag) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = ray::protocol::CreateWaitForDirectActorCallArgsRequest(
+      fbb, to_flatbuf(fbb, object_ids), tag);
+  fbb.Finish(message);
+  return conn_->WriteMessage(MessageType::WaitForDirectActorCallArgsRequest, &fbb);
+}
+
 ray::Status RayletClient::PushError(const ray::JobID &job_id, const std::string &type,
                                     const std::string &error_message, double timestamp) {
   flatbuffers::FlatBufferBuilder fbb;

src/ray/raylet/raylet_client.h

@@ -118,6 +118,15 @@ class RayletClient {
                    int64_t timeout_milliseconds, bool wait_local,
                    const TaskID &current_task_id, WaitResultPair *result);
 
+  /// Wait for the given objects, asynchronously. The core worker is notified when
+  /// the wait completes.
+  ///
+  /// \param object_ids The objects to wait for.
+  /// \param tag Value that will be sent to the core worker via gRPC on completion.
+  /// \return ray::Status.
+  ray::Status WaitForDirectActorCallArgs(const std::vector<ObjectID> &object_ids,
+                                         int64_t tag);
+
   /// Push an error to the relevant driver.
   ///
   /// \param The ID of the job_id that the error is for.

src/ray/raylet/worker.cc

@@ -4,6 +4,8 @@
 
 #include "ray/raylet/format/node_manager_generated.h"
 #include "ray/raylet/raylet.h"
+#include "src/ray/protobuf/direct_actor.grpc.pb.h"
+#include "src/ray/protobuf/direct_actor.pb.h"
 
 namespace ray {
 
@@ -24,6 +26,8 @@ Worker::Worker(const WorkerID &worker_id, pid_t pid, const Language &language, i
   if (port_ > 0) {
     rpc_client_ = std::unique_ptr<rpc::WorkerTaskClient>(
         new rpc::WorkerTaskClient("127.0.0.1", port_, client_call_manager_));
+    direct_rpc_client_ = std::unique_ptr<rpc::DirectActorClient>(
+        new rpc::DirectActorClient("127.0.0.1", port_, client_call_manager_));
   }
 }
 
@@ -150,6 +154,21 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
   }
 }
 
+void Worker::DirectActorCallArgWaitComplete(int64_t tag) {
+  RAY_CHECK(port_ > 0);
+  rpc::DirectActorCallArgWaitCompleteRequest request;
+  request.set_tag(tag);
+  auto status = direct_rpc_client_->DirectActorCallArgWaitComplete(
+      request, [](Status status, const rpc::DirectActorCallArgWaitCompleteReply &reply) {
+        if (!status.ok()) {
+          RAY_LOG(ERROR) << "Failed to send wait complete: " << status.ToString();
+        }
+      });
+  if (!status.ok()) {
+    RAY_LOG(ERROR) << "Failed to send wait complete: " << status.ToString();
+  }
+}
+
 }  // namespace raylet
 
 }  // end namespace ray

src/ray/raylet/worker.h

@@ -8,6 +8,7 @@
 #include "ray/common/task/scheduling_resources.h"
 #include "ray/common/task/task.h"
 #include "ray/common/task/task_common.h"
+#include "ray/rpc/worker/direct_actor_client.h"
 #include "ray/rpc/worker/worker_client.h"
 
 namespace ray {
@@ -62,6 +63,7 @@ class Worker {
 
   void AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
                   const std::function<void(Status)> finish_assign_callback);
+  void DirectActorCallArgWaitComplete(int64_t tag);
 
  private:
   /// The worker's ID.
@@ -100,6 +102,8 @@ class Worker {
   rpc::ClientCallManager &client_call_manager_;
   /// The rpc client to send tasks to this worker.
   std::unique_ptr<rpc::WorkerTaskClient> rpc_client_;
+  /// The rpc client to send tasks to the direct actor service.
+  std::unique_ptr<rpc::DirectActorClient> direct_rpc_client_;
 };
 
 }  // namespace raylet

src/ray/rpc/worker/direct_actor_client.h

@@ -34,6 +34,19 @@ class DirectActorClient : public std::enable_shared_from_this<DirectActorClient>
     return std::shared_ptr<DirectActorClient>(instance);
   }
 
+  /// Constructor.
+  ///
+  /// \param[in] address Address of the direct actor server.
+  /// \param[in] port Port of the direct actor server.
+  /// \param[in] client_call_manager The `ClientCallManager` used for managing requests.
+  DirectActorClient(const std::string &address, const int port,
+                    ClientCallManager &client_call_manager)
+      : client_call_manager_(client_call_manager) {
+    std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
+        address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
+    stub_ = DirectActorService::NewStub(channel);
+  };
+
   /// Push a task.
   ///
   /// \param[in] request The request message.
@@ -44,12 +57,33 @@ class DirectActorClient : public std::enable_shared_from_this<DirectActorClient>
     request->set_sequence_number(request->task_spec().actor_task_spec().actor_counter());
     {
       std::lock_guard<std::mutex> lock(mutex_);
+      if (request->task_spec().caller_id() != cur_caller_id_) {
+        // We are running a new task, reset the seq no counter.
+        max_finished_seq_no_ = -1;
+        cur_caller_id_ = request->task_spec().caller_id();
+      }
       send_queue_.push_back(std::make_pair(std::move(request), callback));
     }
     SendRequests();
     return ray::Status::OK();
   }
 
+  /// Notify a wait has completed for direct actor call arguments.
+  ///
+  /// \param[in] request The request message.
+  /// \param[in] callback The callback function that handles reply.
+  /// \return if the rpc call succeeds
+  ray::Status DirectActorCallArgWaitComplete(
+      const DirectActorCallArgWaitCompleteRequest &request,
+      const ClientCallback<DirectActorCallArgWaitCompleteReply> &callback) {
+    auto call = client_call_manager_.CreateCall<DirectActorService,
+                                                DirectActorCallArgWaitCompleteRequest,
+                                                DirectActorCallArgWaitCompleteReply>(
+        *stub_, &DirectActorService::Stub::PrepareAsyncDirectActorCallArgWaitComplete,
+        request, callback);
+    return call->GetStatus();
+  }
+
   /// Send as many pending tasks as possible. This method is thread-safe.
   ///
   /// The client will guarantee no more than kMaxBytesInFlight bytes of RPCs are being
@@ -93,19 +127,6 @@ class DirectActorClient : public std::enable_shared_from_this<DirectActorClient>
   }
 
  private:
-  /// Constructor.
-  ///
-  /// \param[in] address Address of the direct actor server.
-  /// \param[in] port Port of the direct actor server.
-  /// \param[in] client_call_manager The `ClientCallManager` used for managing requests.
-  DirectActorClient(const std::string &address, const int port,
-                    ClientCallManager &client_call_manager)
-      : client_call_manager_(client_call_manager) {
-    std::shared_ptr<grpc::Channel> channel = grpc::CreateChannel(
-        address + ":" + std::to_string(port), grpc::InsecureChannelCredentials());
-    stub_ = DirectActorService::NewStub(channel);
-  };
-
   /// Protects against unsafe concurrent access from the callback thread.
   std::mutex mutex_;
 
@@ -124,6 +145,10 @@ class DirectActorClient : public std::enable_shared_from_this<DirectActorClient>
 
   /// The max sequence number we have processed responses for.
   int64_t max_finished_seq_no_ GUARDED_BY(mutex_) = -1;
+
+  /// The task id we are currently sending requests for. When this changes,
+  /// the max finished seq no counter is reset.
+  std::string cur_caller_id_;
 };
 
 }  // namespace rpc

src/ray/rpc/worker/direct_actor_server.h

@@ -22,6 +22,16 @@ class DirectActorHandler {
   /// \param[in] done_callback The callback to be called when the request is done.
   virtual void HandlePushTask(const PushTaskRequest &request, PushTaskReply *reply,
                               SendReplyCallback send_reply_callback) = 0;
+
+  /// Handle a wait reply for direct actor call arg dependencies.
+  ///
+  /// \param[in] request The request message.
+  /// \param[out] reply The reply message.
+  /// \param[in] send_replay_callback The callback to be called when the request is done.
+  virtual void HandleDirectActorCallArgWaitComplete(
+      const rpc::DirectActorCallArgWaitCompleteRequest &request,
+      rpc::DirectActorCallArgWaitCompleteReply *reply,
+      rpc::SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcServer` for `WorkerService`.
@@ -48,10 +58,20 @@ class DirectActorGrpcService : public GrpcService {
                                   PushTaskReply>(
             service_, &DirectActorService::AsyncService::RequestPushTask,
             service_handler_, &DirectActorHandler::HandlePushTask, cq, main_service_));
-
-    // Set `PushTask`'s accept concurrency to 100.
     server_call_factories_and_concurrencies->emplace_back(
         std::move(push_task_call_Factory), 100);
+
+    // Initialize the Factory for `DirectActorCallArgWaitComplete` requests.
+    std::unique_ptr<ServerCallFactory> wait_complete_call_Factory(
+        new ServerCallFactoryImpl<DirectActorService, DirectActorHandler,
+                                  DirectActorCallArgWaitCompleteRequest,
+                                  DirectActorCallArgWaitCompleteReply>(
+            service_,
+            &DirectActorService::AsyncService::RequestDirectActorCallArgWaitComplete,
+            service_handler_, &DirectActorHandler::HandleDirectActorCallArgWaitComplete,
+            cq, main_service_));
+    server_call_factories_and_concurrencies->emplace_back(
+        std::move(wait_complete_call_Factory), 100);
   }
 
  private: