[Core] ray list tasks filter state and name on gcs side

python/ray/util/state/state_manager.py

@@ -307,6 +307,10 @@ async def get_all_task_info(
                 req_filters.job_id = JobID(hex_to_binary(value)).binary()
             elif key == "task_id":
                 req_filters.task_ids.append(TaskID(hex_to_binary(value)).binary())
+            elif key == "name":
+                req_filters.name = value
+            elif key == "state":
+                req_filters.state = value
             else:
                 continue
 

src/ray/gcs/gcs_server/gcs_task_manager.cc

@@ -112,7 +112,7 @@ void GcsTaskManager::GcsTaskManagerStorage::MarkTaskAttemptFailedIfNeeded(
   // We could mark the task as failed even if might not have state updates yet (i.e. only
   // profiling events are reported).
   auto state_updates = task_events.mutable_state_updates();
-  state_updates->set_failed_ts(failed_ts);
+  (*state_updates->mutable_state_ts())[ray::rpc::TaskStatus::FAILED] = failed_ts;
   state_updates->mutable_error_info()->CopyFrom(error_info);
 }
 
@@ -423,6 +423,28 @@ void GcsTaskManager::HandleGetTaskEvents(rpc::GetTaskEventsRequest request,
       return false;
     }
 
+    if (filters.has_state()) {
+      const google::protobuf::EnumDescriptor *task_status_descriptor =
+          ray::rpc::TaskStatus_descriptor();
+
+      // Figure out the latest state of a task.
+      ray::rpc::TaskStatus state = ray::rpc::TaskStatus::NIL;
+      if (task_event.has_state_updates()) {
+        for (int i = task_status_descriptor->value_count() - 1; i >= 0; --i) {
+          if (task_event.state_updates().state_ts().contains(
+                  task_status_descriptor->value(i)->number())) {
+            state = static_cast<ray::rpc::TaskStatus>(
+                task_status_descriptor->value(i)->number());
+            break;
+          }
+        }
+      }
+
+      if (filters.state() != task_status_descriptor->FindValueByNumber(state)->name()) {
+        return false;
+      }
+    }
+
     return true;
   };
 

src/ray/gcs/pb_util.h

@@ -304,7 +304,8 @@ inline bool IsTaskTerminated(const rpc::TaskEvents &task_event) {
   }
 
   const auto &state_updates = task_event.state_updates();
-  return state_updates.has_finished_ts() || state_updates.has_failed_ts();
+  return state_updates.state_ts().contains(rpc::TaskStatus::FINISHED) ||
+         state_updates.state_ts().contains(rpc::TaskStatus::FAILED);
 }
 
 inline size_t NumProfileEvents(const rpc::TaskEvents &task_event) {
@@ -335,7 +336,7 @@ inline bool IsTaskFinished(const rpc::TaskEvents &task_event) {
   }
 
   const auto &state_updates = task_event.state_updates();
-  return state_updates.has_finished_ts();
+  return state_updates.state_ts().contains(rpc::TaskStatus::FINISHED);
 }
 
 /// Fill the rpc::TaskStateUpdate with the timestamps according to the status change.
@@ -346,39 +347,11 @@ inline bool IsTaskFinished(const rpc::TaskEvents &task_event) {
 inline void FillTaskStatusUpdateTime(const ray::rpc::TaskStatus &task_status,
                                      int64_t timestamp,
                                      ray::rpc::TaskStateUpdate *state_updates) {
-  switch (task_status) {
-  case rpc::TaskStatus::PENDING_ARGS_AVAIL: {
-    state_updates->set_pending_args_avail_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::SUBMITTED_TO_WORKER: {
-    state_updates->set_submitted_to_worker_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::PENDING_NODE_ASSIGNMENT: {
-    state_updates->set_pending_node_assignment_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::FINISHED: {
-    state_updates->set_finished_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::FAILED: {
-    state_updates->set_failed_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::RUNNING: {
-    state_updates->set_running_ts(timestamp);
-    break;
-  }
-  case rpc::TaskStatus::NIL: {
+  if (task_status == rpc::TaskStatus::NIL) {
     // Not status change.
-    break;
-  }
-  default: {
-    UNREACHABLE;
-  }
+    return;
   }
+  (*state_updates->mutable_state_ts())[task_status] = timestamp;
 }
 
 inline std::string FormatPlacementGroupLabelName(const std::string &pg_id) {

src/ray/protobuf/gcs.proto

@@ -220,18 +220,6 @@ message TaskLogInfo {
 message TaskStateUpdate {
   // Node that runs the task.
   optional bytes node_id = 1;
-  // Timestamp when status changes to PENDING_ARGS_AVAIL.
-  optional int64 pending_args_avail_ts = 2;
-  // Timestamp when status changes to PENDING_NODE_ASSIGNMENT.
-  optional int64 pending_node_assignment_ts = 3;
-  // Timestamp when status changes to SUBMITTED_TO_WORKER.
-  optional int64 submitted_to_worker_ts = 4;
-  // Timestamp when status changes to RUNNING.
-  optional int64 running_ts = 5;
-  // Timestamp when status changes to FINISHED.
-  optional int64 finished_ts = 6;
-  // Timestamp when status changes to FAILED.
-  optional int64 failed_ts = 7;
   // Worker that runs the task.
   optional bytes worker_id = 8;
   // Task faulure info.
@@ -244,6 +232,9 @@ message TaskStateUpdate {
   optional int32 worker_pid = 12;
   // Is task paused by debugger.
   optional bool is_debugger_paused = 13;
+  // Key is the integer value of TaskStatus enum (protobuf doesn't support Enum as key).
+  // Value is the timestamp when status changes to the target status indicated by the key.
+  map<int32, int64> state_ts = 14;
 }
 
 // Represents events and state changes from a single task run.

src/ray/protobuf/gcs_service.proto

@@ -719,6 +719,8 @@ message GetTaskEventsRequest {
     optional string name = 4;
     // True if task events from driver (only profiling events) should be excluded.
     optional bool exclude_driver = 5;
+    // Latest state of the task.
+    optional string state = 6;
   }
 
   // Maximum number of TaskEvents to return.