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

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(

src/ray/core_worker/task_submission/normal_task_submitter.h

@@ -95,7 +95,8 @@ class NormalTaskSubmitter {
       const JobID &job_id,
       std::shared_ptr<LeaseRequestRateLimiter> lease_request_rate_limiter,
       const TensorTransportGetter &tensor_transport_getter,
-      boost::asio::steady_timer cancel_timer)
+      boost::asio::steady_timer cancel_timer,
+      ray::observability::MetricInterface &scheduler_placement_time_s_histogram)
       : rpc_address_(std::move(rpc_address)),
         local_raylet_client_(std::move(local_raylet_client)),
         raylet_client_pool_(std::move(raylet_client_pool)),
@@ -109,7 +110,8 @@ class NormalTaskSubmitter {
         core_worker_client_pool_(std::move(core_worker_client_pool)),
         job_id_(job_id),
         lease_request_rate_limiter_(std::move(lease_request_rate_limiter)),
-        cancel_retry_timer_(std::move(cancel_timer)) {}
+        cancel_retry_timer_(std::move(cancel_timer)),
+        scheduler_placement_time_s_histogram_(scheduler_placement_time_s_histogram) {}
 
   /// Schedule a task for direct submission to a worker.
   void SubmitTask(TaskSpecification task_spec);
@@ -362,6 +364,8 @@ class NormalTaskSubmitter {
 
   // Retries cancelation requests if they were not successful.
   boost::asio::steady_timer cancel_retry_timer_ ABSL_GUARDED_BY(mu_);
+
+  ray::observability::MetricInterface &scheduler_placement_time_s_histogram_;
 };
 
 }  // namespace core

src/ray/core_worker/task_submission/tests/normal_task_submitter_test.cc

@@ -32,6 +32,7 @@
 #include "ray/core_worker/store_provider/memory_store/memory_store.h"
 #include "ray/core_worker_rpc_client/core_worker_client_pool.h"
 #include "ray/core_worker_rpc_client/fake_core_worker_client.h"
+#include "ray/observability/fake_metric.h"
 #include "ray/raylet_rpc_client/fake_raylet_client.h"
 #include "ray/raylet_rpc_client/raylet_client_interface.h"
 
@@ -494,7 +495,8 @@ class NormalTaskSubmitterTest : public testing::Test {
         JobID::Nil(),
         rate_limiter,
         [](const ObjectID &object_id) { return rpc::TensorTransport::OBJECT_STORE; },
-        boost::asio::steady_timer(io_context));
+        boost::asio::steady_timer(io_context),
+        fake_scheduler_placement_time_s_histogram_);
   }
 
   NodeID local_node_id;
