[core][metric/04] redefine STATS_tasks using Metric class

Signed-off-by: Cuong Nguyen <can@anyscale.com>

Author: can-anyscale

src/ray/core_worker/BUILD.bazel

@@ -28,13 +28,13 @@ ray_cc_library(
         ":generator_waiter",
         ":grpc_service",
         ":memory_store",
+        ":metrics",
         ":object_recovery_manager",
         ":plasma_store_provider",
         ":profile_event",
         ":reference_count",
         ":shutdown_coordinator",
         ":task_event_buffer",
-        ":metrics",
         "//src/ray/common:protobuf_utils",
         "//src/ray/common/cgroup:cgroup_context",
         "//src/ray/common/cgroup:cgroup_manager",

src/ray/core_worker/core_worker.cc

@@ -168,8 +168,8 @@ JobID GetProcessJobID(const CoreWorkerOptions &options) {
   return options.job_id;
 }
 
-TaskCounter::TaskCounter(ray::observability::MetricInterface &metric_tasks)
-    : metric_tasks_(metric_tasks) {
+TaskCounter::TaskCounter(ray::observability::MetricInterface &task_by_state_counter)
+    : task_by_state_counter_(task_by_state_counter) {
   counter_.SetOnChangeCallback(
       [this](const std::tuple<std::string, TaskStatusType, bool>
                  &key) ABSL_EXCLUSIVE_LOCKS_REQUIRED(&mu_) mutable {
@@ -184,31 +184,31 @@ TaskCounter::TaskCounter(ray::observability::MetricInterface &metric_tasks)
         const auto is_retry_label = is_retry ? "1" : "0";
         // RUNNING_IN_RAY_GET/WAIT are sub-states of RUNNING, so we need to subtract
         // them out to avoid double-counting.
-        metric_tasks_.Record(
+        task_by_state_counter_.Record(
             running_total - num_in_get - num_in_wait,
             {{"State"sv, rpc::TaskStatus_Name(rpc::TaskStatus::RUNNING)},
              {"Name"sv, func_name},
              {"IsRetry"sv, is_retry_label},
              {"JobId"sv, job_id_},
              {"Source"sv, "executor"}});
         // Negate the metrics recorded from the submitter process for these tasks.
-        metric_tasks_.Record(
+        task_by_state_counter_.Record(
             -running_total,
             {{"State"sv, rpc::TaskStatus_Name(rpc::TaskStatus::SUBMITTED_TO_WORKER)},
              {"Name"sv, func_name},
              {"IsRetry"sv, is_retry_label},
              {"JobId"sv, job_id_},
              {"Source"sv, "executor"}});
         // Record sub-state for get.
-        metric_tasks_.Record(
+        task_by_state_counter_.Record(
             num_in_get,
             {{"State"sv, rpc::TaskStatus_Name(rpc::TaskStatus::RUNNING_IN_RAY_GET)},
              {"Name"sv, func_name},
              {"IsRetry"sv, is_retry_label},
              {"JobId"sv, job_id_},
              {"Source"sv, "executor"}});
         // Record sub-state for wait.
-        metric_tasks_.Record(
+        task_by_state_counter_.Record(
             num_in_wait,
             {{"State"sv, rpc::TaskStatus_Name(rpc::TaskStatus::RUNNING_IN_RAY_WAIT)},
              {"Name"sv, func_name},
@@ -322,7 +322,7 @@ CoreWorker::CoreWorker(
     instrumented_io_context &task_execution_service,
     std::unique_ptr<worker::TaskEventBuffer> task_event_buffer,
     uint32_t pid,
-    ray::observability::MetricInterface &metric_tasks)
+    ray::observability::MetricInterface &task_by_state_counter)
     : options_(std::move(options)),
       get_call_site_(RayConfig::instance().record_ref_creation_sites()
                          ? options_.get_lang_stack
@@ -358,15 +358,14 @@ CoreWorker::CoreWorker(
       actor_id_(ActorID::Nil()),
       task_queue_length_(0),
       num_executed_tasks_(0),
-      task_counter_(metric_tasks),
       task_execution_service_(task_execution_service),
       exiting_detail_(std::nullopt),
       max_direct_call_object_size_(RayConfig::instance().max_direct_call_object_size()),
+      task_counter_(task_by_state_counter),
       task_event_buffer_(std::move(task_event_buffer)),
       pid_(pid),
       actor_shutdown_callback_(std::move(options_.actor_shutdown_callback)),
-      runtime_env_json_serialization_cache_(kDefaultSerializationCacheCap),
-      metric_tasks_(metric_tasks) {
+      runtime_env_json_serialization_cache_(kDefaultSerializationCacheCap) {
   // Initialize task receivers.
   if (options_.worker_type == WorkerType::WORKER || options_.is_local_mode) {
     RAY_CHECK(options_.task_execution_callback != nullptr);

src/ray/core_worker/core_worker.h

@@ -69,7 +69,7 @@ class TaskCounter {
   enum class TaskStatusType { kPending, kRunning, kFinished };
 
  public:
-  TaskCounter(ray::observability::MetricInterface &metric_tasks);
+  explicit TaskCounter(ray::observability::MetricInterface &task_by_state_counter);
 
   void BecomeActor(const std::string &actor_name) {
     absl::MutexLock l(&mu_);
@@ -128,8 +128,8 @@ class TaskCounter {
   std::string actor_name_ ABSL_GUARDED_BY(mu_);
   int64_t num_tasks_running_ ABSL_GUARDED_BY(mu_) = 0;
 
-  // Metrics
-  ray::observability::MetricInterface &metric_tasks_;
+  // Metric to track the number of tasks by state. Tags: State, Name, IsRetry, Source.
+  ray::observability::MetricInterface &task_by_state_counter_;
 };
 
 struct TaskToRetry {
@@ -194,7 +194,7 @@ class CoreWorker {
              instrumented_io_context &task_execution_service,
              std::unique_ptr<worker::TaskEventBuffer> task_event_buffer,
              uint32_t pid,
-             ray::observability::MetricInterface &metric_tasks);
+             ray::observability::MetricInterface &task_by_state_counter);
 
   CoreWorker(CoreWorker const &) = delete;
 
@@ -1892,6 +1892,8 @@ class CoreWorker {
 
   int64_t max_direct_call_object_size_;
 
+  friend class CoreWorkerTest;
+
   TaskCounter task_counter_;
 
   /// Used to guarantee that submitting actor task is thread safe.
@@ -1941,8 +1943,5 @@ class CoreWorker {
   std::mutex gcs_client_node_cache_populated_mutex_;
   std::condition_variable gcs_client_node_cache_populated_cv_;
   bool gcs_client_node_cache_populated_ = false;
-
-  // Metrics
-  ray::observability::MetricInterface &metric_tasks_;
 };
 }  // namespace ray::core

src/ray/core_worker/core_worker_process.cc

@@ -457,7 +457,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         return core_worker->core_worker_client_pool_->GetOrConnect(addr.value());
       },
       gcs_client,
-      metric_tasks_);
+      task_by_state_counter_);
 
   auto on_excess_queueing = [this](const ActorID &actor_id, uint64_t num_queued) {
     auto timestamp = std::chrono::duration_cast<std::chrono::seconds>(
@@ -662,7 +662,7 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
                                    task_execution_service_,
                                    std::move(task_event_buffer),
                                    pid,
-                                   metric_tasks_);
+                                   task_by_state_counter_);
   return core_worker;
 }
 

src/ray/core_worker/core_worker_process.h

@@ -184,8 +184,7 @@ class CoreWorkerProcessImpl {
   /// The client to export metrics to the metrics agent.
   std::unique_ptr<ray::rpc::MetricsAgentClient> metrics_agent_client_;
 
-  /// Metrics
-  ray::stats::Gauge metric_tasks_{GetTaskMetric()};
+  ray::stats::Gauge task_by_state_counter_{GetTaskMetric()};
 };
 }  // namespace core
 }  // namespace ray

src/ray/core_worker/metrics.h

@@ -0,0 +1,44 @@
+// Copyright 2025 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "ray/stats/metric.h"
+
+namespace ray {
+namespace core {
+
+inline ray::stats::Gauge GetTaskMetric() {
+  /// Tracks tasks by state, including pending, running, and finished tasks.
+  /// This metric may be recorded from multiple components processing the task in Ray,
+  /// including the submitting core worker, executor core worker, and pull manager.
+  ///
+  /// To avoid metric collection conflicts between components reporting on the same task,
+  /// we use the "Source" required label.
+  return ray::stats::Gauge{
+      /*name=*/"tasks",
+      /*description=*/"Current number of tasks currently in a particular state.",
+      /*unit=*/"",
+      // State: the task state, as described by rpc::TaskState proto in common.proto.
+      // Name: the name of the function called (Keep in sync with the
+      // TASK_OR_ACTOR_NAME_TAG_KEY in
+      // python/ray/_private/telemetry/metric_cardinality.py) Source: component reporting,
+      // e.g., "core_worker", "executor", or "pull_manager". IsRetry: whether this task is
+      // a retry.
+      /*tag_keys=*/{"State", "Name", "Source", "IsRetry", "JobId"},
+  };
+}
+
+}  // namespace core
+}  // namespace ray

src/ray/core_worker/task_manager.h

@@ -31,12 +31,12 @@
 #include "ray/core_worker/task_event_buffer.h"
 #include "ray/core_worker/task_manager_interface.h"
 #include "ray/gcs/gcs_client/gcs_client.h"
+#include "ray/observability/metric_interface.h"
 #include "ray/stats/metric_defs.h"
 #include "ray/util/counter_map.h"
 #include "src/ray/protobuf/common.pb.h"
 #include "src/ray/protobuf/core_worker.pb.h"
 #include "src/ray/protobuf/gcs.pb.h"
-#include "ray/observability/metric_interface.h"
 
 namespace ray {
 namespace core {
@@ -186,7 +186,7 @@ class TaskManager : public TaskManagerInterface {
       std::function<std::shared_ptr<ray::rpc::CoreWorkerClientInterface>(const ActorID &)>
           client_factory,
       std::shared_ptr<gcs::GcsClient> gcs_client,
-      ray::observability::MetricInterface &metric_tasks)
+      ray::observability::MetricInterface &task_by_state_counter)
       : in_memory_store_(in_memory_store),
         reference_counter_(reference_counter),
         put_in_local_plasma_callback_(std::move(put_in_local_plasma_callback)),
@@ -197,16 +197,16 @@ class TaskManager : public TaskManagerInterface {
         task_event_buffer_(task_event_buffer),
         get_actor_rpc_client_callback_(std::move(client_factory)),
         gcs_client_(std::move(gcs_client)),
-        metric_tasks_(metric_tasks) {
+        task_by_state_counter_(task_by_state_counter) {
     task_counter_.SetOnChangeCallback(
         [this](const std::tuple<std::string, rpc::TaskStatus, bool> &key)
             ABSL_EXCLUSIVE_LOCKS_REQUIRED(&mu_) {
-              metric_tasks_.Record(
+              task_by_state_counter_.Record(
                   task_counter_.Get(key),
                   {{"State"sv, rpc::TaskStatus_Name(std::get<1>(key))},
                    {"Name"sv, std::get<0>(key)},
                    {"IsRetry"sv, std::get<2>(key) ? "1" : "0"},
-                   {"Source"sv, "owner"sv}});
+                   {"Source"sv, "owner"}});
             });
     reference_counter_.SetReleaseLineageCallback(
         [this](const ObjectID &object_id, std::vector<ObjectID> *ids_to_release) {
@@ -802,8 +802,8 @@ class TaskManager : public TaskManagerInterface {
 
   std::shared_ptr<gcs::GcsClient> gcs_client_;
 
-  // Metrics
-  ray::observability::MetricInterface &metric_tasks_;
+  // Metric to track the number of tasks by state. Tags: State, Name, IsRetry, Source.
+  observability::MetricInterface &task_by_state_counter_;
 
   friend class TaskManagerTest;
 };

src/ray/core_worker/task_submission/actor_task_submitter.h

@@ -440,6 +440,8 @@ class ActorTaskSubmitter : public ActorTaskSubmitterInterface {
   instrumented_io_context &io_service_;
 
   std::shared_ptr<ReferenceCounterInterface> reference_counter_;
+
+  friend class CoreWorkerTest;
 };
 
 }  // namespace core

src/ray/core_worker/tests/BUILD.bazel

@@ -98,6 +98,7 @@ ray_cc_test(
         "//src/ray/core_worker:task_event_buffer",
         "//src/ray/core_worker:task_manager",
         "//src/ray/gcs/gcs_client:gcs_client_lib",
+        "//src/ray/observability:fake_metric",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
     ],
@@ -253,6 +254,7 @@ ray_cc_test(
         "//src/ray/core_worker:memory_store",
         "//src/ray/core_worker:reference_count",
         "//src/ray/ipc:fake_raylet_ipc_client",
+        "//src/ray/observability:fake_metric",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
     ],

src/ray/core_worker/tests/core_worker_test.cc

@@ -41,6 +41,7 @@
 #include "ray/core_worker/task_submission/actor_task_submitter.h"
 #include "ray/core_worker/task_submission/normal_task_submitter.h"
 #include "ray/ipc/fake_raylet_ipc_client.h"
+#include "ray/observability/fake_metric.h"
 #include "ray/rpc/worker/core_worker_client_pool.h"
 
 namespace ray {
@@ -160,7 +161,8 @@ class CoreWorkerTest : public ::testing::Test {
         [](const ActorID &actor_id) {
           return std::make_shared<rpc::CoreWorkerClientInterface>();
         },
-        mock_gcs_client);
+        mock_gcs_client,
+        fake_task_by_state_counter_);
 
     auto object_recovery_manager = std::make_unique<ObjectRecoveryManager>(
         rpc_address_,
@@ -245,7 +247,8 @@ class CoreWorkerTest : public ::testing::Test {
                                                 std::move(actor_manager),
                                                 task_execution_service_,
                                                 std::move(task_event_buffer),
-                                                getpid());
+                                                getpid(),
+                                                fake_task_by_state_counter_);
   }
 
  protected:
@@ -263,6 +266,7 @@ class CoreWorkerTest : public ::testing::Test {
   ActorTaskSubmitter *actor_task_submitter_;
   std::shared_ptr<TaskManager> task_manager_;
   std::shared_ptr<CoreWorker> core_worker_;
+  ray::observability::FakeMetric fake_task_by_state_counter_;
 };
 
 std::shared_ptr<RayObject> MakeRayObject(const std::string &data_str,
@@ -278,6 +282,22 @@ std::shared_ptr<RayObject> MakeRayObject(const std::string &data_str,
   return std::make_shared<RayObject>(data, metadata, std::vector<rpc::ObjectReference>());
 }
 
+TEST_F(CoreWorkerTest, RecordMetrics) {
+  std::vector<std::shared_ptr<RayObject>> results;
+  auto status = core_worker_->Get({}, -1, results);
+  ASSERT_TRUE(status.ok());
+  // disconnect to trigger metric recording
+  core_worker_->Disconnect(rpc::WorkerExitType::SYSTEM_ERROR, "test", nullptr);
+  auto tag_to_value = fake_task_by_state_counter_.GetTagToValue();
+  // 4 states: RUNNING, SUBMITTED_TO_WORKER, RUNNING_IN_RAY_GET and RUNNING_IN_RAY_WAIT
+  ASSERT_EQ(tag_to_value.size(), 4);
+  for (auto &[key, value] : tag_to_value) {
+    ASSERT_EQ(key.at("Name"), "Unknown task");
+    ASSERT_EQ(key.at("Source"), "executor");
+    ASSERT_EQ(key.at("IsRetry"), "0");
+  }
+}
+
 TEST_F(CoreWorkerTest, HandleGetObjectStatusIdempotency) {
   auto object_id = ObjectID::FromRandom();
   auto ray_object = MakeRayObject("test_data", "meta");