[Core] Add PENDING_ACTOR_TASK_ARGS_FETCH and PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY TaskStatus

python/ray/tests/test_actor_cancel.py

@@ -222,7 +222,7 @@ async def g(self):
             list_tasks(
                 filters=[
                     ("name", "=", "Actor.g"),
-                    ("STATE", "=", "SUBMITTED_TO_WORKER"),
+                    ("STATE", "=", "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY"),
                 ]
             )
         )
@@ -238,9 +238,8 @@ async def g(self):
             ray.get(ref)
 
     # Verify the task is submitted to the worker and never executed
-    # assert task.state == "SUBMITTED_TO_WORKER"
     for task in tasks:
-        assert task.state == "SUBMITTED_TO_WORKER"
+        assert task.state == "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY"
 
 
 def test_async_actor_cancel_after_task_finishes(shutdown_only):

python/ray/tests/test_state_api.py

@@ -2632,7 +2632,7 @@ def test_list_actor_tasks(shutdown_only):
     ray.init(num_cpus=2)
     job_id = ray.get_runtime_context().get_job_id()
 
-    @ray.remote
+    @ray.remote(max_concurrency=2)
     class Actor:
         def call(self):
             import time
@@ -2650,18 +2650,19 @@ def verify():
         for task in tasks:
             assert task["actor_id"] == actor_id
         # Actor.__init__: 1 finished
-        # Actor.call: 1 running, 9 waiting for execution (queued).
+        # Actor.call: 2 running, 8 waiting for execution (queued).
         assert len(tasks) == 11
         assert (
             len(
                 list(
                     filter(
-                        lambda task: task["state"] == "SUBMITTED_TO_WORKER",
+                        lambda task: task["state"]
+                        == "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY",
                         tasks,
                     )
                 )
             )
-            == 9
+            == 8
         )
         assert (
             len(
@@ -2694,7 +2695,7 @@ def verify():
                     )
                 )
             )
-            == 1
+            == 2
         )
 
         # Filters with actor id.

python/ray/util/state/custom_types.py

@@ -39,6 +39,8 @@
     "PENDING_OBJ_STORE_MEM_AVAIL",
     "PENDING_ARGS_FETCH",
     "SUBMITTED_TO_WORKER",
+    "PENDING_ACTOR_TASK_ARGS_FETCH",
+    "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY",
     "RUNNING",
     "RUNNING_IN_RAY_GET",
     "RUNNING_IN_RAY_WAIT",

src/ray/common/task/task_spec.cc

@@ -454,7 +454,7 @@ int TaskSpecification::MaxActorConcurrency() const {
   return message_->actor_creation_task_spec().max_concurrency();
 }
 
-std::string TaskSpecification::ConcurrencyGroupName() const {
+const std::string &TaskSpecification::ConcurrencyGroupName() const {
   RAY_CHECK(IsActorTask());
   return message_->concurrency_group_name();
 }

src/ray/common/task/task_spec.h

@@ -481,7 +481,7 @@ class TaskSpecification : public MessageWrapper<rpc::TaskSpec> {
   // Concurrency groups of the actor.
   std::vector<ConcurrencyGroup> ConcurrencyGroups() const;
 
-  std::string ConcurrencyGroupName() const;
+  const std::string &ConcurrencyGroupName() const;
 
   bool ExecuteOutOfOrder() const;
 

src/ray/core_worker/core_worker.cc

@@ -168,6 +168,9 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
 #endif
   }
 
+  task_event_buffer_ = std::make_unique<worker::TaskEventBufferImpl>(
+      std::make_shared<gcs::GcsClient>(options_.gcs_options));
+
   // Initialize task receivers.
   if (options_.worker_type == WorkerType::WORKER || options_.is_local_mode) {
     RAY_CHECK(options_.task_execution_callback != nullptr);
@@ -182,9 +185,11 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
                                   std::placeholders::_7,
                                   std::placeholders::_8);
     task_receiver_ = std::make_unique<TaskReceiver>(
-        worker_context_, task_execution_service_, execute_task, [this] {
-          return local_raylet_client_->ActorCreationTaskDone();
-        });
+        worker_context_,
+        task_execution_service_,
+        *task_event_buffer_,
+        execute_task,
+        [this] { return local_raylet_client_->ActorCreationTaskDone(); });
   }
 
   // Initialize raylet client.
@@ -281,9 +286,6 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
   RAY_CHECK_OK(gcs_client_->Connect(io_service_));
   RegisterToGcs(options_.worker_launch_time_ms, options_.worker_launched_time_ms);
 
