[GCS] Add actor checkpoint related methods to accessor (#6605)

src/ray/gcs/accessor.h

@@ -77,6 +77,37 @@ class ActorInfoAccessor {
   virtual Status AsyncUnsubscribe(const ActorID &actor_id,
                                   const StatusCallback &done) = 0;
 
+  /// Add actor checkpoint data to GCS asynchronously.
+  ///
+  /// \param data_ptr The checkpoint data that will be added to GCS.
+  /// \param callback The callback that will be called after add finishes.
+  /// \return Status
+  /// TODO(micafan) When the GCS backend is redis,
+  /// the checkpoint of the same actor needs to be updated serially,
+  /// otherwise the checkpoint may be overwritten. This issue will be resolved if
+  /// necessary.
+  virtual Status AsyncAddCheckpoint(
+      const std::shared_ptr<rpc::ActorCheckpointData> &data_ptr,
+      const StatusCallback &callback) = 0;
+
+  /// Get actor checkpoint data from GCS asynchronously.
+  ///
+  /// \param checkpoint_id The ID of checkpoint to lookup in GCS.
+  /// \param callback The callback that will be called after lookup finishes.
+  /// \return Status
+  virtual Status AsyncGetCheckpoint(
+      const ActorCheckpointID &checkpoint_id,
+      const OptionalItemCallback<rpc::ActorCheckpointData> &callback) = 0;
+
+  /// Get actor checkpoint id data from GCS asynchronously.
+  ///
+  /// \param actor_id The ID of actor to lookup in GCS.
+  /// \param callback The callback that will be called after lookup finishes.
+  /// \return Status
+  virtual Status AsyncGetCheckpointID(
+      const ActorID &actor_id,
+      const OptionalItemCallback<rpc::ActorCheckpointIdData> &callback) = 0;
+
  protected:
   ActorInfoAccessor() = default;
 };

src/ray/gcs/redis_accessor.cc

@@ -134,6 +134,79 @@ Status RedisActorInfoAccessor::AsyncUnsubscribe(const ActorID &actor_id,
   return actor_sub_executor_.AsyncUnsubscribe(subscribe_id_, actor_id, done);
 }
 
+Status RedisActorInfoAccessor::AsyncAddCheckpoint(
+    const std::shared_ptr<ActorCheckpointData> &data_ptr,
+    const StatusCallback &callback) {
+  auto on_add_data_done = [callback, data_ptr, this](
+                              RedisGcsClient *client,
+                              const ActorCheckpointID &checkpoint_id,
+                              const ActorCheckpointData &data) {
+    ActorID actor_id = ActorID::FromBinary(data_ptr->actor_id());
+    Status status = AsyncAddCheckpointID(actor_id, checkpoint_id, callback);
+    if (!status.ok()) {
+      callback(status);
+    }
+  };
+
+  ActorCheckpointID checkpoint_id =
+      ActorCheckpointID::FromBinary(data_ptr->checkpoint_id());
+  ActorCheckpointTable &actor_cp_table = client_impl_->actor_checkpoint_table();
+  return actor_cp_table.Add(JobID::Nil(), checkpoint_id, data_ptr, on_add_data_done);
+}
+
+Status RedisActorInfoAccessor::AsyncGetCheckpoint(
+    const ActorCheckpointID &checkpoint_id,
+    const OptionalItemCallback<ActorCheckpointData> &callback) {
+  RAY_CHECK(callback != nullptr);
+  auto on_success = [callback](RedisGcsClient *client,
+                               const ActorCheckpointID &checkpoint_id,
+                               const ActorCheckpointData &checkpoint_data) {
+    boost::optional<ActorCheckpointData> optional(checkpoint_data);
+    callback(Status::OK(), std::move(optional));
+  };
+
+  auto on_failure = [callback](RedisGcsClient *client,
+                               const ActorCheckpointID &checkpoint_id) {
+    boost::optional<ActorCheckpointData> optional;
+    callback(Status::Invalid("Invalid checkpoint id."), std::move(optional));
+  };
+
+  ActorCheckpointTable &actor_cp_table = client_impl_->actor_checkpoint_table();
+  return actor_cp_table.Lookup(JobID::Nil(), checkpoint_id, on_success, on_failure);
+}
+
+Status RedisActorInfoAccessor::AsyncGetCheckpointID(
+    const ActorID &actor_id,
+    const OptionalItemCallback<ActorCheckpointIdData> &callback) {
+  RAY_CHECK(callback != nullptr);
+  auto on_success = [callback](RedisGcsClient *client, const ActorID &actor_id,
+                               const ActorCheckpointIdData &data) {
+    boost::optional<ActorCheckpointIdData> optional(data);
+    callback(Status::OK(), std::move(optional));
+  };
+
+  auto on_failure = [callback](RedisGcsClient *client, const ActorID &actor_id) {
+    boost::optional<ActorCheckpointIdData> optional;
+    callback(Status::Invalid("Checkpoint not found."), std::move(optional));
+  };
+
+  ActorCheckpointIdTable &cp_id_table = client_impl_->actor_checkpoint_id_table();
+  return cp_id_table.Lookup(JobID::Nil(), actor_id, on_success, on_failure);
+}
+
+Status RedisActorInfoAccessor::AsyncAddCheckpointID(
+    const ActorID &actor_id, const ActorCheckpointID &checkpoint_id,
+    const StatusCallback &callback) {
+  ActorCheckpointIdTable::WriteCallback on_done = nullptr;
+  if (callback != nullptr) {
+    on_done = [callback](RedisGcsClient *client, const ActorID &actor_id,
+                         const ActorCheckpointIdData &data) { callback(Status::OK()); };
+  }
+
+  ActorCheckpointIdTable &cp_id_table = client_impl_->actor_checkpoint_id_table();
+  return cp_id_table.AddCheckpointId(JobID::Nil(), actor_id, checkpoint_id, on_done);
+}
+
 RedisJobInfoAccessor::RedisJobInfoAccessor(RedisGcsClient *client_impl)
     : client_impl_(client_impl), job_sub_executor_(client_impl->job_table()) {}
 

src/ray/gcs/redis_accessor.h

@@ -46,6 +46,27 @@ class RedisActorInfoAccessor : public ActorInfoAccessor {
 
   Status AsyncUnsubscribe(const ActorID &actor_id, const StatusCallback &done) override;
 
+  Status AsyncAddCheckpoint(const std::shared_ptr<ActorCheckpointData> &data_ptr,
+                            const StatusCallback &callback) override;
+
+  Status AsyncGetCheckpoint(
+      const ActorCheckpointID &checkpoint_id,
+      const OptionalItemCallback<ActorCheckpointData> &callback) override;
+
+  Status AsyncGetCheckpointID(
+      const ActorID &actor_id,
+      const OptionalItemCallback<ActorCheckpointIdData> &callback) override;
+
+ private:
+  /// Add checkpoint id to GCS asynchronously.
+  ///
+  /// \param actor_id The ID of actor that the checkpoint belongs to.
+  /// \param checkpoint_id The ID of checkpoint that will be added to GCS.
+  /// \return Status
+  Status AsyncAddCheckpointID(const ActorID &actor_id,
+                              const ActorCheckpointID &checkpoint_id,
+                              const StatusCallback &callback);
+
  private:
   RedisGcsClient *client_impl_{nullptr};
   // Use a random ClientID for actor subscription. Because:

src/ray/gcs/redis_gcs_client.h

@@ -18,7 +18,7 @@ namespace gcs {
 class RedisContext;
 
 class RAY_EXPORT RedisGcsClient : public GcsClient {
-  // TODO(micafan) Will remove those friend class / method after we replace RedisGcsClient
+  // TODO(micafan) Will remove those friend classes after we replace RedisGcsClient
   // with interface class GcsClient in raylet.
   friend class RedisActorInfoAccessor;
   friend class RedisJobInfoAccessor;
@@ -30,6 +30,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   friend class TaskTableTestHelper;
   friend class ClientTableTestHelper;
   friend class SetTestHelper;
+  friend class ActorCheckpointIdTable;
 
  public:
   /// Constructor of RedisGcsClient.
@@ -64,8 +65,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   TaskLeaseTable &task_lease_table();
   ErrorTable &error_table();
   ProfileTable &profile_table();
-  ActorCheckpointTable &actor_checkpoint_table();
-  ActorCheckpointIdTable &actor_checkpoint_id_table();
   DynamicResourceTable &resource_table();
   /// Used only for direct calls. Tasks submitted through the raylet transport
   /// should use Actors(), which has a requirement on the order in which
@@ -93,8 +92,10 @@ class RAY_EXPORT RedisGcsClient : public GcsClient {
   /// one event loop should be attached at a time.
   void Attach(boost::asio::io_service &io_service);
 
-  /// This method will be deprecated, use method Actors() instead.
+  /// The following three methods will be deprecated, use method Actors() instead.
   ActorTable &actor_table();
+  ActorCheckpointTable &actor_checkpoint_table();
+  ActorCheckpointIdTable &actor_checkpoint_id_table();
   /// This method will be deprecated, use method Jobs() instead.
   JobTable &job_table();
   /// This method will be deprecated, use method Objects() instead

src/ray/gcs/tables.cc

@@ -728,8 +728,9 @@ std::string ClientTable::DebugString() const {
 
 Status ActorCheckpointIdTable::AddCheckpointId(const JobID &job_id,
                                                const ActorID &actor_id,
-                                               const ActorCheckpointID &checkpoint_id) {
-  auto lookup_callback = [this, checkpoint_id, job_id, actor_id](
+                                               const ActorCheckpointID &checkpoint_id,
+                                               const WriteCallback &done) {
+  auto lookup_callback = [this, checkpoint_id, job_id, actor_id, done](
                              ray::gcs::RedisGcsClient *client, const ActorID &id,
                              const ActorCheckpointIdData &data) {
     std::shared_ptr<ActorCheckpointIdData> copy =
@@ -744,16 +745,16 @@ Status ActorCheckpointIdTable::AddCheckpointId(const JobID &job_id,
       copy->mutable_timestamps()->erase(copy->mutable_timestamps()->begin());
       client_->actor_checkpoint_table().Delete(job_id, to_delete);
     }
-    RAY_CHECK_OK(Add(job_id, actor_id, copy, nullptr));
+    RAY_CHECK_OK(Add(job_id, actor_id, copy, done));
   };
-  auto failure_callback = [this, checkpoint_id, job_id, actor_id](
+  auto failure_callback = [this, checkpoint_id, job_id, actor_id, done](
                               ray::gcs::RedisGcsClient *client, const ActorID &id) {
     std::shared_ptr<ActorCheckpointIdData> data =
         std::make_shared<ActorCheckpointIdData>();
     data->set_actor_id(id.Binary());
     data->add_timestamps(absl::GetCurrentTimeNanos() / 1000000);
     *data->add_checkpoint_ids() = checkpoint_id.Binary();
-    RAY_CHECK_OK(Add(job_id, actor_id, data, nullptr));
+    RAY_CHECK_OK(Add(job_id, actor_id, data, done));
   };
   return Lookup(job_id, actor_id, lookup_callback, failure_callback);
 }

src/ray/gcs/tables.h

@@ -786,7 +786,8 @@ class ActorCheckpointIdTable : public Table<ActorID, ActorCheckpointIdData> {
   /// \param checkpoint_id ID of the checkpoint.
   /// \return Status.
   Status AddCheckpointId(const JobID &job_id, const ActorID &actor_id,
-                         const ActorCheckpointID &checkpoint_id);
+                         const ActorCheckpointID &checkpoint_id,
+                         const WriteCallback &done);
 };
 
 namespace raylet {

src/ray/gcs/test/redis_actor_info_accessor_test.cc

@@ -26,7 +26,28 @@ class ActorInfoAccessorTest : public AccessorTestBase<ActorID, ActorTableData> {
       actor->set_actor_id(actor_id.Binary());
       id_to_data_[actor_id] = actor;
     }
+    GenCheckpointData();
   }
+
+  void GenCheckpointData() {
+    for (const auto item : id_to_data_) {
+      const ActorID &id = item.first;
+      ActorCheckpointList checkpoints;
+      for (size_t i = 0; i < checkpoint_number_; ++i) {
+        ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom();
+        auto checkpoint = std::make_shared<ActorCheckpointData>();
+        checkpoint->set_actor_id(id.Binary());
+        checkpoint->set_checkpoint_id(checkpoint_id.Binary());
+        checkpoint->set_execution_dependency(checkpoint_id.Binary());
+        checkpoints.emplace_back(checkpoint);
+      }
+      id_to_checkpoints_[id] = std::move(checkpoints);
+    }
+  }
+
+  typedef std::vector<std::shared_ptr<ActorCheckpointData>> ActorCheckpointList;
+  std::unordered_map<ActorID, ActorCheckpointList> id_to_checkpoints_;
+  size_t checkpoint_number_{2};
 };
 
 TEST_F(ActorInfoAccessorTest, RegisterAndGet) {
@@ -99,6 +120,66 @@ TEST_F(ActorInfoAccessorTest, Subscribe) {
   WaitPendingDone(sub_pending_count, wait_pending_timeout_);
 }
 
+TEST_F(ActorInfoAccessorTest, GetActorCheckpointTest) {
+  ActorInfoAccessor &actor_accessor = gcs_client_->Actors();
+  auto on_add_done = [this](Status status) {
+    RAY_CHECK_OK(status);
+    --pending_count_;
+  };
+  for (size_t index = 0; index < checkpoint_number_; ++index) {
+    for (const auto &actor_checkpoints : id_to_checkpoints_) {
+      const ActorCheckpointList &checkpoints = actor_checkpoints.second;
+      const auto &checkpoint = checkpoints[index];
+      ++pending_count_;
+      Status status = actor_accessor.AsyncAddCheckpoint(checkpoint, on_add_done);
+      RAY_CHECK_OK(status);
+    }
+    WaitPendingDone(wait_pending_timeout_);
+  }
+
+  for (const auto &actor_checkpoints : id_to_checkpoints_) {
+    const ActorCheckpointList &checkpoints = actor_checkpoints.second;
+    for (const auto &checkpoint : checkpoints) {
+      ActorCheckpointID checkpoint_id =
+          ActorCheckpointID::FromBinary(checkpoint->checkpoint_id());
+      auto on_get_done = [this, checkpoint_id](
+                             Status status,
+                             const boost::optional<ActorCheckpointData> &result) {
+        RAY_CHECK(result);
+        ActorCheckpointID result_checkpoint_id =
+            ActorCheckpointID::FromBinary(result->checkpoint_id());
+        ASSERT_EQ(checkpoint_id, result_checkpoint_id);
+        --pending_count_;
+      };
+      ++pending_count_;
+      Status status = actor_accessor.AsyncGetCheckpoint(checkpoint_id, on_get_done);
+      RAY_CHECK_OK(status);
+    }
+  }
+  WaitPendingDone(wait_pending_timeout_);
+
+  for (const auto &actor_checkpoints : id_to_checkpoints_) {
+    const ActorID &actor_id = actor_checkpoints.first;
+    const ActorCheckpointList &checkpoints = actor_checkpoints.second;
+    auto on_get_done = [this, &checkpoints](
+                           Status status,
+                           const boost::optional<ActorCheckpointIdData> &result) {
+      RAY_CHECK(result);
+      ASSERT_EQ(checkpoints.size(), result->checkpoint_ids_size());
+      for (size_t i = 0; i < checkpoints.size(); ++i) {
+        const std::string checkpoint_id_str = checkpoints[i]->checkpoint_id();
+        const std::string &result_checkpoint_id_str = result->checkpoint_ids(i);
+        ASSERT_EQ(checkpoint_id_str, result_checkpoint_id_str);
+      }
+      --pending_count_;
+    };
+    ++pending_count_;
+    Status status = actor_accessor.AsyncGetCheckpointID(actor_id, on_get_done);
+    RAY_CHECK_OK(status);
+  }
+  WaitPendingDone(wait_pending_timeout_);
+}
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/protobuf/gcs.proto

@@ -240,20 +240,22 @@ message JobTableData {
 // is the snapshot of an actor's state in the actor registration.
 // See `actor_registration.h` for more detailed explanation of these fields.
 message ActorCheckpointData {
+  // ID of this checkpoint.
+  bytes checkpoint_id = 1;
   // ID of this actor.
-  bytes actor_id = 1;
+  bytes actor_id = 2;
   // The dummy object ID of actor's most recently executed task.
-  bytes execution_dependency = 2;
+  bytes execution_dependency = 3;
   // A list of IDs of this actor's handles.
-  repeated bytes handle_ids = 3;
+  repeated bytes handle_ids = 4;
   // The task counters of the above handles.
-  repeated uint64 task_counters = 4;
+  repeated uint64 task_counters = 5;
   // The frontier dependencies of the above handles.
-  repeated bytes frontier_dependencies = 5;
+  repeated bytes frontier_dependencies = 6;
   // A list of unreleased dummy objects from this actor.
-  repeated bytes unreleased_dummy_objects = 6;
+  repeated bytes unreleased_dummy_objects = 7;
   // The numbers of dependencies for the above unreleased dummy objects.
-  repeated uint32 num_dummy_object_dependencies = 7;
+  repeated uint32 num_dummy_object_dependencies = 8;
 }
 
 // This table stores the actor-to-available-checkpoint-ids mapping.

src/ray/raylet/actor_registration.cc

@@ -120,6 +120,9 @@ std::shared_ptr<ActorCheckpointData> ActorRegistration::GenerateCheckpointData(
     checkpoint_data->add_unreleased_dummy_objects(entry.first.Binary());
     checkpoint_data->add_num_dummy_object_dependencies(entry.second);
   }
+
+  ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom();
+  checkpoint_data->set_checkpoint_id(checkpoint_id.Binary());
   return checkpoint_data;
 }
 

src/ray/raylet/node_manager.cc

@@ -1368,7 +1368,6 @@ void NodeManager::ProcessPrepareActorCheckpointRequest(
   std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
   RAY_CHECK(worker && worker->GetActorId() == actor_id);
 
-  ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom();
   std::shared_ptr<ActorCheckpointData> checkpoint_data;
   if (actor_entry->second.GetTableData().is_direct_call()) {
     checkpoint_data =
@@ -1383,17 +1382,15 @@ void NodeManager::ProcessPrepareActorCheckpointRequest(
   }
 
   // Write checkpoint data to GCS.
-  RAY_CHECK_OK(gcs_client_->actor_checkpoint_table().Add(
-      JobID::Nil(), checkpoint_id, checkpoint_data,
-      [worker, actor_id, this](ray::gcs::RedisGcsClient *client,
-                               const ActorCheckpointID &checkpoint_id,
-                               const ActorCheckpointData &data) {
+  RAY_CHECK_OK(gcs_client_->Actors().AsyncAddCheckpoint(
+      checkpoint_data, [worker, checkpoint_data](Status status) {
+        ActorCheckpointID checkpoint_id =
+            ActorCheckpointID::FromBinary(checkpoint_data->checkpoint_id());
+        RAY_CHECK(status.ok()) << "Add checkpoint failed, actor is "
+                               << worker->GetActorId() << " checkpoint_id is "
+                               << checkpoint_id;
         RAY_LOG(DEBUG) << "Checkpoint " << checkpoint_id << " saved for actor "
                        << worker->GetActorId();
-        // Save this actor-to-checkpoint mapping, and remove old checkpoints associated
-        // with this actor.
-        RAY_CHECK_OK(gcs_client_->actor_checkpoint_id_table().AddCheckpointId(
-            JobID::Nil(), actor_id, checkpoint_id));
         // Send reply to worker.
         flatbuffers::FlatBufferBuilder fbb;
         auto reply = ray::protocol::CreatePrepareActorCheckpointReply(
@@ -2487,15 +2484,16 @@ void NodeManager::FinishAssignedActorCreationTask(const ActorID &parent_actor_id
     checkpoint_id_to_restore_.erase(actor_id);
     RAY_LOG(DEBUG) << "Looking up checkpoint " << checkpoint_id << " for actor "
                    << actor_id;
-    RAY_CHECK_OK(gcs_client_->actor_checkpoint_table().Lookup(
-        JobID::Nil(), checkpoint_id,
-        [this, actor_id, new_actor_info, update_callback](
-            ray::gcs::RedisGcsClient *client, const UniqueID &checkpoint_id,
-            const ActorCheckpointData &checkpoint_data) {
+    RAY_CHECK_OK(gcs_client_->Actors().AsyncGetCheckpoint(
+        checkpoint_id,
+        [this, checkpoint_id, actor_id, new_actor_info, update_callback](
+            Status status, const boost::optional<ActorCheckpointData> &checkpoint_data) {
+          RAY_CHECK(checkpoint_data) << "Couldn't find checkpoint " << checkpoint_id
+                                     << " for actor " << actor_id << " in GCS.";
           RAY_LOG(INFO) << "Restoring registration for actor " << actor_id
                         << " from checkpoint " << checkpoint_id;
           ActorRegistration actor_registration =
-              ActorRegistration(*new_actor_info, checkpoint_data);
+              ActorRegistration(*new_actor_info, *checkpoint_data);
           // Mark the unreleased dummy objects in the checkpoint frontier as local.
           for (const auto &entry : actor_registration.GetDummyObjects()) {
             HandleObjectLocal(entry.first);
@@ -2504,10 +2502,6 @@ void NodeManager::FinishAssignedActorCreationTask(const ActorID &parent_actor_id
           // The actor was created before.
           RAY_CHECK_OK(gcs_client_->Actors().AsyncUpdate(actor_id, new_actor_info,
                                                          update_callback));
-        },
-        [actor_id](ray::gcs::RedisGcsClient *client, const UniqueID &checkpoint_id) {
-          RAY_LOG(FATAL) << "Couldn't find checkpoint " << checkpoint_id << " for actor "
-                         << actor_id << " in GCS.";
         }));
   } else {
     // The actor did not resume from a checkpoint. Immediately notify the