diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index b2196c4f3..b8fd80fa0 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -78,7 +78,6 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, client_call_manager_(new rpc::ClientCallManager(io_service_)), heartbeat_timer_(io_service_), core_worker_server_(WorkerTypeString(worker_type), 0 /* let grpc choose a port */), - gcs_client_(gcs_options), memory_store_(std::make_shared()), memory_store_provider_(memory_store_), task_execution_service_work_(task_execution_service_), @@ -95,7 +94,8 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, } // Initialize gcs client. - RAY_CHECK_OK(gcs_client_.Connect(io_service_)); + gcs_client_ = std::make_shared(gcs_options); + RAY_CHECK_OK(gcs_client_->Connect(io_service_)); // Initialize profiler. profiler_ = std::make_shared(worker_context_, node_ip_address, @@ -171,7 +171,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, std::shared_ptr data = std::make_shared(); data->mutable_task()->mutable_task_spec()->CopyFrom(builder.Build().GetMessage()); - RAY_CHECK_OK(gcs_client_.raylet_task_table().Add(job_id, task_id, data, nullptr)); + RAY_CHECK_OK(gcs_client_->raylet_task_table().Add(job_id, task_id, data, nullptr)); SetCurrentTaskId(task_id); } @@ -209,7 +209,7 @@ void CoreWorker::Shutdown() { void CoreWorker::Disconnect() { io_service_.stop(); - gcs_client_.Disconnect(); + gcs_client_->Disconnect(); if (raylet_client_) { RAY_IGNORE_EXPR(raylet_client_->Disconnect()); } @@ -232,7 +232,7 @@ void CoreWorker::SetCurrentTaskId(const TaskID &task_id) { // Clear all actor handles at the end of each non-actor task. if (actor_id_.IsNil() && task_id.IsNil()) { for (const auto &handle : actor_handles_) { - RAY_CHECK_OK(gcs_client_.Actors().AsyncUnsubscribe(handle.first, nullptr)); + RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(handle.first, nullptr)); } actor_handles_.clear(); } @@ -642,7 +642,7 @@ bool CoreWorker::AddActorHandle(std::unique_ptr actor_handle) { it->second->Reset(); } } else if (actor_data.state() == gcs::ActorTableData::DEAD) { - RAY_CHECK_OK(gcs_client_.Actors().AsyncUnsubscribe(actor_id, nullptr)); + RAY_CHECK_OK(gcs_client_->Actors().AsyncUnsubscribe(actor_id, nullptr)); // We cannot erase the actor handle here because clients can still // submit tasks to dead actors. } @@ -655,7 +655,7 @@ bool CoreWorker::AddActorHandle(std::unique_ptr actor_handle) { << ", port: " << actor_data.port(); }; - RAY_CHECK_OK(gcs_client_.Actors().AsyncSubscribe( + RAY_CHECK_OK(gcs_client_->Actors().AsyncSubscribe( actor_id, actor_notification_callback, nullptr)); } return inserted; diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 24ebe42c4..386e8f52d 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -1,9 +1,7 @@ #ifndef RAY_CORE_WORKER_CORE_WORKER_H #define RAY_CORE_WORKER_CORE_WORKER_H -#include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" -#include "absl/synchronization/mutex.h" #include "ray/common/buffer.h" #include "ray/core_worker/actor_handle.h" @@ -321,11 +319,11 @@ class CoreWorker { const std::vector> &metadatas, std::vector> *return_objects); - /** - * The following methods are handlers for the core worker's gRPC server, which follow - * a macro-generated call convention. These are executed on the io_service_ and - * post work to the appropriate event loop. - */ + /// + /// The following methods are handlers for the core worker's gRPC server, which follow + /// a macro-generated call convention. These are executed on the io_service_ and + /// post work to the appropriate event loop. + /// /// Implements gRPC server handler. void HandleAssignTask(const rpc::AssignTaskRequest &request, @@ -462,7 +460,7 @@ class CoreWorker { rpc::GrpcServer core_worker_server_; // Client to the GCS shared by core worker interfaces. - gcs::RedisGcsClient gcs_client_; + std::shared_ptr gcs_client_; // Client to the raylet shared by core worker interfaces. std::unique_ptr raylet_client_; diff --git a/src/ray/core_worker/profiling.cc b/src/ray/core_worker/profiling.cc index 81e682421..cb55d5ac6 100644 --- a/src/ray/core_worker/profiling.cc +++ b/src/ray/core_worker/profiling.cc @@ -6,7 +6,7 @@ namespace ray { namespace worker { -ProfileEvent::ProfileEvent(const std::shared_ptr profiler, +ProfileEvent::ProfileEvent(const std::shared_ptr &profiler, const std::string &event_type) : profiler_(profiler) { rpc_event_.set_event_type(event_type); @@ -14,11 +14,11 @@ ProfileEvent::ProfileEvent(const std::shared_ptr profiler, } Profiler::Profiler(WorkerContext &worker_context, const std::string &node_ip_address, - boost::asio::io_service &io_service, gcs::RedisGcsClient &gcs_client) + boost::asio::io_service &io_service, + const std::shared_ptr &gcs_client) : io_service_(io_service), timer_(io_service_, boost::asio::chrono::seconds(1)), gcs_client_(gcs_client) { - absl::MutexLock l(&mu_); rpc_profile_data_.set_component_type(WorkerTypeString(worker_context.GetWorkerType())); rpc_profile_data_.set_component_id(worker_context.GetWorkerID().Binary()); rpc_profile_data_.set_node_ip_address(node_ip_address); @@ -26,16 +26,16 @@ Profiler::Profiler(WorkerContext &worker_context, const std::string &node_ip_add } void Profiler::AddEvent(const rpc::ProfileTableData::ProfileEvent &event) { - absl::MutexLock l(&mu_); + absl::MutexLock lock(&mutex_); rpc_profile_data_.add_profile_events()->CopyFrom(event); } void Profiler::FlushEvents() { - absl::MutexLock l(&mu_); + absl::MutexLock lock(&mutex_); if (rpc_profile_data_.profile_events_size() != 0) { // TODO(edoakes): this should be migrated to use the new GCS client interface // instead of the raw table interface once it's ready. - if (!gcs_client_.profile_table().AddProfileEventBatch(rpc_profile_data_).ok()) { + if (!gcs_client_->profile_table().AddProfileEventBatch(rpc_profile_data_).ok()) { RAY_LOG(WARNING) << "Failed to push profile events to GCS."; } else { RAY_LOG(DEBUG) << "Pushed " << rpc_profile_data_.profile_events_size() diff --git a/src/ray/core_worker/profiling.h b/src/ray/core_worker/profiling.h index f523ca63f..913e72800 100644 --- a/src/ray/core_worker/profiling.h +++ b/src/ray/core_worker/profiling.h @@ -1,6 +1,7 @@ #ifndef RAY_CORE_WORKER_PROFILING_H #define RAY_CORE_WORKER_PROFILING_H +#include "absl/base/thread_annotations.h" #include "absl/synchronization/mutex.h" #include "absl/time/clock.h" @@ -14,14 +15,18 @@ namespace worker { class Profiler { public: Profiler(WorkerContext &worker_context, const std::string &node_ip_address, - boost::asio::io_service &io_service, gcs::RedisGcsClient &gcs_client); + boost::asio::io_service &io_service, + const std::shared_ptr &gcs_client); // Add an event to the queue to be flushed periodically. - void AddEvent(const rpc::ProfileTableData::ProfileEvent &event); + void AddEvent(const rpc::ProfileTableData::ProfileEvent &event) LOCKS_EXCLUDED(mutex_); private: // Flush all of the events that have been added since last flush to the GCS. - void FlushEvents(); + void FlushEvents() LOCKS_EXCLUDED(mutex_); + + // Mutex guarding rpc_profile_data_. + absl::Mutex mutex_; // ASIO IO service event loop. Must be started by the caller. boost::asio::io_service &io_service_; @@ -31,28 +36,32 @@ class Profiler { // RPC message containing profiling data. Holds the queue of profile events // until they are flushed. - rpc::ProfileTableData rpc_profile_data_ GUARDED_BY(mu_); + rpc::ProfileTableData rpc_profile_data_ GUARDED_BY(mutex_); - gcs::RedisGcsClient &gcs_client_; - - absl::Mutex mu_; + // Client to the GCS used to push profile events to it. + std::shared_ptr gcs_client_; }; class ProfileEvent { public: - ProfileEvent(const std::shared_ptr profiler, const std::string &event_type); + ProfileEvent(const std::shared_ptr &profiler, const std::string &event_type); + // Set the end time for the event and add it to the profiler. ~ProfileEvent() { rpc_event_.set_end_time(absl::GetCurrentTimeNanos() / 1e9); profiler_->AddEvent(rpc_event_); } + // Set extra metadata for the event, which could change during the event. void SetExtraData(const std::string &extra_data) { rpc_event_.set_extra_data(extra_data); } private: - const std::shared_ptr profiler_; + // shared_ptr to the profiler that this event will be added to when it is destructed. + std::shared_ptr profiler_; + + // Underlying proto data structure that holds the event data. rpc::ProfileTableData::ProfileEvent rpc_event_; };