[Stats] enable core worker stats (ray-project#9355)

BUILD.bazel

@@ -705,6 +705,7 @@ cc_library(
  ":ray_util",
  ":raylet_client_lib",
  ":service_based_gcs_client_lib",
+ ":stats_lib",
  ":worker_cc_proto",
  ":worker_rpc",
  "@boost//:fiber",

src/ray/common/ray_config_def.h

@@ -325,6 +325,10 @@ RAY_CONFIG(int64_t, enable_metrics_collection, true)
 /// Whether start the Plasma Store as a Raylet thread.
 RAY_CONFIG(bool, put_small_object_in_memory_store, false)
 
+/// Metric agent port for reporting, default -1 means no such agent will be
+/// listening.
+RAY_CONFIG(int, metrics_agent_port, -1)
+
 /// Maximum number of tasks that can be in flight between an owner and a worker for which
 /// the owner has been granted a lease. A value >1 is used when we want to enable
 /// pipelining task submission.

src/ray/core_worker/core_worker.cc

@@ -21,6 +21,7 @@
 #include "ray/core_worker/transport/direct_actor_transport.h"
 #include "ray/core_worker/transport/raylet_transport.h"
 #include "ray/gcs/gcs_client/service_based_gcs_client.h"
+#include "ray/stats/stats.h"
 #include "ray/util/process.h"
 #include "ray/util/util.h"
 
@@ -127,6 +128,19 @@ CoreWorkerProcess::CoreWorkerProcess(const CoreWorkerOptions &options)
  CreateWorker();
  }
  }
+
+ // Assume stats module will be initialized exactly once in once process.
+ // So it must be called in CoreWorkerProcess constructor and will be reused
+ // by all of core worker.
+ RAY_LOG(DEBUG) << "Stats setup in core worker.";
+ // Initialize stats in core worker global tags.
+ const ray::stats::TagsType global_tags = {{ray::stats::ComponentKey, "core_worker"},
+ {ray::stats::VersionKey, "0.9.0.dev0"}};
+
+ // NOTE(lingxuan.zlx): We assume RayConfig is initialized before it's used.
+ // RayConfig is generated in Java_io_ray_runtime_RayNativeRuntime_nativeInitialize
+ // for java worker or in constructor of CoreWorker for python worker.
+ ray::stats::Init(global_tags, RayConfig::instance().metrics_agent_port());
 }
 
 CoreWorkerProcess::~CoreWorkerProcess() {
@@ -136,6 +150,9 @@ CoreWorkerProcess::~CoreWorkerProcess() {
  absl::ReaderMutexLock lock(&worker_map_mutex_);
  RAY_CHECK(workers_.empty());
  }
+ RAY_LOG(DEBUG) << "Stats stop in core worker.";
+ // Shutdown stats module if worker process exits.
+ ray::stats::Shutdown();
  if (options_.enable_logging) {
  RayLog::ShutDownRayLog();
  }

src/ray/core_worker/core_worker.h

@@ -270,6 +270,7 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
  CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_id);
 
  CoreWorker(CoreWorker const &) = delete;
+
  void operator=(CoreWorker const &other) = delete;
 
  ///

src/ray/raylet/main.cc

@@ -215,6 +215,12 @@ int main(int argc, char *argv[]) {
  << object_manager_config.rpc_service_threads_number
  << ", object_chunk_size = "
  << object_manager_config.object_chunk_size;
+ // Initialize stats.
+ const ray::stats::TagsType global_tags = {
+ {ray::stats::ComponentKey, "raylet"},
+ {ray::stats::VersionKey, "0.9.0.dev0"},
+ {ray::stats::NodeAddressKey, node_ip_address}};
+ ray::stats::Init(global_tags, metrics_agent_port);
 
  // Initialize the node manager.
  server.reset(new ray::raylet::Raylet(
@@ -224,13 +230,6 @@ int main(int argc, char *argv[]) {
  server->Start();
  }));
 
- // Initialize stats.
- const ray::stats::TagsType global_tags = {
- {ray::stats::JobNameKey, "raylet"},
- {ray::stats::VersionKey, "0.9.0.dev0"},
- {ray::stats::NodeAddressKey, node_ip_address}};
- ray::stats::Init(global_tags, metrics_agent_port);
-
  // Destroy the Raylet on a SIGTERM. The pointer to main_service is
  // guaranteed to be valid since this function will run the event loop
  // instead of returning immediately.

src/ray/stats/stats.h

@@ -74,6 +74,7 @@ static inline void Init(const TagsType &global_tags, const int metrics_agent_por
  RAY_LOG(INFO) << "Disabled stats.";
  return;
  }
+ RAY_LOG(DEBUG) << "Initialized stats";
 
  metrics_io_service_pool = std::make_shared<IOServicePool>(1);
  metrics_io_service_pool->Run();
@@ -113,7 +114,6 @@ static inline void Shutdown() {
  exporter = nullptr;
  StatsConfig::instance().SetIsInitialized(false);
 }
-
 } // namespace stats
 
 } // namespace ray

src/ray/stats/tag_defs.h

@@ -20,6 +20,8 @@
 using TagKeyType = opencensus::tags::TagKey;
 using TagsType = std::vector<std::pair<opencensus::tags::TagKey, std::string>>;
 
+static const TagKeyType ComponentKey = TagKeyType::Register("Component");
+
 static const TagKeyType JobNameKey = TagKeyType::Register("JobName");
 
 static const TagKeyType CustomKey = TagKeyType::Register("CustomKey");