[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

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

src/ray/core_worker/core_worker.cc

@@ -167,8 +167,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 {
@@ -183,31 +183,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},
@@ -321,7 +321,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
@@ -357,15 +357,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

@@ -80,7 +80,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_);
@@ -138,9 +138,7 @@ class TaskCounter {
   // Used for actor state tracking.
   std::string actor_name_ ABSL_GUARDED_BY(mu_);
   int64_t num_tasks_running_ ABSL_GUARDED_BY(mu_) = 0;
-
-  // Metrics
-  ray::observability::MetricInterface &metric_tasks_;
+  ray::observability::MetricInterface &task_by_state_counter_;
 };
 
 struct TaskToRetry {
@@ -205,7 +203,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;
 
@@ -1954,8 +1952,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

@@ -455,7 +455,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>(
@@ -660,7 +660,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

@@ -183,8 +183,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) {
@@ -803,7 +803,7 @@ class TaskManager : public TaskManagerInterface {
   std::shared_ptr<gcs::GcsClient> gcs_client_;
 
   // Metrics
-  ray::observability::MetricInterface &metric_tasks_;
+  ray::observability::MetricInterface &task_by_state_counter_;
 
   friend class TaskManagerTest;
 };

src/ray/core_worker/tests/BUILD.bazel

@@ -97,6 +97,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",
     ],
@@ -251,6 +252,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

@@ -40,6 +40,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 {
@@ -159,7 +160,8 @@ class CoreWorkerHandleGetObjectStatusTest : 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,
@@ -243,7 +245,8 @@ class CoreWorkerHandleGetObjectStatusTest : public ::testing::Test {
                                                 std::move(actor_manager),
                                                 task_execution_service_,
                                                 std::move(task_event_buffer),
-                                                getpid());
+                                                getpid(),
+                                                fake_task_by_state_counter_);
   }
 
  protected:
@@ -260,6 +263,7 @@ class CoreWorkerHandleGetObjectStatusTest : public ::testing::Test {
   std::shared_ptr<ReferenceCounter> reference_counter_;
   std::shared_ptr<CoreWorkerMemoryStore> memory_store_;
   std::shared_ptr<CoreWorker> core_worker_;
+  ray::observability::FakeMetric fake_task_by_state_counter_;
 };
 
 std::shared_ptr<RayObject> MakeRayObject(const std::string &data_str,

src/ray/core_worker/tests/task_manager_test.cc

@@ -31,6 +31,7 @@
 #include "ray/core_worker/reference_count.h"
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 #include "ray/core_worker/task_event_buffer.h"
+#include "ray/observability/fake_metric.h"
 
 namespace ray {
 namespace core {
@@ -181,7 +182,8 @@ class TaskManagerTest : public ::testing::Test {
                 -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> {
               return nullptr;
             },
-            mock_gcs_client_) {}
+            mock_gcs_client_,
+            fake_task_by_state_counter_) {}
 
   virtual void TearDown() { AssertNoLeaks(); }
 
@@ -228,6 +230,7 @@ class TaskManagerTest : public ::testing::Test {
   uint32_t last_delay_ms_ = 0;
   bool last_object_recovery_ = false;
   std::unordered_set<ObjectID> stored_in_plasma;
+  ray::observability::FakeMetric fake_task_by_state_counter_;
 };
 
 class TaskManagerLineageTest : public TaskManagerTest {
@@ -1396,7 +1399,8 @@ TEST_F(TaskManagerTest, PlasmaPut_ObjectStoreFull_FailsTaskAndWritesError) {
       [](const ActorID &) -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> {
         return nullptr;
       },
-      mock_gcs_client_);
+      mock_gcs_client_,
+      fake_task_by_state_counter_);
 
   rpc::Address caller_address;
   auto spec = CreateTaskHelper(1, {});
@@ -1459,7 +1463,8 @@ TEST_F(TaskManagerTest, PlasmaPut_TransientFull_RetriesThenSucceeds) {
       [](const ActorID &) -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> {
         return nullptr;
       },
-      mock_gcs_client_);
+      mock_gcs_client_,
+      fake_task_by_state_counter_);
 
   rpc::Address caller_address;
   auto spec = CreateTaskHelper(1, {});
@@ -1520,7 +1525,8 @@ TEST_F(TaskManagerTest, DynamicReturn_PlasmaPutFailure_FailsTaskImmediately) {
       [](const ActorID &) -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> {
         return nullptr;
       },
-      mock_gcs_client_);
+      mock_gcs_client_,
+      fake_task_by_state_counter_);
 
   auto spec = CreateTaskHelper(1, {}, /*dynamic_returns=*/true);
   dyn_mgr.AddPendingTask(addr_, spec, "", /*num_retries=*/0);