[direct task] Remove timeout for resolving futures that were deserialized (#6337)

* Reply GetObjectStatus once the task completes

* Remove timeout-based future resolution

* fix

* Update core_worker.h

src/ray/core_worker/core_worker.cc

@@ -213,7 +213,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
           },
           memory_store_, task_manager_, local_raylet_id,
           RayConfig::instance().worker_lease_timeout_milliseconds()));
-  future_resolver_.reset(new FutureResolver(memory_store_, client_factory, io_service_));
+  future_resolver_.reset(new FutureResolver(memory_store_, client_factory));
 }
 
 CoreWorker::~CoreWorker() {
@@ -966,16 +966,24 @@ void CoreWorker::HandleGetObjectStatus(const rpc::GetObjectStatusRequest &reques
     // We may have owned this object in the past, but we are now executing some
     // other task or actor.
     reply->set_status(rpc::GetObjectStatusReply::WRONG_OWNER);
+    send_reply_callback(Status::OK(), nullptr, nullptr);
   } else {
+    // We own the task. Reply back to the borrower once the object has been
+    // created.
+    // TODO: We could probably just send the object value if it is small
+    // enough and we have it local.
+    reply->set_status(rpc::GetObjectStatusReply::CREATED);
     if (task_manager_->IsTaskPending(object_id.TaskId())) {
-      reply->set_status(rpc::GetObjectStatusReply::PENDING);
+      // The task is pending. Send the reply once the task finishes.
+      memory_store_->GetAsync(object_id,
+                              [send_reply_callback](std::shared_ptr<RayObject> obj) {
+                                send_reply_callback(Status::OK(), nullptr, nullptr);
+                              });
     } else {
-      // TODO: We could probably just send the object value if it is small
-      // enough and we have it local.
-      reply->set_status(rpc::GetObjectStatusReply::CREATED);
+      // The task is done. Send the reply immediately.
+      send_reply_callback(Status::OK(), nullptr, nullptr);
     }
   }
-  send_reply_callback(Status::OK(), nullptr, nullptr);
 }
 
 void CoreWorker::YieldCurrentFiber(FiberEvent &event) {

src/ray/core_worker/core_worker.h

@@ -30,7 +30,7 @@
   RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5)                       \
   RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999)                      \
   RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100) \
-  RAY_CORE_WORKER_RPC_HANDLER(GetObjectStatus, 100)
+  RAY_CORE_WORKER_RPC_HANDLER(GetObjectStatus, 9999)
 
 namespace ray {
 

src/ray/core_worker/future_resolver.cc

@@ -10,45 +10,27 @@ void FutureResolver::ResolveFutureAsync(const ObjectID &object_id, const TaskID
   if (it == owner_clients_.end()) {
     auto client = std::shared_ptr<rpc::CoreWorkerClientInterface>(
         client_factory_({owner_address.ip_address(), owner_address.port()}));
-    owner_clients_.emplace(owner_id, std::move(client));
+    it = owner_clients_.emplace(owner_id, std::move(client)).first;
   }
 
-  // This timer will get deallocated once the future has been resolved.
-  auto timer = std::make_shared<boost::asio::deadline_timer>(io_service_);
-  AttemptFutureResolution(object_id, owner_id, std::move(timer));
-}
-
-void FutureResolver::AttemptFutureResolution(
-    const ObjectID &object_id, const TaskID &owner_id,
-    std::shared_ptr<boost::asio::deadline_timer> timer) {
-  auto &owner_client = owner_clients_[owner_id];
   rpc::GetObjectStatusRequest request;
   request.set_object_id(object_id.Binary());
   request.set_owner_id(owner_id.Binary());
-  auto status = owner_client->GetObjectStatus(
-      request, [this, object_id, owner_id, timer](
-                   const Status &status, const rpc::GetObjectStatusReply &reply) {
-        if (!status.ok() || reply.status() != rpc::GetObjectStatusReply::PENDING) {
-          // Either the owner is gone or the owner replied that the object has
-          // been created. In both cases, we can now try to fetch the object via
-          // plasma.
-          RAY_CHECK_OK(in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA),
-                                             object_id));
-        } else {
-          // Try again later.
-          timer->expires_from_now(
-              boost::posix_time::milliseconds(wait_future_resolution_milliseconds_));
-          timer->async_wait(
-              [this, object_id, owner_id, timer](const boost::system::error_code &error) {
-                absl::MutexLock lock(&mu_);
-                AttemptFutureResolution(object_id, owner_id, std::move(timer));
-              });
+  RAY_CHECK_OK(it->second->GetObjectStatus(
+      request,
+      [this, object_id](const Status &status, const rpc::GetObjectStatusReply &reply) {
+        if (!status.ok() || reply.status() == rpc::GetObjectStatusReply::WRONG_OWNER) {
+          RAY_LOG(ERROR)
+              << "Error retrieving the value of object ID " << object_id
+              << " that was deserialized. Probably, the task or actor that created the "
+                 "object ID initially (via ray.put or task submission) has exited.";
         }
-      });
-  if (!status.ok()) {
-    RAY_CHECK_OK(
-        in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id));
-  }
+        // Either the owner is gone or the owner replied that the object has
+        // been created. In both cases, we can now try to fetch the object via
+        // plasma.
+        RAY_CHECK_OK(in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA),
+                                           object_id));
+      }));
 }
 
 }  // namespace ray

