[core] Reduce assertion check in NodeManager on worker not found (ray…

…-project#41841) Avoid crashing the raylet when a worker that sends "available" IPC is not found. Related to ray-project#41477, although there may be a Tune/Train issue there as well. Signed-off-by: Stephanie Wang <[email protected]>
c21 · Dec 14, 2023 · 98d2e86 · 98d2e86
1 parent dde5461
commit 98d2e86
Show file tree

Hide file tree

Showing 10 changed files with 24 additions and 29 deletions.
diff --git a/src/ray/common/client_connection.h b/src/ray/common/client_connection.h
@@ -227,6 +227,8 @@ class ClientConnection : public ServerConnection {
  /// ProcessClientMessage handler will be called.
  void ProcessMessages();
 
+ const std::string GetDebugLabel() const { return debug_label_; }
+
  protected:
  /// A protected constructor for a node client connection.
  ClientConnection(MessageHandler &message_handler,

diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc
@@ -146,7 +146,7 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
  std::placeholders::_8);
  direct_task_receiver_ = std::make_unique<CoreWorkerDirectTaskReceiver>(
  worker_context_, task_execution_service_, execute_task, [this] {
- return local_raylet_client_->TaskDone();
+ return local_raylet_client_->ActorCreationTaskDone();
  });
  }
 

diff --git a/src/ray/core_worker/test/direct_actor_transport_test.cc b/src/ray/core_worker/test/direct_actor_transport_test.cc
@@ -720,12 +720,13 @@ class MockWorkerContext : public WorkerContext {
 
 class MockCoreWorkerDirectTaskReceiver : public CoreWorkerDirectTaskReceiver {
  public:
- MockCoreWorkerDirectTaskReceiver(WorkerContext &worker_context,
- instrumented_io_context &main_io_service,
- const TaskHandler &task_handler,
- const OnTaskDone &task_done)
+ MockCoreWorkerDirectTaskReceiver(
+ WorkerContext &worker_context,
+ instrumented_io_context &main_io_service,
+ const TaskHandler &task_handler,
+ const OnActorCreationTaskDone &actor_creation_task_done_)
  : CoreWorkerDirectTaskReceiver(
- worker_context, main_io_service, task_handler, task_done) {}
+ worker_context, main_io_service, task_handler, actor_creation_task_done_) {}
 
  void UpdateConcurrencyGroupsCache(const ActorID &actor_id,
  const std::vector<ConcurrencyGroup> &cgs) {

diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc
@@ -174,7 +174,7 @@ void CoreWorkerDirectTaskReceiver::HandleTask(
  // Tell raylet that an actor creation task has finished execution, so that
  // raylet can publish actor creation event to GCS, and mark this worker as
  // actor, thus if this worker dies later raylet will restart the actor.
- RAY_CHECK_OK(task_done_());
+ RAY_CHECK_OK(actor_creation_task_done_());
  if (status.IsCreationTaskError()) {
  RAY_LOG(WARNING) << "Actor creation task finished with errors, task_id: "
  << task_spec.TaskId()

diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h
@@ -61,16 +61,16 @@ class CoreWorkerDirectTaskReceiver {
  bool *is_retryable_error,
  std::string *application_error)>;
 
- using OnTaskDone = std::function<Status()>;
+ using OnActorCreationTaskDone = std::function<Status()>;
 
  CoreWorkerDirectTaskReceiver(WorkerContext &worker_context,
  instrumented_io_context &main_io_service,
  const TaskHandler &task_handler,
- const OnTaskDone &task_done)
+ const OnActorCreationTaskDone &actor_creation_task_done_)
  : worker_context_(worker_context),
  task_handler_(task_handler),
  task_main_io_service_(main_io_service),
- task_done_(task_done),
+ actor_creation_task_done_(actor_creation_task_done_),
  pool_manager_(std::make_shared<ConcurrencyGroupManager<BoundedExecutor>>()),
  fiber_state_manager_(nullptr) {}
 
@@ -127,7 +127,7 @@ class CoreWorkerDirectTaskReceiver {
  /// The IO event loop for running tasks on.
  instrumented_io_context &task_main_io_service_;
  /// The callback function to be invoked when finishing a task.
- OnTaskDone task_done_;
+ OnActorCreationTaskDone actor_creation_task_done_;
  /// Shared pool for producing new core worker clients.
  std::shared_ptr<rpc::CoreWorkerClientPool> client_pool_;
  /// Address of our RPC server.

diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs
@@ -24,7 +24,7 @@ enum MessageType:int {
  SubmitTask = 1,
  // Notify the raylet that a task has finished. This is sent from a
  // worker to a raylet.
- TaskDone,
+ ActorCreationTaskDone,
  // Log a message to the event table. This is sent from a worker to a raylet.
  EventLogMessage,
  // Send an initial connection message to the raylet. This is sent

diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc
@@ -1223,8 +1223,11 @@ void NodeManager::ProcessClientMessage(const std::shared_ptr<ClientConnection> &
  case protocol::MessageType::AnnounceWorkerPort: {
  ProcessAnnounceWorkerPortMessage(client, message_data);
  } break;
- case protocol::MessageType::TaskDone: {
- HandleWorkerAvailable(client);
+ case protocol::MessageType::ActorCreationTaskDone: {
+ if (registered_worker) {
+ // Worker may send this message after it was disconnected.
+ HandleWorkerAvailable(registered_worker);
+ }
  } break;
  case protocol::MessageType::DisconnectClient: {
  ProcessDisconnectClientMessage(client, message_data);
@@ -1400,15 +1403,10 @@ void NodeManager::ProcessAnnounceWorkerPortMessage(
  worker->Connect(port);
  if (is_worker) {
  worker_pool_.OnWorkerStarted(worker);
- HandleWorkerAvailable(worker->Connection());
+ HandleWorkerAvailable(worker);
  }
 }
 
-void NodeManager::HandleWorkerAvailable(const std::shared_ptr<ClientConnection> &client) {
- std::shared_ptr<WorkerInterface> worker = worker_pool_.GetRegisteredWorker(client);
- HandleWorkerAvailable(worker);
-}
-
 void NodeManager::HandleWorkerAvailable(const std::shared_ptr<WorkerInterface> &worker) {
  RAY_CHECK(worker);
 

diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h
@@ -445,12 +445,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
  void ProcessAnnounceWorkerPortMessage(const std::shared_ptr<ClientConnection> &client,
  const uint8_t *message_data);
 
- /// Handle the case that a worker is available.
- ///
- /// \param client The connection for the worker.
- /// \return Void.
- void HandleWorkerAvailable(const std::shared_ptr<ClientConnection> &client);
-
  /// Handle the case that a worker is available.
  ///
  /// \param worker The pointer to the worker

diff --git a/src/ray/raylet_client/raylet_client.cc b/src/ray/raylet_client/raylet_client.cc
@@ -200,8 +200,8 @@ Status raylet::RayletClient::AnnounceWorkerPort(int port) {
  return conn_->WriteMessage(MessageType::AnnounceWorkerPort, &fbb);
 }
 
-Status raylet::RayletClient::TaskDone() {
- return conn_->WriteMessage(MessageType::TaskDone);
+Status raylet::RayletClient::ActorCreationTaskDone() {
+ return conn_->WriteMessage(MessageType::ActorCreationTaskDone);
 }
 
 Status raylet::RayletClient::FetchOrReconstruct(

diff --git a/src/ray/raylet_client/raylet_client.h b/src/ray/raylet_client/raylet_client.h
@@ -304,7 +304,7 @@ class RayletClient : public RayletClientInterface {
  /// Tell the raylet that the client has finished executing a task.
  ///
  /// \return ray::Status.
- ray::Status TaskDone();
+ ray::Status ActorCreationTaskDone();
 
  /// Tell the raylet to reconstruct or fetch objects.
  ///