[core][stats-die/03] kill STATS in core worker component

Signed-off-by: Cuong Nguyen <[email protected]>

Author: can-anyscale

src/ray/common/BUILD.bazel

@@ -234,6 +234,7 @@ ray_cc_library(
         "//src/ray/common/scheduling:resource_set",
         "//src/ray/common/scheduling:scheduling_class_util",
         "//src/ray/flatbuffers:node_manager_generated",
+        "//src/ray/observability:metric_interface",
         "//src/ray/util:container_util",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -425,6 +426,6 @@ ray_cc_library(
     name = "metrics",
     hdrs = ["metrics.h"],
     deps = [
-        "//src/ray/stats:stats_lib",
+        "//src/ray/stats:stats_metric",
     ],
 )

src/ray/common/metrics.h

@@ -73,4 +73,17 @@ inline ray::stats::Gauge GetObjectStoreMemoryGaugeMetric() {
   };
 }
 
+inline ray::stats::Histogram GetSchedulerPlacementTimeSHistogramMetric() {
+  return ray::stats::Histogram{
+      /*name=*/"scheduler_placement_time_s",
+      /*description=*/
+      "The time it takes for a worklod (task, actor, placement group) to "
+      "be placed. This is the time from when the tasks dependencies are "
+      "resolved to when it actually reserves resources on a node to run.",
+      /*unit=*/"s",
+      /*boundaries=*/{0.1, 1, 10, 100, 1000, 10000},
+      /*tag_keys=*/{"WorkloadType"},
+  };
+}
+
 }  // namespace ray

src/ray/common/task/task_spec.cc

@@ -23,7 +23,6 @@
 
 #include "ray/common/ray_config.h"
 #include "ray/common/runtime_env_common.h"
-#include "ray/stats/metric_defs.h"
 #include "ray/util/logging.h"
 
 namespace ray {
@@ -602,17 +601,16 @@ bool TaskSpecification::IsRetriable() const {
   return true;
 }
 
-void TaskSpecification::EmitTaskMetrics() const {
+void TaskSpecification::EmitTaskMetrics(
+    ray::observability::MetricInterface &scheduler_placement_time_s_histogram) const {
   double duration_s = (GetMessage().lease_grant_timestamp_ms() -
                        GetMessage().dependency_resolution_timestamp_ms()) /
                       1000;
 
   if (IsActorCreationTask()) {
-    stats::STATS_scheduler_placement_time_s.Record(duration_s,
-                                                   {{"WorkloadType", "Actor"}});
+    scheduler_placement_time_s_histogram.Record(duration_s, {{"WorkloadType", "Actor"}});
   } else {
-    stats::STATS_scheduler_placement_time_s.Record(duration_s,
-                                                   {{"WorkloadType", "Task"}});
+    scheduler_placement_time_s_histogram.Record(duration_s, {{"WorkloadType", "Task"}});
   }
 }
 

src/ray/common/task/task_spec.h

@@ -29,6 +29,7 @@
 #include "ray/common/scheduling/resource_set.h"
 #include "ray/common/scheduling/scheduling_class_util.h"
 #include "ray/common/task/task_common.h"
+#include "ray/observability/metric_interface.h"
 
 extern "C" {
 #include "ray/thirdparty/sha256.h"
@@ -357,7 +358,8 @@ class TaskSpecification : public MessageWrapper<rpc::TaskSpec> {
   /// \return true if the task or actor is retriable.
   bool IsRetriable() const;
 
-  void EmitTaskMetrics() const;
+  void EmitTaskMetrics(
+      ray::observability::MetricInterface &scheduler_placement_time_s_histogram) const;
 
   /// \return true if task events from this task should be reported.
   bool EnableTaskEvents() const;

src/ray/core_worker/core_worker_process.cc

@@ -470,7 +470,8 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         return core_worker->core_worker_client_pool_->GetOrConnect(*addr);
       },
       gcs_client,
-      task_by_state_gauge_,
+      *task_by_state_gauge_,
+      *total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/
       [this](const ObjectID &object_id) {
         auto core_worker = GetCoreWorker();
@@ -552,7 +553,8 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
         // OBJECT_STORE.
         return rpc::TensorTransport::OBJECT_STORE;
       },
-      boost::asio::steady_timer(io_service_));
+      boost::asio::steady_timer(io_service_),
+      *scheduler_placement_time_s_histogram_);
 
   auto report_locality_data_callback = [this](
                                            const ObjectID &object_id,
@@ -683,8 +685,8 @@ std::shared_ptr<CoreWorker> CoreWorkerProcessImpl::CreateCoreWorker(
                                    task_execution_service_,
                                    std::move(task_event_buffer),
                                    pid,
-                                   task_by_state_gauge_,
-                                   actor_by_state_gauge_);
+                                   *task_by_state_gauge_,
+                                   *actor_by_state_gauge_);
   return core_worker;
 }
 
@@ -790,6 +792,14 @@ CoreWorkerProcessImpl::CoreWorkerProcessImpl(const CoreWorkerOptions &options)
 
   // We need init stats before using it/spawning threads.
   stats::Init(global_tags, options_.metrics_agent_port, worker_id_);