src/ray/core_worker/future_resolver.h

@@ -10,22 +10,13 @@
 
 namespace ray {
 
-/// Max time between requests to the owner to check whether the object is still
-/// being computed.
-const int kWaitObjectEvictionMilliseconds = 100;
-
 // Resolve values for futures that were given to us before the value
 // was available. This class is thread-safe.
 class FutureResolver {
  public:
-  FutureResolver(
-      std::shared_ptr<CoreWorkerMemoryStore> store, rpc::ClientFactoryFn client_factory,
-      boost::asio::io_service &io_service,
-      int wait_future_resolution_milliseconds = kWaitObjectEvictionMilliseconds)
-      : in_memory_store_(store),
-        client_factory_(client_factory),
-        io_service_(io_service),
-        wait_future_resolution_milliseconds_(wait_future_resolution_milliseconds) {}
+  FutureResolver(std::shared_ptr<CoreWorkerMemoryStore> store,
+                 rpc::ClientFactoryFn client_factory)
+      : in_memory_store_(store), client_factory_(client_factory) {}
 
   /// Resolve the value for a future. This will periodically contact the given
   /// owner until the owner dies or the owner has finished creating the object.
@@ -40,24 +31,12 @@ class FutureResolver {
                           const rpc::Address &owner_address);
 
  private:
-  // Attempt to contact the owner to ask about the future's current status.
-  void AttemptFutureResolution(const ObjectID &object_id, const TaskID &owner_id,
-                               std::shared_ptr<boost::asio::deadline_timer> timer)
-      EXCLUSIVE_LOCKS_REQUIRED(mu_);
-
-  /// Used to set timers.
-  boost::asio::io_service &io_service_;
-
   /// Used to store values of resolved futures.
   std::shared_ptr<CoreWorkerMemoryStore> in_memory_store_;
 
   /// Factory for producing new core worker clients.
   const rpc::ClientFactoryFn client_factory_;
 
-  /// The amount of time to wait between requests to a future's owner to get
-  /// the object's current status.
-  const int wait_future_resolution_milliseconds_;
-
   /// Protects against concurrent access to internal state.
   absl::Mutex mu_;
 

src/ray/protobuf/core_worker.proto

@@ -94,9 +94,8 @@ message GetObjectStatusRequest {
 
 message GetObjectStatusReply {
   enum ObjectStatus {
-    PENDING = 0;
-    CREATED = 1;
-    WRONG_OWNER = 2;
+    CREATED = 0;
+    WRONG_OWNER = 1;
   }
   ObjectStatus status = 1;
 }