@@ -511,6 +513,7 @@ class NormalTaskSubmitterTest : public testing::Test {
   std::unique_ptr<MockLeasePolicy> lease_policy;
   MockLeasePolicy *lease_policy_ptr = nullptr;
   instrumented_io_context io_context;
+  ray::observability::FakeHistogram fake_scheduler_placement_time_s_histogram_;
 };
 
 TEST_F(NormalTaskSubmitterTest, TestLocalityAwareSubmitOneTask) {
@@ -1430,6 +1433,7 @@ void TestSchedulingKey(const std::shared_ptr<CoreWorkerMemoryStore> store,
                        const TaskSpecification &same2,
                        const TaskSpecification &different) {
   rpc::Address address;
+  ray::observability::FakeHistogram fake_scheduler_placement_time_s_histogram_;
   auto local_node_id = NodeID::FromRandom();
   auto raylet_client = std::make_shared<MockRayletClient>();
   auto raylet_client_pool = std::make_shared<rpc::RayletClientPool>(
@@ -1457,7 +1461,8 @@ void TestSchedulingKey(const std::shared_ptr<CoreWorkerMemoryStore> store,
       JobID::Nil(),
       std::make_shared<StaticLeaseRequestRateLimiter>(1),
       [](const ObjectID &object_id) { return rpc::TensorTransport::OBJECT_STORE; },
-      boost::asio::steady_timer(io_context));
+      boost::asio::steady_timer(io_context),
+      fake_scheduler_placement_time_s_histogram_);
 
   submitter.SubmitTask(same1);
   submitter.SubmitTask(same2);

src/ray/core_worker/tests/core_worker_test.cc

@@ -187,6 +187,7 @@ class CoreWorkerTest : public ::testing::Test {
         },
         mock_gcs_client_,
         fake_task_by_state_gauge_,
+        fake_total_lineage_bytes_gauge_,
         /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
     auto object_recovery_manager = std::make_unique<ObjectRecoveryManager>(
@@ -222,7 +223,8 @@ class CoreWorkerTest : public ::testing::Test {
         JobID::Nil(),
         lease_request_rate_limiter,
         [](const ObjectID &object_id) { return rpc::TensorTransport::OBJECT_STORE; },
-        boost::asio::steady_timer(io_service_));
+        boost::asio::steady_timer(io_service_),
+        fake_scheduler_placement_time_s_histogram_);
 
     auto actor_task_submitter = std::make_unique<ActorTaskSubmitter>(
         *core_worker_client_pool,
@@ -297,6 +299,8 @@ class CoreWorkerTest : public ::testing::Test {
   std::shared_ptr<CoreWorker> core_worker_;
   ray::observability::FakeGauge fake_task_by_state_gauge_;
   ray::observability::FakeGauge fake_actor_by_state_gauge_;
+  ray::observability::FakeGauge fake_total_lineage_bytes_gauge_;
+  ray::observability::FakeHistogram fake_scheduler_placement_time_s_histogram_;
   std::unique_ptr<FakePeriodicalRunner> fake_periodical_runner_;
 
   // Controllable time for testing publisher timeouts

src/ray/core_worker/tests/task_manager_test.cc

@@ -183,6 +183,7 @@ class TaskManagerTest : public ::testing::Test {
             },
             mock_gcs_client_,
             fake_task_by_state_counter_,
+            fake_total_lineage_bytes_gauge_,
             /*free_actor_object_callback=*/[](const ObjectID &object_id) {}) {}
 
   virtual void TearDown() { AssertNoLeaks(); }
@@ -230,6 +231,7 @@ class TaskManagerTest : public ::testing::Test {
   uint32_t last_delay_ms_ = 0;
   std::unordered_set<ObjectID> stored_in_plasma;
   ray::observability::FakeGauge fake_task_by_state_counter_;
+  ray::observability::FakeGauge fake_total_lineage_bytes_gauge_;
 };
 
 class TaskManagerLineageTest : public TaskManagerTest {
@@ -1404,6 +1406,7 @@ TEST_F(TaskManagerTest, PlasmaPut_ObjectStoreFull_FailsTaskAndWritesError) {
       },
       mock_gcs_client_,
       fake_task_by_state_counter_,
+      fake_total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
   rpc::Address caller_address;
@@ -1467,6 +1470,7 @@ TEST_F(TaskManagerTest, PlasmaPut_TransientFull_RetriesThenSucceeds) {
       },
       mock_gcs_client_,
       fake_task_by_state_counter_,
+      fake_total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
   rpc::Address caller_address;
@@ -1528,6 +1532,7 @@ TEST_F(TaskManagerTest, DynamicReturn_PlasmaPutFailure_FailsTaskImmediately) {
       },
       mock_gcs_client_,
       fake_task_by_state_counter_,
+      fake_total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
   auto spec = CreateTaskHelper(1, {}, /*dynamic_returns=*/true);
@@ -3016,6 +3021,7 @@ TEST_F(TaskManagerTest, TestRetryErrorMessageSentToCallback) {
           -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> { return nullptr; },
       mock_gcs_client_,
       fake_task_by_state_counter_,
+      fake_total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
   // Create a task with retries enabled
@@ -3096,6 +3102,7 @@ TEST_F(TaskManagerTest, TestErrorLogWhenPushErrorCallbackFails) {
           -> std::shared_ptr<ray::rpc::CoreWorkerClientInterface> { return nullptr; },
       mock_gcs_client_,
       fake_task_by_state_counter_,
+      fake_total_lineage_bytes_gauge_,
       /*free_actor_object_callback=*/[](const ObjectID &object_id) {});
 
   // Create a task that will be retried

src/ray/gcs/BUILD.bazel

@@ -273,6 +273,7 @@ ray_cc_library(
     deps = [
         "//src/ray/common:bundle_spec",
         "//src/ray/common:id",
+        "//src/ray/common:metrics",
         "//src/ray/protobuf:gcs_service_cc_proto",
         "//src/ray/stats:stats_lib",
         "//src/ray/util:counter_map",

src/ray/gcs/gcs_actor_scheduler.cc

@@ -35,6 +35,7 @@ GcsActorScheduler::GcsActorScheduler(
     GcsActorSchedulerSuccessCallback schedule_success_handler,
     rpc::RayletClientPool &raylet_client_pool,
     rpc::CoreWorkerClientPool &worker_client_pool,
+    ray::observability::MetricInterface &scheduler_placement_time_s_histogram,
     std::function<void(const NodeID &, const rpc::ResourcesData &)>
         normal_task_resources_changed_callback)
     : io_context_(io_context),
@@ -45,6 +46,7 @@ GcsActorScheduler::GcsActorScheduler(
       schedule_success_handler_(std::move(schedule_success_handler)),
       raylet_client_pool_(raylet_client_pool),
       worker_client_pool_(worker_client_pool),
+      scheduler_placement_time_s_histogram_(scheduler_placement_time_s_histogram),
       normal_task_resources_changed_callback_(
           std::move(normal_task_resources_changed_callback)) {
   RAY_CHECK(schedule_failure_handler_ != nullptr && schedule_success_handler_ != nullptr);
@@ -393,7 +395,8 @@ void GcsActorScheduler::HandleWorkerLeaseGrantedReply(
     actor->UpdateAddress(leased_worker->GetAddress());
     actor->GetMutableActorTableData()->set_pid(reply.worker_pid());
     actor->GetMutableTaskSpec()->set_lease_grant_timestamp_ms(current_sys_time_ms());
-    actor->GetCreationTaskSpecification().EmitTaskMetrics();
+    actor->GetCreationTaskSpecification().EmitTaskMetrics(
+        scheduler_placement_time_s_histogram_);
     // Make sure to connect to the client before persisting actor info to GCS.
     // Without this, there could be a possible race condition. Related issues:
     // https://github.com/ray-project/ray/pull/9215/files#r449469320