+  task_by_state_gauge_ = std::unique_ptr<ray::stats::Gauge>(
+      new ray::stats::Gauge(GetTaskByStateGaugeMetric()));
+  actor_by_state_gauge_ = std::unique_ptr<ray::stats::Gauge>(
+      new ray::stats::Gauge(GetActorByStateGaugeMetric()));
+  total_lineage_bytes_gauge_ = std::unique_ptr<ray::stats::Gauge>(
+      new ray::stats::Gauge(GetTotalLineageBytesGaugeMetric()));
+  scheduler_placement_time_s_histogram_ = std::unique_ptr<ray::stats::Histogram>(
+      new ray::stats::Histogram(GetSchedulerPlacementTimeSHistogramMetric()));
 
   // Initialize event framework before starting up worker.
   if (RayConfig::instance().event_log_reporter_enabled() && !options_.log_dir.empty()) {

src/ray/core_worker/core_worker_process.h

@@ -185,8 +185,10 @@ class CoreWorkerProcessImpl {
   /// The client to export metrics to the metrics agent.
   std::unique_ptr<ray::rpc::MetricsAgentClient> metrics_agent_client_;
 
-  ray::stats::Gauge task_by_state_gauge_{GetTaskByStateGaugeMetric()};
-  ray::stats::Gauge actor_by_state_gauge_{GetActorByStateGaugeMetric()};
+  std::unique_ptr<ray::stats::Gauge> task_by_state_gauge_;
+  std::unique_ptr<ray::stats::Gauge> actor_by_state_gauge_;
+  std::unique_ptr<ray::stats::Gauge> total_lineage_bytes_gauge_;
+  std::unique_ptr<ray::stats::Histogram> scheduler_placement_time_s_histogram_;
 };
 }  // namespace core
 }  // namespace ray

src/ray/core_worker/metrics.h

@@ -41,5 +41,15 @@ inline ray::stats::Gauge GetTaskByStateGaugeMetric() {
   };
 }
 
+inline ray::stats::Gauge GetTotalLineageBytesGaugeMetric() {
+  return ray::stats::Gauge{
+      /*name=*/"total_lineage_bytes",
+      /*description=*/
+      "Total amount of memory used to store task specs for lineage reconstruction.",
+      /*unit=*/"",
+      /*tag_keys=*/{},
+  };
+}
+
 }  // namespace core
 }  // namespace ray

src/ray/core_worker/task_manager.cc

@@ -1781,7 +1781,7 @@ void TaskManager::FillTaskInfo(rpc::GetCoreWorkerStatsReply *reply,
 
 void TaskManager::RecordMetrics() {
   absl::MutexLock lock(&mu_);
-  ray::stats::STATS_total_lineage_bytes.Record(total_lineage_footprint_bytes_);
+  total_lineage_bytes_gauge_.Record(total_lineage_footprint_bytes_);
   task_counter_.FlushOnChangeCallbacks();
 }
 

src/ray/core_worker/task_manager.h

@@ -33,7 +33,6 @@
 #include "ray/core_worker_rpc_client/core_worker_client_interface.h"
 #include "ray/gcs_rpc_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"
@@ -187,6 +186,7 @@ class TaskManager : public TaskManagerInterface {
           const ActorID &)> get_actor_rpc_client_callback,
       std::shared_ptr<gcs::GcsClient> gcs_client,
       ray::observability::MetricInterface &task_by_state_counter,
+      ray::observability::MetricInterface &total_lineage_bytes_gauge,
       FreeActorObjectCallback free_actor_object_callback)
       : in_memory_store_(in_memory_store),
         reference_counter_(reference_counter),
@@ -199,6 +199,7 @@ class TaskManager : public TaskManagerInterface {
         get_actor_rpc_client_callback_(std::move(get_actor_rpc_client_callback)),
         gcs_client_(std::move(gcs_client)),
         task_by_state_counter_(task_by_state_counter),
+        total_lineage_bytes_gauge_(total_lineage_bytes_gauge),
         free_actor_object_callback_(std::move(free_actor_object_callback)) {
     task_counter_.SetOnChangeCallback(
         [this](const std::tuple<std::string, rpc::TaskStatus, bool> &key)
@@ -812,6 +813,10 @@ class TaskManager : public TaskManagerInterface {
   // - Source: component reporting, e.g., "core_worker", "executor", or "pull_manager"
   observability::MetricInterface &task_by_state_counter_;
 
+  /// Metric to track the total amount of memory used to store task specs for lineage
+  /// reconstruction.
+  observability::MetricInterface &total_lineage_bytes_gauge_;
+
   /// Callback to free GPU object from the in-actor object store.
   FreeActorObjectCallback free_actor_object_callback_;
 

src/ray/core_worker/task_submission/normal_task_submitter.cc

@@ -177,7 +177,7 @@ void NormalTaskSubmitter::OnWorkerIdle(
       scheduling_key_entry.num_busy_workers++;
 
       task_spec.GetMutableMessage().set_lease_grant_timestamp_ms(current_sys_time_ms());
-      task_spec.EmitTaskMetrics();
+      task_spec.EmitTaskMetrics(scheduler_placement_time_s_histogram_);
 
       executing_tasks_.emplace(task_spec.TaskId(), addr);
       PushNormalTask(