-  // Initialize the task state event buffer.
-  task_event_buffer_ = std::make_unique<worker::TaskEventBufferImpl>(
-      std::make_shared<gcs::GcsClient>(options_.gcs_options));
   if (RayConfig::instance().task_events_report_interval_ms() > 0) {
     if (!task_event_buffer_->Start().ok()) {
       RAY_CHECK(!task_event_buffer_->Enabled()) << "TaskEventBuffer should be disabled.";
@@ -638,6 +640,9 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
           RAY_LOG(INFO) << "Event stats:\n\n"
                         << io_service_.stats().StatsString() << "\n\n"
                         << "-----------------\n"
+                        << "Task execution event stats:\n"
+                        << task_execution_service_.stats().StatsString() << "\n\n"
+                        << "-----------------\n"
                         << "Task Event stats:\n"
                         << task_event_buffer_->DebugString() << "\n";
         },
@@ -2941,13 +2946,13 @@ Status CoreWorker::ExecuteTask(
             ? worker::TaskStatusEvent::TaskStateUpdate(actor_repr_name, pid_)
             : worker::TaskStatusEvent::TaskStateUpdate(pid_);
     RAY_UNUSED(
-        task_manager_->RecordTaskStatusEventIfNeeded(task_spec.TaskId(),
-                                                     worker_context_.GetCurrentJobID(),
-                                                     task_spec.AttemptNumber(),
-                                                     task_spec,
-                                                     rpc::TaskStatus::RUNNING,
-                                                     /* include_task_info */ false,
-                                                     update));
+        task_event_buffer_->RecordTaskStatusEventIfNeeded(task_spec.TaskId(),
+                                                          task_spec.JobId(),
+                                                          task_spec.AttemptNumber(),
+                                                          task_spec,
+                                                          rpc::TaskStatus::RUNNING,
+                                                          /* include_task_info */ false,
+                                                          update));
 
     worker_context_.SetCurrentTask(task_spec);
     SetCurrentTaskId(task_spec.TaskId(), task_spec.AttemptNumber(), task_spec.GetName());
@@ -4677,7 +4682,7 @@ void CoreWorker::RecordTaskLogStart(const TaskID &task_id,
   auto current_task = worker_context_.GetCurrentTask();
   RAY_CHECK(current_task)
       << "We should have set the current task spec while executing the task.";
-  RAY_UNUSED(task_manager_->RecordTaskStatusEventIfNeeded(
+  RAY_UNUSED(task_event_buffer_->RecordTaskStatusEventIfNeeded(
       task_id,
       worker_context_.GetCurrentJobID(),
       attempt_number,
@@ -4701,7 +4706,7 @@ void CoreWorker::RecordTaskLogEnd(const TaskID &task_id,
   auto current_task = worker_context_.GetCurrentTask();
   RAY_CHECK(current_task)
       << "We should have set the current task spec before executing the task.";
-  RAY_UNUSED(task_manager_->RecordTaskStatusEventIfNeeded(
+  RAY_UNUSED(task_event_buffer_->RecordTaskStatusEventIfNeeded(
       task_id,
       worker_context_.GetCurrentJobID(),
       attempt_number,
@@ -4719,7 +4724,7 @@ void CoreWorker::UpdateTaskIsDebuggerPaused(const TaskID &task_id,
       << "We should have set the current task spec before executing the task.";
   RAY_LOG(DEBUG).WithField(current_task_it->second.TaskId())
       << "Task is paused by debugger set to " << is_debugger_paused;
-  RAY_UNUSED(task_manager_->RecordTaskStatusEventIfNeeded(
+  RAY_UNUSED(task_event_buffer_->RecordTaskStatusEventIfNeeded(
       task_id,
       worker_context_.GetCurrentJobID(),
       current_task_it->second.AttemptNumber(),

src/ray/core_worker/core_worker.h

@@ -223,30 +223,6 @@ class TaskCounter {
     }
   }
 
-  std::unordered_map<std::string, std::vector<int64_t>> AsMap() const {
-    absl::MutexLock l(&mu_);
-    std::unordered_map<std::string, std::vector<int64_t>> total_counts;
-
-    counter_.ForEachEntry(
-        [&total_counts](const std::tuple<std::string, TaskStatusType, bool> &key,
-                        int64_t value) mutable {
-          auto func_name = std::get<0>(key);
-          auto status = std::get<1>(key);
-          total_counts[func_name].resize(3, 0);
-          if (status == kPending) {
-            total_counts[func_name][0] = value;
-          } else if (status == kRunning) {
-            total_counts[func_name][1] = value;
-          } else if (status == kFinished) {
-            total_counts[func_name][2] = value;
-          } else {
-            RAY_CHECK(false) << "Invalid task status type " << status;
-          }
-        });
-
-    return total_counts;
-  }
-
  private:
   mutable absl::Mutex mu_;
   // Tracks all tasks submitted to this worker by state, is_retry.

src/ray/core_worker/task_event_buffer.cc

@@ -202,6 +202,34 @@ void TaskProfileEvent::ToRpcTaskExportEvents(
   event_entry->set_extra_data(std::move(extra_data_));
 }
 
+bool TaskEventBuffer::RecordTaskStatusEventIfNeeded(
+    const TaskID &task_id,
+    const JobID &job_id,
+    int32_t attempt_number,
+    const TaskSpecification &spec,
+    rpc::TaskStatus status,
+    bool include_task_info,
+    absl::optional<const TaskStatusEvent::TaskStateUpdate> state_update) {
+  if (!Enabled()) {
+    return false;
+  }
+  if (!spec.EnableTaskEvents()) {
+    return false;
+  }
+
+  auto task_event = std::make_unique<TaskStatusEvent>(
+      task_id,
+      job_id,
+      attempt_number,
+      status,
+      /* timestamp */ absl::GetCurrentTimeNanos(),
+      include_task_info ? std::make_shared<const TaskSpecification>(spec) : nullptr,
+      std::move(state_update));
+
+  AddTaskEvent(std::move(task_event));
+  return true;
+}
+
 TaskEventBufferImpl::TaskEventBufferImpl(std::shared_ptr<gcs::GcsClient> gcs_client)
     : work_guard_(boost::asio::make_work_guard(io_service_)),
       periodical_runner_(io_service_),

src/ray/core_worker/task_event_buffer.h

@@ -87,7 +87,7 @@ class TaskStatusEvent : public TaskEvent {
   struct TaskStateUpdate {
     TaskStateUpdate() {}
 
-    TaskStateUpdate(const absl::optional<const rpc::RayErrorInfo> &error_info)
+    TaskStateUpdate(const std::optional<const rpc::RayErrorInfo> &error_info)
         : error_info_(error_info) {}
 
     TaskStateUpdate(const NodeID &node_id, const WorkerID &worker_id)
@@ -107,19 +107,19 @@ class TaskStatusEvent : public TaskEvent {
     friend class TaskStatusEvent;
 
     /// Node id if it's a SUBMITTED_TO_WORKER status change.
-    const absl::optional<NodeID> node_id_ = absl::nullopt;
+    const std::optional<NodeID> node_id_ = std::nullopt;
     /// Worker id if it's a SUBMITTED_TO_WORKER status change.
-    const absl::optional<WorkerID> worker_id_ = absl::nullopt;
+    const std::optional<WorkerID> worker_id_ = std::nullopt;
     /// Task error info.
-    const absl::optional<rpc::RayErrorInfo> error_info_ = absl::nullopt;
+    const std::optional<rpc::RayErrorInfo> error_info_ = std::nullopt;
     /// Task log info.
-    const absl::optional<rpc::TaskLogInfo> task_log_info_ = absl::nullopt;
+    const std::optional<rpc::TaskLogInfo> task_log_info_ = std::nullopt;
     /// Actor task repr name.
     const std::string actor_repr_name_ = "";
     /// Worker's pid if it's a RUNNING status change.
-    const absl::optional<uint32_t> pid_ = absl::nullopt;
-    /// Is the task is paused by the debugger.
-    const absl::optional<bool> is_debugger_paused_ = absl::nullopt;
+    const std::optional<uint32_t> pid_ = std::nullopt;
+    /// If the task is paused by the debugger.
+    const std::optional<bool> is_debugger_paused_ = std::nullopt;
   };
 
   explicit TaskStatusEvent(
@@ -129,7 +129,7 @@ class TaskStatusEvent : public TaskEvent {
       const rpc::TaskStatus &task_status,
       int64_t timestamp,
       const std::shared_ptr<const TaskSpecification> &task_spec = nullptr,
-      absl::optional<const TaskStateUpdate> state_update = absl::nullopt);
+      std::optional<const TaskStateUpdate> state_update = std::nullopt);
 
   void ToRpcTaskEvents(rpc::TaskEvents *rpc_task_events) override;
 
@@ -146,7 +146,7 @@ class TaskStatusEvent : public TaskEvent {
   /// Pointer to the task spec.
   const std::shared_ptr<const TaskSpecification> task_spec_ = nullptr;
   /// Optional task state update
-  const absl::optional<const TaskStateUpdate> state_update_ = absl::nullopt;
+  const std::optional<const TaskStateUpdate> state_update_ = std::nullopt;
 };
 
 /// TaskProfileEvent is generated when `RAY_enable_timeline` is on.
@@ -221,6 +221,28 @@ class TaskEventBuffer {
  public:
   virtual ~TaskEventBuffer() = default;
 
+  /// Update task status change for the task attempt in TaskEventBuffer if needed.
+  ///
+  /// It skips the reporting when:
+  ///   1. when the enable_task_events for the task is false in TaskSpec.
+  ///   2. when the task event reporting is disabled on the worker (through ray config,
+  ///   i.e., RAY_task_events_report_interval_ms=0).
+  ///
+  /// \param attempt_number Attempt number for the task attempt.
+  /// \param spec corresponding TaskSpecification of the task
+  /// \param status the changed status.
+  /// \param state_update optional task state updates.
+  /// \return true if the event is recorded, false otherwise.
+  bool RecordTaskStatusEventIfNeeded(
+      const TaskID &task_id,
+      const JobID &job_id,
+      int32_t attempt_number,
+      const TaskSpecification &spec,
+      rpc::TaskStatus status,
+      bool include_task_info = false,
+      absl::optional<const TaskStatusEvent::TaskStateUpdate> state_update =
+          absl::nullopt);
+
   /// Add a task event to be reported.
   ///
   /// \param task_events Task events.