mirror of
https://github.com/wassname/ray.git
synced 2026-07-25 13:30:52 +08:00
[gRPC] Migrate raylet client implementation to grpc (#5120)
This commit is contained in:
+49
-18
@@ -52,7 +52,7 @@ proto_library(
|
||||
|
||||
cc_proto_library(
|
||||
name = "node_manager_cc_proto",
|
||||
deps = ["node_manager_proto"],
|
||||
deps = [":node_manager_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
@@ -62,7 +62,21 @@ proto_library(
|
||||
|
||||
cc_proto_library(
|
||||
name = "object_manager_cc_proto",
|
||||
deps = ["object_manager_proto"],
|
||||
deps = [":object_manager_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "raylet_proto",
|
||||
srcs = ["src/ray/protobuf/raylet.proto"],
|
||||
deps = [
|
||||
":common_proto",
|
||||
":gcs_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_proto_library(
|
||||
name = "raylet_cc_proto",
|
||||
deps = [":raylet_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
@@ -91,7 +105,7 @@ cc_proto_library(
|
||||
|
||||
# === Begin of rpc definitions ===
|
||||
|
||||
# grpc common lib
|
||||
# GRPC common lib.
|
||||
cc_library(
|
||||
name = "grpc_common_lib",
|
||||
srcs = glob([
|
||||
@@ -141,7 +155,7 @@ cc_grpc_library(
|
||||
deps = [":object_manager_cc_proto"],
|
||||
)
|
||||
|
||||
# Object manager server and client.
|
||||
# Object manager rpc server and client.
|
||||
cc_library(
|
||||
name = "object_manager_rpc",
|
||||
hdrs = glob([
|
||||
@@ -157,7 +171,35 @@ cc_library(
|
||||
],
|
||||
)
|
||||
|
||||
# worker gRPC lib.
|
||||
# Raylet gRPC lib.
|
||||
cc_grpc_library(
|
||||
name = "raylet_cc_grpc",
|
||||
srcs = [":raylet_proto"],
|
||||
grpc_only = True,
|
||||
deps = [":raylet_cc_proto"],
|
||||
)
|
||||
|
||||
# Raylet rpc server and client.
|
||||
cc_library(
|
||||
name = "raylet_rpc",
|
||||
srcs = glob([
|
||||
"src/ray/rpc/raylet/*.cc",
|
||||
]),
|
||||
hdrs = glob([
|
||||
"src/ray/rpc/raylet/*.h",
|
||||
"src/ray/raylet/*.h",
|
||||
]),
|
||||
copts = COPTS,
|
||||
deps = [
|
||||
":grpc_common_lib",
|
||||
":ray_common",
|
||||
":raylet_cc_grpc",
|
||||
"@boost//:asio",
|
||||
"@com_github_grpc_grpc//:grpc++",
|
||||
],
|
||||
)
|
||||
|
||||
# Worker gRPC lib.
|
||||
cc_grpc_library(
|
||||
name = "worker_cc_grpc",
|
||||
srcs = [":worker_proto"],
|
||||
@@ -165,7 +207,7 @@ cc_grpc_library(
|
||||
deps = [":worker_cc_proto"],
|
||||
)
|
||||
|
||||
# worker server and client.
|
||||
# Worker server and client.
|
||||
cc_library(
|
||||
name = "worker_rpc",
|
||||
hdrs = glob([
|
||||
@@ -201,7 +243,6 @@ cc_library(
|
||||
copts = COPTS,
|
||||
deps = [
|
||||
":common_cc_proto",
|
||||
":node_manager_fbs",
|
||||
":ray_util",
|
||||
"@boost//:asio",
|
||||
"@com_github_grpc_grpc//:grpc++",
|
||||
@@ -300,11 +341,11 @@ cc_library(
|
||||
deps = [
|
||||
":common_cc_proto",
|
||||
":gcs",
|
||||
":node_manager_fbs",
|
||||
":node_manager_rpc",
|
||||
":object_manager",
|
||||
":ray_common",
|
||||
":ray_util",
|
||||
":raylet_rpc",
|
||||
":stats_lib",
|
||||
":worker_rpc",
|
||||
"@boost//:asio",
|
||||
@@ -375,7 +416,6 @@ cc_test(
|
||||
srcs = ["src/ray/raylet/lineage_cache_test.cc"],
|
||||
copts = COPTS,
|
||||
deps = [
|
||||
":node_manager_fbs",
|
||||
":raylet_lib",
|
||||
"@com_google_googletest//:gtest_main",
|
||||
],
|
||||
@@ -386,7 +426,6 @@ cc_test(
|
||||
srcs = ["src/ray/raylet/reconstruction_policy_test.cc"],
|
||||
copts = COPTS,
|
||||
deps = [
|
||||
":node_manager_fbs",
|
||||
":object_manager",
|
||||
":raylet_lib",
|
||||
"@com_google_googletest//:gtest_main",
|
||||
@@ -576,7 +615,6 @@ cc_library(
|
||||
deps = [
|
||||
":gcs_cc_proto",
|
||||
":hiredis",
|
||||
":node_manager_fbs",
|
||||
":node_manager_rpc",
|
||||
":ray_common",
|
||||
":ray_util",
|
||||
@@ -632,13 +670,6 @@ flatbuffer_cc_library(
|
||||
out_prefix = "src/ray/common/",
|
||||
)
|
||||
|
||||
flatbuffer_cc_library(
|
||||
name = "node_manager_fbs",
|
||||
srcs = ["src/ray/raylet/format/node_manager.fbs"],
|
||||
flatc_args = FLATC_ARGS,
|
||||
out_prefix = "src/ray/raylet/format/",
|
||||
)
|
||||
|
||||
flatbuffer_cc_library(
|
||||
name = "object_manager_fbs",
|
||||
srcs = ["src/ray/object_manager/format/object_manager.fbs"],
|
||||
|
||||
@@ -37,9 +37,9 @@ public class RayletClientImpl implements RayletClient {
|
||||
private long client = 0;
|
||||
|
||||
// TODO(qwang): JobId parameter can be removed once we embed jobId in driverId.
|
||||
public RayletClientImpl(String schedulerSockName, UniqueId clientId,
|
||||
public RayletClientImpl(String schedulerSockName, UniqueId workerId,
|
||||
boolean isWorker, JobId jobId) {
|
||||
client = nativeInit(schedulerSockName, clientId.getBytes(),
|
||||
client = nativeInit(schedulerSockName, workerId.getBytes(),
|
||||
isWorker, jobId.getBytes());
|
||||
}
|
||||
|
||||
|
||||
@@ -5,6 +5,13 @@ from __future__ import print_function
|
||||
import os
|
||||
import sys
|
||||
|
||||
# MUST import ray._raylet before pyarrow to initialize some global variables.
|
||||
# It seems the library related to memory allocation in pyarrow will destroy the
|
||||
# initialization of grpc if we import pyarrow at first.
|
||||
# NOTE(JoeyJiang): See https://github.com/ray-project/ray/issues/5219 for more
|
||||
# details.
|
||||
import ray._raylet
|
||||
|
||||
if "pyarrow" in sys.modules:
|
||||
raise ImportError("Ray must be imported before pyarrow because Ray "
|
||||
"requires a specific version of pyarrow (which is "
|
||||
|
||||
@@ -220,14 +220,14 @@ cdef class RayletClient:
|
||||
cdef unique_ptr[CRayletClient] client
|
||||
|
||||
def __cinit__(self, raylet_socket,
|
||||
ClientID client_id,
|
||||
WorkerID worker_id,
|
||||
c_bool is_worker,
|
||||
JobID job_id):
|
||||
# We know that we are using Python, so just skip the language
|
||||
# parameter.
|
||||
# TODO(suquark): Should we allow unicode chars in "raylet_socket"?
|
||||
self.client.reset(new CRayletClient(
|
||||
raylet_socket.encode("ascii"), client_id.native(), is_worker,
|
||||
raylet_socket.encode("ascii"), worker_id.native(), is_worker,
|
||||
job_id.native(), LANGUAGE_PYTHON))
|
||||
|
||||
def disconnect(self):
|
||||
@@ -374,7 +374,7 @@ cdef class RayletClient:
|
||||
|
||||
@property
|
||||
def client_id(self):
|
||||
return ClientID(self.client.get().GetClientID().Binary())
|
||||
return ClientID(self.client.get().GetWorkerId().Binary())
|
||||
|
||||
@property
|
||||
def job_id(self):
|
||||
|
||||
@@ -88,16 +88,16 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
|
||||
|
||||
|
||||
cdef extern from "ray/protobuf/common.pb.h" nogil:
|
||||
cdef cppclass CLanguage "Language":
|
||||
cdef cppclass CLanguage "ray::rpc::Language":
|
||||
pass
|
||||
|
||||
|
||||
# This is a workaround for C++ enum class since Cython has no corresponding
|
||||
# representation.
|
||||
cdef extern from "ray/protobuf/common.pb.h" namespace "Language" nogil:
|
||||
cdef CLanguage LANGUAGE_PYTHON "Language::PYTHON"
|
||||
cdef CLanguage LANGUAGE_CPP "Language::CPP"
|
||||
cdef CLanguage LANGUAGE_JAVA "Language::JAVA"
|
||||
cdef extern from "ray/protobuf/common.pb.h" namespace "ray::rpc::Language" nogil:
|
||||
cdef CLanguage LANGUAGE_PYTHON "ray::rpc::Language::PYTHON"
|
||||
cdef CLanguage LANGUAGE_CPP "ray::rpc::Language::CPP"
|
||||
cdef CLanguage LANGUAGE_JAVA "ray::rpc::Language::JAVA"
|
||||
|
||||
|
||||
cdef extern from "ray/common/task/scheduling_resources.h" \
|
||||
|
||||
@@ -23,14 +23,14 @@ from ray.includes.task cimport CTaskSpec
|
||||
|
||||
|
||||
cdef extern from "ray/protobuf/gcs.pb.h" nogil:
|
||||
cdef cppclass GCSProfileEvent "ProfileTableData::ProfileEvent":
|
||||
cdef cppclass GCSProfileEvent "ray::rpc::ProfileTableData::ProfileEvent":
|
||||
void set_event_type(const c_string &value)
|
||||
void set_start_time(double value)
|
||||
void set_end_time(double value)
|
||||
c_string set_extra_data(const c_string &value)
|
||||
GCSProfileEvent()
|
||||
|
||||
cdef cppclass GCSProfileTableData "ProfileTableData":
|
||||
cdef cppclass GCSProfileTableData "ray::rpc::ProfileTableData":
|
||||
void set_component_type(const c_string &value)
|
||||
void set_component_id(const c_string &value)
|
||||
void set_node_ip_address(const c_string &value)
|
||||
@@ -43,13 +43,12 @@ ctypedef unordered_map[c_string, c_vector[pair[int64_t, double]]] \
|
||||
ctypedef pair[c_vector[CObjectID], c_vector[CObjectID]] WaitResultPair
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/raylet_client.h" nogil:
|
||||
cdef cppclass CRayletClient "RayletClient":
|
||||
cdef extern from "ray/rpc/raylet/raylet_client.h" namespace "ray::rpc" nogil:
|
||||
cdef cppclass CRayletClient "ray::rpc::RayletClient":
|
||||
CRayletClient(const c_string &raylet_socket,
|
||||
const CClientID &client_id,
|
||||
const CWorkerID &worker_id,
|
||||
c_bool is_worker, const CJobID &job_id,
|
||||
const CLanguage &language)
|
||||
CRayStatus Disconnect()
|
||||
CRayStatus SubmitTask(const CTaskSpec &task_spec)
|
||||
CRayStatus GetTask(unique_ptr[CTaskSpec] *task_spec)
|
||||
CRayStatus TaskDone()
|
||||
@@ -73,7 +72,8 @@ cdef extern from "ray/raylet/raylet_client.h" nogil:
|
||||
const CActorID &actor_id, const CActorCheckpointID &checkpoint_id)
|
||||
CRayStatus SetResource(const c_string &resource_name, const double capacity, const CClientID &client_Id)
|
||||
CLanguage GetLanguage() const
|
||||
CClientID GetClientID() const
|
||||
CWorkerID GetWorkerId() const
|
||||
CJobID GetJobID() const
|
||||
c_bool IsWorker() const
|
||||
CRayStatus Disconnect()
|
||||
const ResourceMappingType &GetResourceIDs() const
|
||||
|
||||
@@ -458,11 +458,9 @@ print("success")
|
||||
# Make sure the first driver ran to completion.
|
||||
assert "success" in out
|
||||
|
||||
nonexistent_id_bytes = _random_string()
|
||||
nonexistent_id_hex = ray.utils.binary_to_hex(nonexistent_id_bytes)
|
||||
# Define a driver that creates one task that depends on a nonexistent
|
||||
# object. This task will be queued as waiting to execute.
|
||||
driver_script = """
|
||||
driver_script_template = """
|
||||
import time
|
||||
import ray
|
||||
ray.init(redis_address="{}")
|
||||
@@ -472,11 +470,15 @@ def g(x):
|
||||
g.remote(ray.ObjectID(ray.utils.hex_to_binary("{}")))
|
||||
time.sleep(1)
|
||||
print("success")
|
||||
""".format(redis_address, nonexistent_id_hex)
|
||||
"""
|
||||
|
||||
# Create some drivers and let them exit and make sure everything is
|
||||
# still alive.
|
||||
for _ in range(3):
|
||||
nonexistent_id_bytes = _random_string()
|
||||
nonexistent_id_hex = ray.utils.binary_to_hex(nonexistent_id_bytes)
|
||||
driver_script = driver_script_template.format(redis_address,
|
||||
nonexistent_id_hex)
|
||||
out = run_string_as_driver(driver_script)
|
||||
# Simulate the nonexistent dependency becoming available.
|
||||
ray.worker.global_worker.put_object(
|
||||
@@ -484,10 +486,8 @@ print("success")
|
||||
# Make sure the first driver ran to completion.
|
||||
assert "success" in out
|
||||
|
||||
nonexistent_id_bytes = _random_string()
|
||||
nonexistent_id_hex = ray.utils.binary_to_hex(nonexistent_id_bytes)
|
||||
# Define a driver that calls `ray.wait` on a nonexistent object.
|
||||
driver_script = """
|
||||
driver_script_template = """
|
||||
import time
|
||||
import ray
|
||||
ray.init(redis_address="{}")
|
||||
@@ -497,11 +497,15 @@ def g():
|
||||
g.remote()
|
||||
time.sleep(1)
|
||||
print("success")
|
||||
""".format(redis_address, nonexistent_id_hex)
|
||||
"""
|
||||
|
||||
# Create some drivers and let them exit and make sure everything is
|
||||
# still alive.
|
||||
for _ in range(3):
|
||||
nonexistent_id_bytes = _random_string()
|
||||
nonexistent_id_hex = ray.utils.binary_to_hex(nonexistent_id_bytes)
|
||||
driver_script = driver_script_template.format(redis_address,
|
||||
nonexistent_id_hex)
|
||||
out = run_string_as_driver(driver_script)
|
||||
# Simulate the nonexistent dependency becoming available.
|
||||
ray.worker.global_worker.put_object(
|
||||
|
||||
@@ -41,7 +41,6 @@ import ray.state
|
||||
from ray import (
|
||||
ActorHandleID,
|
||||
ActorID,
|
||||
ClientID,
|
||||
WorkerID,
|
||||
JobID,
|
||||
ObjectID,
|
||||
@@ -1923,7 +1922,7 @@ def connect(node,
|
||||
|
||||
worker.raylet_client = ray._raylet.RayletClient(
|
||||
node.raylet_socket_name,
|
||||
ClientID(worker.worker_id),
|
||||
WorkerID(worker.worker_id),
|
||||
(mode == WORKER_MODE),
|
||||
worker.current_job_id,
|
||||
)
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
#include <google/protobuf/map.h>
|
||||
#include <google/protobuf/repeated_field.h>
|
||||
#include <grpcpp/grpcpp.h>
|
||||
|
||||
#include "status.h"
|
||||
|
||||
namespace ray {
|
||||
@@ -73,6 +72,17 @@ inline std::vector<T> VectorFromProtobuf(
|
||||
return std::vector<T>(pb_repeated.begin(), pb_repeated.end());
|
||||
}
|
||||
|
||||
template <typename Message>
|
||||
using AddFunction = void (Message::*)(const ::std::string &value);
|
||||
/// Add a vector of type ID to protobuf message.
|
||||
template <typename ID, typename Message>
|
||||
inline void IdVectorToProtobuf(const std::vector<ID> &ids, Message &message,
|
||||
AddFunction<Message> add_func) {
|
||||
for (const auto &id : ids) {
|
||||
(message.*add_func)(id.Binary());
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts a Protobuf `RepeatedField` to a vector of IDs.
|
||||
template <class ID>
|
||||
inline std::vector<ID> IdVectorFromProtobuf(
|
||||
|
||||
@@ -20,8 +20,12 @@ RAY_CONFIG(int64_t, ray_cookie, 0x5241590000000000)
|
||||
/// warning is logged that the handler is taking too long.
|
||||
RAY_CONFIG(int64_t, handler_warning_timeout_ms, 100)
|
||||
|
||||
/// The duration between heartbeats. These are sent by the raylet.
|
||||
/// The duration between heartbeats. This value is used for both worker and raylet.
|
||||
RAY_CONFIG(int64_t, heartbeat_timeout_milliseconds, 100)
|
||||
/// Worker heartbeats also use `heartbeat_timeout_milliseconds` as timer timeout period.
|
||||
/// If a worker has not sent a heartbeat in the last `num_worker_heartbeats_timeout`
|
||||
/// heartbeat intervals, raylet will mark this worker as dead.
|
||||
RAY_CONFIG(int64_t, num_worker_heartbeats_timeout, 30)
|
||||
/// If a component has not sent a heartbeat in the last num_heartbeats_timeout
|
||||
/// heartbeat intervals, the raylet monitor process will report
|
||||
/// it as dead to the db_client table.
|
||||
@@ -152,6 +156,9 @@ RAY_CONFIG(uint32_t, num_actor_checkpoints_to_keep, 20)
|
||||
/// Maximum number of ids in one batch to send to GCS to delete keys.
|
||||
RAY_CONFIG(uint32_t, maximum_gcs_deletion_batch_size, 1000)
|
||||
|
||||
/// Number of times for a raylet client to retry to register.
|
||||
RAY_CONFIG(int, num_raylet_client_retry_times, 25)
|
||||
|
||||
/// When getting objects from object store, print a warning every this number of attempts.
|
||||
RAY_CONFIG(uint32_t, object_store_get_warn_per_num_attempts, 50)
|
||||
|
||||
|
||||
@@ -674,37 +674,22 @@ std::string ResourceIdSet::ToString() const {
|
||||
return return_string;
|
||||
}
|
||||
|
||||
std::vector<flatbuffers::Offset<protocol::ResourceIdSetInfo>> ResourceIdSet::ToFlatbuf(
|
||||
flatbuffers::FlatBufferBuilder &fbb) const {
|
||||
std::vector<flatbuffers::Offset<protocol::ResourceIdSetInfo>> return_message;
|
||||
std::vector<rpc::ResourceIdSetInfo> ResourceIdSet::ToProtobuf() const {
|
||||
std::vector<rpc::ResourceIdSetInfo> resources;
|
||||
for (auto const &resource_pair : available_resources_) {
|
||||
std::vector<int64_t> resource_ids;
|
||||
std::vector<double> resource_fractions;
|
||||
rpc::ResourceIdSetInfo resource_id_set_info;
|
||||
resource_id_set_info.set_resource_name(resource_pair.first);
|
||||
for (auto whole_id : resource_pair.second.WholeIds()) {
|
||||
resource_ids.push_back(whole_id);
|
||||
resource_fractions.push_back(1);
|
||||
resource_id_set_info.add_resource_ids(whole_id);
|
||||
resource_id_set_info.add_resource_fractions(1);
|
||||
}
|
||||
|
||||
for (auto const &fractional_pair : resource_pair.second.FractionalIds()) {
|
||||
resource_ids.push_back(fractional_pair.first);
|
||||
resource_fractions.push_back(fractional_pair.second.ToDouble());
|
||||
resource_id_set_info.add_resource_ids(fractional_pair.first);
|
||||
resource_id_set_info.add_resource_fractions(fractional_pair.second.ToDouble());
|
||||
}
|
||||
|
||||
auto resource_id_set_message = protocol::CreateResourceIdSetInfo(
|
||||
fbb, fbb.CreateString(resource_pair.first), fbb.CreateVector(resource_ids),
|
||||
fbb.CreateVector(resource_fractions));
|
||||
|
||||
return_message.push_back(resource_id_set_message);
|
||||
resources.emplace_back(resource_id_set_info);
|
||||
}
|
||||
|
||||
return return_message;
|
||||
}
|
||||
|
||||
const std::string ResourceIdSet::Serialize() const {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto resource_id_set_flatbuf = ToFlatbuf(fbb);
|
||||
fbb.Finish(fbb.CreateVector(resource_id_set_flatbuf));
|
||||
return std::string(fbb.GetBufferPointer(), fbb.GetBufferPointer() + fbb.GetSize());
|
||||
return resources;
|
||||
}
|
||||
|
||||
/// SchedulingResources class implementation
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/protobuf/common.pb.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
@@ -422,21 +422,13 @@ class ResourceIdSet {
|
||||
/// \return A human-readable string version of the object.
|
||||
std::string ToString() const;
|
||||
|
||||
/// \brief Serialize this object using flatbuffers.
|
||||
/// \brief Convert this object to a vector of protobuf `ResourceIdSetInfo`s.
|
||||
///
|
||||
/// \param fbb A flatbuffer builder object.
|
||||
/// \return A flatbuffer serialized version of this object.
|
||||
std::vector<flatbuffers::Offset<ray::protocol::ResourceIdSetInfo>> ToFlatbuf(
|
||||
flatbuffers::FlatBufferBuilder &fbb) const;
|
||||
|
||||
/// \brief Serialize this object as a string.
|
||||
///
|
||||
/// \return A serialized string of this object.
|
||||
/// TODO(zhijunfu): this can be removed after raylet client is migrated to grpc.
|
||||
const std::string Serialize() const;
|
||||
/// \return A vector inclusing resource id set infos.
|
||||
std::vector<rpc::ResourceIdSetInfo> ToProtobuf() const;
|
||||
|
||||
private:
|
||||
/// A mapping from reosurce name to a set of resource IDs for that resource.
|
||||
/// A mapping from resource name to a set of resource IDs for that resource.
|
||||
std::unordered_map<std::string, ResourceIds> available_resources_;
|
||||
};
|
||||
|
||||
|
||||
@@ -3,9 +3,9 @@
|
||||
|
||||
#include <inttypes.h>
|
||||
|
||||
#include "ray/common/task/task_common.h"
|
||||
#include "ray/common/task/task_execution_spec.h"
|
||||
#include "ray/common/task/task_spec.h"
|
||||
#include "ray/protobuf/common.pb.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
#include "ray/common/buffer.h"
|
||||
#include "ray/common/id.h"
|
||||
#include "ray/common/task/task_spec.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
#include "ray/util/util.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
@@ -27,9 +27,9 @@ CoreWorker::CoreWorker(
|
||||
// so that the worker (java/python .etc) can retrieve and handle the error
|
||||
// instead of crashing.
|
||||
raylet_client_ = std::unique_ptr<RayletClient>(new RayletClient(
|
||||
raylet_socket_, ClientID::FromBinary(worker_context_.GetWorkerID().Binary()),
|
||||
(worker_type_ == WorkerType::WORKER), worker_context_.GetCurrentJobID(), language_,
|
||||
rpc_server_port));
|
||||
raylet_socket_, WorkerID::FromBinary(worker_context_.GetWorkerID().Binary()),
|
||||
(worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(),
|
||||
language_, rpc_server_port));
|
||||
}
|
||||
|
||||
} // namespace ray
|
||||
|
||||
@@ -7,10 +7,12 @@
|
||||
#include "ray/core_worker/object_interface.h"
|
||||
#include "ray/core_worker/task_execution.h"
|
||||
#include "ray/core_worker/task_interface.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
using rpc::RayletClient;
|
||||
|
||||
/// The root class that contains all the core and language-independent functionalities
|
||||
/// of the worker. This class is supposed to be used to implement app-language (Java,
|
||||
/// Python, etc) workers.
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
#include "ray/common/buffer.h"
|
||||
#include "ray/core_worker/context.h"
|
||||
#include "ray/core_worker/core_worker.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
|
||||
#include <boost/asio.hpp>
|
||||
#include <boost/asio/error.hpp>
|
||||
|
||||
@@ -5,6 +5,8 @@
|
||||
#include "ray/core_worker/lib/java/jni_utils.h"
|
||||
#include "ray/core_worker/object_interface.h"
|
||||
|
||||
using ray::rpc::RayletClient;
|
||||
|
||||
inline ray::CoreWorkerObjectInterface *GetObjectInterfaceFromPointer(
|
||||
jlong nativeObjectInterfacePointer) {
|
||||
return reinterpret_cast<ray::CoreWorkerObjectInterface *>(nativeObjectInterfacePointer);
|
||||
|
||||
@@ -11,6 +11,8 @@
|
||||
|
||||
namespace ray {
|
||||
|
||||
using rpc::RayletClient;
|
||||
|
||||
class CoreWorker;
|
||||
class CoreWorkerStoreProvider;
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
#include "ray/common/status.h"
|
||||
#include "ray/core_worker/common.h"
|
||||
#include "ray/core_worker/store_provider/store_provider.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
|
||||
@@ -8,10 +8,12 @@
|
||||
#include "ray/core_worker/common.h"
|
||||
#include "ray/core_worker/store_provider/local_plasma_provider.h"
|
||||
#include "ray/core_worker/store_provider/store_provider.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
using rpc::RayletClient;
|
||||
|
||||
class CoreWorker;
|
||||
|
||||
/// The class provides implementations for accessing plasma store, which includes both
|
||||
|
||||
@@ -15,6 +15,8 @@
|
||||
|
||||
namespace ray {
|
||||
|
||||
using rpc::RayletClient;
|
||||
|
||||
class CoreWorker;
|
||||
|
||||
/// Options of a non-actor-creation task.
|
||||
|
||||
@@ -56,8 +56,8 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
|
||||
// rpc reply first before the NotifyUnblocked message arrives,
|
||||
// as they use different connections, the `TaskDone` message is sent
|
||||
// to raylet via the same connection so the order is guaranteed.
|
||||
raylet_client_->TaskDone();
|
||||
// send rpc reply.
|
||||
RAY_UNUSED(raylet_client_->TaskDone());
|
||||
// Send rpc reply.
|
||||
send_reply_callback(status, nullptr, nullptr);
|
||||
}
|
||||
|
||||
|
||||
@@ -5,11 +5,13 @@
|
||||
|
||||
#include "ray/core_worker/object_interface.h"
|
||||
#include "ray/core_worker/transport/transport.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
#include "ray/rpc/worker/worker_server.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
using rpc::RayletClient;
|
||||
|
||||
/// In raylet task submitter and receiver, a task is submitted to raylet, and possibly
|
||||
/// gets forwarded to another raylet on which node the task should be executed, and
|
||||
/// then a worker on that node gets this task and starts executing it.
|
||||
|
||||
@@ -769,9 +769,8 @@ void ObjectManager::SpreadFreeObjectsRequest(
|
||||
const std::vector<std::shared_ptr<rpc::ObjectManagerClient>> &rpc_clients) {
|
||||
// This code path should be called from node manager.
|
||||
rpc::FreeObjectsRequest free_objects_request;
|
||||
for (const auto &e : object_ids) {
|
||||
free_objects_request.add_object_ids(e.Binary());
|
||||
}
|
||||
IdVectorToProtobuf<ObjectID, rpc::FreeObjectsRequest>(
|
||||
object_ids, free_objects_request, &rpc::FreeObjectsRequest::add_object_ids);
|
||||
|
||||
for (auto &rpc_client : rpc_clients) {
|
||||
rpc_client->FreeObjects(free_objects_request, [](const Status &status,
|
||||
|
||||
@@ -11,6 +11,18 @@ enum Language {
|
||||
CPP = 2;
|
||||
}
|
||||
|
||||
// Resource id set info.
|
||||
message ResourceIdSetInfo {
|
||||
// The name of the resource.
|
||||
bytes resource_name = 1;
|
||||
// The resource IDs reserved for this worker.
|
||||
repeated uint64 resource_ids = 2;
|
||||
// The fraction of each resource ID that is reserved for this worker. Note
|
||||
// that the length of this list must be the same as the length of
|
||||
// resource_ids.
|
||||
repeated double resource_fractions = 3;
|
||||
}
|
||||
|
||||
// Type of a worker.
|
||||
enum WorkerType {
|
||||
WORKER = 0;
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
syntax = "proto3";
|
||||
|
||||
package ray.rpc;
|
||||
|
||||
import "src/ray/protobuf/common.proto";
|
||||
import "src/ray/protobuf/gcs.proto";
|
||||
|
||||
/// Service request and reply messages.
|
||||
message RegisterClientRequest {
|
||||
// Indicates the client is a worker or a driver.
|
||||
bool is_worker = 1;
|
||||
// The worker id.
|
||||
bytes worker_id = 2;
|
||||
// The process ID of this worker.
|
||||
uint32 worker_pid = 3;
|
||||
// The job ID.
|
||||
bytes job_id = 4;
|
||||
// Language of this worker.
|
||||
Language language = 5;
|
||||
// Port that this worker is listening on.
|
||||
// If port > 0, then worker will listen to this port and wait for
|
||||
// raylet to push tasks, instead of invoking GetTask().
|
||||
int32 port = 6;
|
||||
}
|
||||
message RegisterClientReply {
|
||||
repeated int32 gpu_ids = 1;
|
||||
}
|
||||
|
||||
message SubmitTaskRequest {
|
||||
TaskSpec task_spec = 1;
|
||||
}
|
||||
message SubmitTaskReply {
|
||||
}
|
||||
|
||||
message DisconnectClientRequest {
|
||||
bytes worker_id = 1;
|
||||
}
|
||||
message DisconnectClientReply {
|
||||
}
|
||||
|
||||
message GetTaskRequest {
|
||||
bytes worker_id = 1;
|
||||
}
|
||||
message GetTaskReply {
|
||||
// A string of bytes representing the task specification.
|
||||
bytes task_spec = 1;
|
||||
// A list of the resources reserved for this worker.
|
||||
repeated ResourceIdSetInfo fractional_resource_ids = 2;
|
||||
}
|
||||
|
||||
message TaskDoneRequest {
|
||||
bytes worker_id = 1;
|
||||
}
|
||||
message TaskDoneReply {
|
||||
}
|
||||
|
||||
message FetchOrReconstructRequest {
|
||||
// List of object IDs of the objects that we want to reconstruct or fetch.
|
||||
repeated bytes object_ids = 1;
|
||||
// Indicates that we only want to fetch objects, not reconstruct them.
|
||||
bool fetch_only = 2;
|
||||
// The current task ID. If fetch_only is false, then this task is blocked.
|
||||
bytes task_id = 3;
|
||||
// The worker ID.
|
||||
bytes worker_id = 4;
|
||||
}
|
||||
message FetchOrReconstructReply {
|
||||
}
|
||||
|
||||
message NotifyUnblockedRequest {
|
||||
bytes worker_id = 1;
|
||||
// The current task ID. This task is no longer blocked.
|
||||
bytes task_id = 2;
|
||||
}
|
||||
message NotifyUnblockedReply {
|
||||
}
|
||||
|
||||
message WaitRequest {
|
||||
// List of object ids we'll be waiting on.
|
||||
repeated bytes object_ids = 1;
|
||||
// Number of objects expected to be returned, if available.
|
||||
uint64 num_ready_objects = 2;
|
||||
// Timeout in milliseconds.
|
||||
int64 timeout = 3;
|
||||
// Whether to wait until objects appear locally.
|
||||
bool wait_local = 4;
|
||||
// The current task ID. If there are less than num_ready_objects local, then
|
||||
// this task is blocked.
|
||||
bytes task_id = 5;
|
||||
// The worker ID.
|
||||
bytes worker_id = 6;
|
||||
}
|
||||
message WaitReply {
|
||||
// List of object ids found.
|
||||
repeated bytes found = 1;
|
||||
// List of object ids not found.
|
||||
repeated bytes remaining = 2;
|
||||
}
|
||||
|
||||
message PushErrorRequest {
|
||||
// The job id that the error is for.
|
||||
bytes job_id = 1;
|
||||
// The type of the error.
|
||||
bytes type = 2;
|
||||
// The error message.
|
||||
bytes error_message = 3;
|
||||
// The timestamp of the error message.
|
||||
double timestamp = 4;
|
||||
}
|
||||
message PushErrorReply {
|
||||
}
|
||||
|
||||
message PushProfileEventsRequest {
|
||||
ProfileTableData profile_table_data = 1;
|
||||
}
|
||||
message PushProfileEventsReply {
|
||||
}
|
||||
|
||||
message FreeObjectsInStoreRequest {
|
||||
// Whether keep this request within the local object store
|
||||
// or send it to all of the object stores.
|
||||
bool local_only = 1;
|
||||
// Whether also delete objects' creating tasks from GCS.
|
||||
bool delete_creating_tasks = 2;
|
||||
// List of object ids to delete from the object store.
|
||||
repeated bytes object_ids = 3;
|
||||
}
|
||||
message FreeObjectsInStoreReply {
|
||||
}
|
||||
|
||||
message PrepareActorCheckpointRequest {
|
||||
bytes actor_id = 1;
|
||||
bytes worker_id = 2;
|
||||
}
|
||||
message PrepareActorCheckpointReply {
|
||||
bytes checkpoint_id = 1;
|
||||
}
|
||||
|
||||
message NotifyActorResumedFromCheckpointRequest {
|
||||
// ID of the actor that resumed.
|
||||
bytes actor_id = 1;
|
||||
// ID of the checkpoint from which the actor was resumed.
|
||||
bytes checkpoint_id = 2;
|
||||
}
|
||||
message NotifyActorResumedFromCheckpointReply {
|
||||
}
|
||||
|
||||
message SetResourceRequest {
|
||||
// Name of the resource to be set.
|
||||
bytes resource_name = 1;
|
||||
// Capacity of the resource to be set.
|
||||
double capacity = 2;
|
||||
// Client ID where this resource will be set.
|
||||
bytes client_id = 3;
|
||||
}
|
||||
message SetResourceReply {
|
||||
}
|
||||
|
||||
message HeartbeatRequest {
|
||||
bytes worker_id = 1;
|
||||
bool is_worker = 2;
|
||||
}
|
||||
message HeartbeatReply {
|
||||
}
|
||||
|
||||
/// Worker-to-raylet RPC service interface.
|
||||
service RayletService {
|
||||
// Register a new worker to the raylet.
|
||||
rpc RegisterClient(RegisterClientRequest) returns (RegisterClientReply);
|
||||
// Submit a task to the raylet.
|
||||
rpc SubmitTask(SubmitTaskRequest) returns (SubmitTaskReply);
|
||||
// Disconnect this client from raylet gracefully.
|
||||
rpc DisconnectClient(DisconnectClientRequest) returns (DisconnectClientReply);
|
||||
// Get a new task from the raylet.
|
||||
rpc GetTask(GetTaskRequest) returns (GetTaskReply);
|
||||
// Notify the raylet that a task is finished.
|
||||
rpc TaskDone(TaskDoneRequest) returns (TaskDoneReply);
|
||||
// Reconstruct or fetch possibly lost objects.
|
||||
rpc FetchOrReconstruct(FetchOrReconstructRequest) returns (FetchOrReconstructReply);
|
||||
// For a worker that was blocked on some object(s), tell the raylet
|
||||
// that the worker is now unblocked.
|
||||
rpc NotifyUnblocked(NotifyUnblockedRequest) returns (NotifyUnblockedReply);
|
||||
// Wait for objects to be ready either from local or remote plasma stores.
|
||||
// The `WaitReply` contains the objects found and objects remaining.
|
||||
rpc Wait(WaitRequest) returns (WaitReply);
|
||||
// Push an error to the relevant driver.
|
||||
rpc PushError(PushErrorRequest) returns (PushErrorReply);
|
||||
// Push some profiling events to the GCS. When sending this message to the
|
||||
// node manager, the message itself is serialized as a ProfileTableData object.
|
||||
rpc PushProfileEvents(PushProfileEventsRequest) returns (PushProfileEventsReply);
|
||||
// Free the objects in plasma objects store.
|
||||
rpc FreeObjectsInStore(FreeObjectsInStoreRequest) returns (FreeObjectsInStoreReply);
|
||||
// Request raylet backend to prepare a checkpoint for an actor.
|
||||
rpc PrepareActorCheckpoint(PrepareActorCheckpointRequest)
|
||||
returns (PrepareActorCheckpointReply);
|
||||
// Notify raylet backend that an actor was resumed from a checkpoint.
|
||||
rpc NotifyActorResumedFromCheckpoint(NotifyActorResumedFromCheckpointRequest)
|
||||
returns (NotifyActorResumedFromCheckpointReply);
|
||||
// Set dynamic custom resource.
|
||||
rpc SetResource(SetResourceRequest) returns (SetResourceReply);
|
||||
// Send a heartbeat message to raylet.
|
||||
rpc Heartbeat(HeartbeatRequest) returns (HeartbeatReply);
|
||||
}
|
||||
@@ -8,9 +8,7 @@ message AssignTaskRequest {
|
||||
// The task to be pushed.
|
||||
Task task = 1;
|
||||
// A list of the resources reserved for this worker.
|
||||
// TODO(zhijunfu): `resource_ids` is represented as
|
||||
// flatbutters-serialized bytes, will be moved to protobuf later.
|
||||
bytes resource_ids = 2;
|
||||
repeated ResourceIdSetInfo resource_ids = 2;
|
||||
}
|
||||
|
||||
message AssignTaskReply {
|
||||
|
||||
@@ -4,13 +4,17 @@
|
||||
|
||||
#include "ray/common/id.h"
|
||||
#include "ray/core_worker/lib/java/jni_utils.h"
|
||||
#include "ray/raylet/raylet_client.h"
|
||||
#include "ray/rpc/raylet/raylet_client.h"
|
||||
#include "ray/util/logging.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
using ray::ClientID;
|
||||
using ray::WorkerID;
|
||||
using ray::rpc::RayletClient;
|
||||
|
||||
/*
|
||||
* Class: org_ray_runtime_raylet_RayletClientImpl
|
||||
* Method: nativeInit
|
||||
@@ -19,7 +23,7 @@ extern "C" {
|
||||
JNIEXPORT jlong JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeInit(
|
||||
JNIEnv *env, jclass, jstring sockName, jbyteArray workerId, jboolean isWorker,
|
||||
jbyteArray jobId) {
|
||||
const auto worker_id = JavaByteArrayToId<ClientID>(env, workerId);
|
||||
const auto worker_id = JavaByteArrayToId<WorkerID>(env, workerId);
|
||||
const auto job_id = JavaByteArrayToId<JobID>(env, jobId);
|
||||
const char *nativeString = env->GetStringUTFChars(sockName, JNI_FALSE);
|
||||
auto raylet_client = new std::unique_ptr<RayletClient>(
|
||||
|
||||
@@ -7,7 +7,6 @@
|
||||
#include "ray/common/task/task_execution_spec.h"
|
||||
#include "ray/common/task/task_spec.h"
|
||||
#include "ray/common/task/task_util.h"
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/raylet/lineage_cache.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
@@ -147,7 +147,7 @@ int main(int argc, char *argv[]) {
|
||||
RAY_LOG(DEBUG) << "Starting object manager with configuration: \n"
|
||||
<< "rpc_service_threads_number = "
|
||||
<< object_manager_config.rpc_service_threads_number
|
||||
<< "object_chunk_size = " << object_manager_config.object_chunk_size;
|
||||
<< ", object_chunk_size = " << object_manager_config.object_chunk_size;
|
||||
|
||||
// Initialize the node manager.
|
||||
boost::asio::io_service main_service;
|
||||
@@ -171,6 +171,7 @@ int main(int argc, char *argv[]) {
|
||||
server.reset();
|
||||
gcs_client->Disconnect();
|
||||
main_service.stop();
|
||||
RAY_LOG(INFO) << "Raylet server received SIGTERM message, shutting down...";
|
||||
};
|
||||
RAY_CHECK_OK(gcs_client->client_table().Disconnect(shutdown_callback));
|
||||
// Give a timeout for this Disconnect operation.
|
||||
|
||||
+300
-358
@@ -6,7 +6,6 @@
|
||||
|
||||
#include "ray/common/common_protocol.h"
|
||||
#include "ray/common/id.h"
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/stats/stats.h"
|
||||
|
||||
namespace {
|
||||
@@ -105,6 +104,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
|
||||
actor_registry_(),
|
||||
node_manager_server_("NodeManager", config.node_manager_port),
|
||||
node_manager_service_(io_service, *this),
|
||||
raylet_service_(io_service, *this),
|
||||
client_call_manager_(io_service) {
|
||||
RAY_CHECK(heartbeat_period_.count() > 0);
|
||||
// Initialize the resource map with own cluster resource configuration.
|
||||
@@ -249,6 +249,7 @@ void NodeManager::KillWorker(std::shared_ptr<Worker> worker) {
|
||||
// up its state before force killing. The client socket will be closed
|
||||
// and the worker struct will be freed after the timeout.
|
||||
kill(worker->Pid(), SIGTERM);
|
||||
worker->MarkAsBeingKilled();
|
||||
|
||||
auto retry_timer = std::make_shared<boost::asio::deadline_timer>(io_service_);
|
||||
auto retry_duration = boost::posix_time::milliseconds(
|
||||
@@ -272,13 +273,10 @@ void NodeManager::HandleJobTableUpdate(const JobID &id,
|
||||
auto workers = worker_pool_.GetWorkersRunningTasksForJob(job_id);
|
||||
|
||||
// Kill all the workers. The actual cleanup for these workers is done
|
||||
// later when we receive the DisconnectClient message from them.
|
||||
// later when the worker heartbeats timeout.
|
||||
for (const auto &worker : workers) {
|
||||
// Clean up any open ray.wait calls that the worker made.
|
||||
task_dependency_manager_.UnsubscribeWaitDependencies(worker->WorkerId());
|
||||
// Mark the worker as dead so further messages from it are ignored
|
||||
// (except DisconnectClient).
|
||||
worker->MarkDead();
|
||||
// Then kill the worker process.
|
||||
KillWorker(worker);
|
||||
}
|
||||
@@ -340,6 +338,18 @@ void NodeManager::Heartbeat() {
|
||||
last_debug_dump_at_ms_ = now_ms;
|
||||
}
|
||||
|
||||
// Check worker heartbeat timeout times.
|
||||
std::vector<std::shared_ptr<Worker>> dead_workers;
|
||||
worker_pool_.TickHeartbeatTimer(RayConfig::instance().num_worker_heartbeats_timeout(),
|
||||
&dead_workers);
|
||||
if (!dead_workers.empty()) {
|
||||
for (const auto &worker : dead_workers) {
|
||||
RAY_LOG(INFO) << "Worker " << worker->GetWorkerId()
|
||||
<< " dead because of timeout, pid: " << worker->Pid();
|
||||
ProcessDisconnectClientMessage(worker->GetWorkerId(), worker->IsBeingKilled());
|
||||
}
|
||||
}
|
||||
|
||||
// Reset the timer.
|
||||
heartbeat_timer_.expires_from_now(heartbeat_period_);
|
||||
heartbeat_timer_.async_wait([this](const boost::system::error_code &error) {
|
||||
@@ -682,9 +692,12 @@ void NodeManager::HandleActorStateTransition(const ActorID &actor_id,
|
||||
}
|
||||
}
|
||||
|
||||
void NodeManager::ProcessNewClient(LocalClientConnection &client) {
|
||||
// The new client is a worker, so begin listening for messages.
|
||||
client.ProcessMessages();
|
||||
void NodeManager::CleanUpTasksForFinishedJob(const JobID &job_id) {
|
||||
auto tasks_to_remove = local_queues_.GetTaskIdsForJob(job_id);
|
||||
task_dependency_manager_.RemoveTasksAndRelatedObjects(tasks_to_remove);
|
||||
// NOTE(swang): SchedulingQueue::RemoveTasks modifies its argument so we must
|
||||
// call it last.
|
||||
local_queues_.RemoveTasks(tasks_to_remove);
|
||||
}
|
||||
|
||||
// A helper function to create a mapping from resource shapes to
|
||||
@@ -701,7 +714,7 @@ std::unordered_map<ResourceSet, ordered_set<TaskID>> MakeTasksWithResources(
|
||||
|
||||
void NodeManager::DispatchTasks(
|
||||
const std::unordered_map<ResourceSet, ordered_set<TaskID>> &tasks_with_resources) {
|
||||
std::unordered_set<TaskID> removed_task_ids;
|
||||
std::unordered_set<TaskID> assigned_task_ids;
|
||||
for (const auto &it : tasks_with_resources) {
|
||||
const auto &task_resources = it.first;
|
||||
for (const auto &task_id : it.second) {
|
||||
@@ -712,146 +725,51 @@ void NodeManager::DispatchTasks(
|
||||
break;
|
||||
}
|
||||
if (AssignTask(task)) {
|
||||
removed_task_ids.insert(task_id);
|
||||
assigned_task_ids.insert(task_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Move the ASSIGNED task to the SWAP queue so that we remember that we have
|
||||
// it queued locally. Once the GetTaskReply has been sent, the task will get
|
||||
// re-queued, depending on whether the message succeeded or not.
|
||||
local_queues_.MoveTasks(removed_task_ids, TaskState::READY, TaskState::SWAP);
|
||||
|
||||
// Move the ASSIGNED task to the RUNNING queue.
|
||||
// We should move task outside `AssignTask` function because removing
|
||||
// task might influence the iterator.
|
||||
local_queues_.MoveTasks(assigned_task_ids, TaskState::READY, TaskState::RUNNING);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessClientMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client, int64_t message_type,
|
||||
const uint8_t *message_data) {
|
||||
auto registered_worker = worker_pool_.GetRegisteredWorker(client);
|
||||
auto message_type_value = static_cast<protocol::MessageType>(message_type);
|
||||
RAY_LOG(DEBUG) << "[Worker] Message "
|
||||
<< protocol::EnumNameMessageType(message_type_value) << "("
|
||||
<< message_type << ") from worker with PID "
|
||||
<< (registered_worker ? std::to_string(registered_worker->Pid())
|
||||
: "nil");
|
||||
if (registered_worker && registered_worker->IsDead()) {
|
||||
// For a worker that is marked as dead (because the job has died already),
|
||||
// all the messages are ignored except DisconnectClient.
|
||||
if ((message_type_value != protocol::MessageType::DisconnectClient) &&
|
||||
(message_type_value != protocol::MessageType::IntentionalDisconnectClient)) {
|
||||
// Listen for more messages.
|
||||
client->ProcessMessages();
|
||||
return;
|
||||
}
|
||||
}
|
||||
void NodeManager::HandleRegisterClientRequest(
|
||||
const rpc::RegisterClientRequest &request, rpc::RegisterClientReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
// Client id in register client is treated as worker id.
|
||||
const WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
auto worker =
|
||||
std::make_shared<Worker>(worker_id, request.worker_pid(), request.language(),
|
||||
request.port(), client_call_manager_);
|
||||
|
||||
switch (message_type_value) {
|
||||
case protocol::MessageType::RegisterClientRequest: {
|
||||
ProcessRegisterClientRequestMessage(client, message_data);
|
||||
} break;
|
||||
case protocol::MessageType::GetTask: {
|
||||
RAY_CHECK(!registered_worker->UsePush());
|
||||
HandleWorkerAvailable(client);
|
||||
} break;
|
||||
case protocol::MessageType::TaskDone: {
|
||||
RAY_CHECK(registered_worker->UsePush());
|
||||
HandleWorkerAvailable(client);
|
||||
} break;
|
||||
case protocol::MessageType::DisconnectClient: {
|
||||
ProcessDisconnectClientMessage(client);
|
||||
// We don't need to receive future messages from this client,
|
||||
// because it's already disconnected.
|
||||
return;
|
||||
} break;
|
||||
case protocol::MessageType::IntentionalDisconnectClient: {
|
||||
ProcessDisconnectClientMessage(client, /* intentional_disconnect = */ true);
|
||||
// We don't need to receive future messages from this client,
|
||||
// because it's already disconnected.
|
||||
return;
|
||||
} break;
|
||||
case protocol::MessageType::SubmitTask: {
|
||||
ProcessSubmitTaskMessage(message_data);
|
||||
} break;
|
||||
case protocol::MessageType::SetResourceRequest: {
|
||||
ProcessSetResourceRequest(client, message_data);
|
||||
} break;
|
||||
case protocol::MessageType::FetchOrReconstruct: {
|
||||
ProcessFetchOrReconstructMessage(client, message_data);
|
||||
} break;
|
||||
case protocol::MessageType::NotifyUnblocked: {
|
||||
auto message = flatbuffers::GetRoot<protocol::NotifyUnblocked>(message_data);
|
||||
HandleTaskUnblocked(client, from_flatbuf<TaskID>(*message->task_id()));
|
||||
} break;
|
||||
case protocol::MessageType::WaitRequest: {
|
||||
ProcessWaitRequestMessage(client, message_data);
|
||||
} break;
|
||||
case protocol::MessageType::PushErrorRequest: {
|
||||
ProcessPushErrorRequestMessage(message_data);
|
||||
} break;
|
||||
case protocol::MessageType::PushProfileEventsRequest: {
|
||||
auto fbs_message = flatbuffers::GetRoot<flatbuffers::String>(message_data);
|
||||
rpc::ProfileTableData profile_table_data;
|
||||
RAY_CHECK(
|
||||
profile_table_data.ParseFromArray(fbs_message->data(), fbs_message->size()));
|
||||
RAY_CHECK_OK(gcs_client_->profile_table().AddProfileEventBatch(profile_table_data));
|
||||
} break;
|
||||
case protocol::MessageType::FreeObjectsInObjectStoreRequest: {
|
||||
auto message = flatbuffers::GetRoot<protocol::FreeObjectsRequest>(message_data);
|
||||
std::vector<ObjectID> object_ids = from_flatbuf<ObjectID>(*message->object_ids());
|
||||
// Clean up objects from the object store.
|
||||
object_manager_.FreeObjects(object_ids, message->local_only());
|
||||
if (message->delete_creating_tasks()) {
|
||||
// Clean up their creating tasks from GCS.
|
||||
std::vector<TaskID> creating_task_ids;
|
||||
for (const auto &object_id : object_ids) {
|
||||
creating_task_ids.push_back(object_id.TaskId());
|
||||
}
|
||||
gcs_client_->raylet_task_table().Delete(JobID::Nil(), creating_task_ids);
|
||||
}
|
||||
} break;
|
||||
case protocol::MessageType::PrepareActorCheckpointRequest: {
|
||||
ProcessPrepareActorCheckpointRequest(client, message_data);
|
||||
} break;
|
||||
case protocol::MessageType::NotifyActorResumedFromCheckpoint: {
|
||||
ProcessNotifyActorResumedFromCheckpoint(message_data);
|
||||
} break;
|
||||
RAY_LOG(DEBUG) << "Received a RegisterClientRequest, worker id: " << worker_id
|
||||
<< ", is worker: " << request.is_worker()
|
||||
<< ", pid: " << request.worker_pid();
|
||||
|
||||
default:
|
||||
RAY_LOG(FATAL) << "Received unexpected message type " << message_type;
|
||||
}
|
||||
|
||||
// Listen for more messages.
|
||||
client->ProcessMessages();
|
||||
}
|
||||
|
||||
void NodeManager::ProcessRegisterClientRequestMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
|
||||
client->Register();
|
||||
auto message = flatbuffers::GetRoot<protocol::RegisterClientRequest>(message_data);
|
||||
Language language = static_cast<Language>(message->language());
|
||||
WorkerID worker_id = from_flatbuf<WorkerID>(*message->worker_id());
|
||||
auto worker = std::make_shared<Worker>(worker_id, message->worker_pid(), language,
|
||||
message->port(), client, client_call_manager_);
|
||||
if (message->is_worker()) {
|
||||
if (request.is_worker()) {
|
||||
// Register the new worker.
|
||||
bool use_push_task = worker->UsePush();
|
||||
auto connection = worker->Connection();
|
||||
worker_pool_.RegisterWorker(std::move(worker));
|
||||
worker_pool_.RegisterWorker(worker_id, std::move(worker));
|
||||
if (use_push_task) {
|
||||
// only call `HandleWorkerAvailable` when push mode is used.
|
||||
HandleWorkerAvailable(connection);
|
||||
HandleWorkerAvailable(worker_id);
|
||||
}
|
||||
} else {
|
||||
// Register the new driver.
|
||||
const JobID job_id = from_flatbuf<JobID>(*message->job_id());
|
||||
// Compute a dummy driver task id from a given driver.
|
||||
const TaskID driver_task_id = TaskID::ComputeDriverTaskId(worker_id);
|
||||
auto driver_task_id = TaskID::ComputeDriverTaskId(worker_id);
|
||||
auto job_id = JobID::FromBinary(request.job_id());
|
||||
worker->AssignTaskId(driver_task_id);
|
||||
worker->AssignJobId(job_id);
|
||||
worker_pool_.RegisterDriver(std::move(worker));
|
||||
worker_pool_.RegisterDriver(worker_id, std::move(worker));
|
||||
local_queues_.AddDriverTaskId(driver_task_id);
|
||||
RAY_CHECK_OK(gcs_client_->job_table().AppendJobData(
|
||||
job_id, /*is_dead=*/false, std::time(nullptr),
|
||||
initial_config_.node_manager_address, message->worker_pid()));
|
||||
initial_config_.node_manager_address, request.worker_pid()));
|
||||
}
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleDisconnectedActor(const ActorID &actor_id, bool was_local,
|
||||
@@ -882,8 +800,8 @@ void NodeManager::HandleDisconnectedActor(const ActorID &actor_id, bool was_loca
|
||||
if (was_local) {
|
||||
// If the actor was local, immediately update the state in actor registry.
|
||||
// So if we receive any actor tasks before we receive GCS notification,
|
||||
// these tasks can be correctly routed to the `MethodsWaitingForActorCreation` queue,
|
||||
// instead of being assigned to the dead actor.
|
||||
// these tasks can be correctly routed to the `MethodsWaitingForActorCreation`
|
||||
// queue, instead of being assigned to the dead actor.
|
||||
HandleActorStateTransition(actor_id, ActorRegistration(new_actor_data));
|
||||
}
|
||||
|
||||
@@ -898,40 +816,62 @@ void NodeManager::HandleDisconnectedActor(const ActorID &actor_id, bool was_loca
|
||||
RAY_CHECK_OK(gcs_client_->Actors().AsyncUpdate(actor_id, actor_notification, done));
|
||||
}
|
||||
|
||||
void NodeManager::HandleWorkerAvailable(
|
||||
const std::shared_ptr<LocalClientConnection> &client) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
|
||||
RAY_CHECK(worker);
|
||||
// If the worker was assigned a task, mark it as finished.
|
||||
if (!worker->GetAssignedTaskId().IsNil()) {
|
||||
FinishAssignedTask(*worker);
|
||||
void NodeManager::HandleGetTaskRequest(const rpc::GetTaskRequest &request,
|
||||
rpc::GetTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
const WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
RAY_LOG(DEBUG) << "Received a GetTaskRequest, worker id " << worker_id << " pid "
|
||||
<< worker->Pid();
|
||||
if (!worker || worker->IsBeingKilled()) {
|
||||
send_reply_callback(Status::Invalid("WorkerBeingKilled"), nullptr, nullptr);
|
||||
return;
|
||||
}
|
||||
RAY_CHECK(!worker->UsePush());
|
||||
|
||||
// Return the worker to the idle pool.
|
||||
worker_pool_.PushWorker(std::move(worker));
|
||||
// Local resource availability changed: invoke scheduling policy for local node.
|
||||
const ClientID &local_client_id = gcs_client_->client_table().GetLocalClientId();
|
||||
cluster_resource_map_[local_client_id].SetLoadResources(
|
||||
local_queues_.GetResourceLoad());
|
||||
// Call task dispatch to assign work to the new worker.
|
||||
DispatchTasks(local_queues_.GetReadyTasksWithResources());
|
||||
// Reply would be sent when assigned a task to the worker successfully.
|
||||
worker->SetGetTaskReplyAndCallback(reply, std::move(send_reply_callback));
|
||||
HandleWorkerAvailable(worker_id);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessDisconnectClientMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client, bool intentional_disconnect) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
|
||||
void NodeManager::HandleTaskDoneRequest(const rpc::TaskDoneRequest &request,
|
||||
rpc::TaskDoneReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
const WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
RAY_LOG(DEBUG) << "Received a TaskDoneRequest from worker " << worker_id;
|
||||
|
||||
auto worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
RAY_CHECK(worker && worker->UsePush());
|
||||
HandleWorkerAvailable(worker_id);
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleDisconnectClientRequest(
|
||||
const rpc::DisconnectClientRequest &request, rpc::DisconnectClientReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
const WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
RAY_LOG(DEBUG) << "Received a DisconnectClientRequest from worker " << worker_id;
|
||||
|
||||
ProcessDisconnectClientMessage(worker_id, true);
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessDisconnectClientMessage(const WorkerID &worker_id,
|
||||
bool intentional_disconnect) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
bool is_worker = false, is_driver = false;
|
||||
if (worker) {
|
||||
// The client is a worker.
|
||||
is_worker = true;
|
||||
} else {
|
||||
worker = worker_pool_.GetRegisteredDriver(client);
|
||||
worker = worker_pool_.GetRegisteredDriver(worker_id);
|
||||
if (worker) {
|
||||
// The client is a driver.
|
||||
is_driver = true;
|
||||
} else {
|
||||
RAY_LOG(INFO) << "Ignoring client disconnect because the client has already "
|
||||
<< "been disconnected.";
|
||||
return;
|
||||
}
|
||||
}
|
||||
RAY_CHECK(!(is_worker && is_driver));
|
||||
@@ -939,9 +879,9 @@ void NodeManager::ProcessDisconnectClientMessage(
|
||||
// If the client has any blocked tasks, mark them as unblocked. In
|
||||
// particular, we are no longer waiting for their dependencies.
|
||||
if (worker) {
|
||||
if (is_worker && worker->IsDead()) {
|
||||
// Don't need to unblock the client if it's a worker and is already dead.
|
||||
// Because in this case, its task is already cleaned up.
|
||||
if (is_worker && worker->IsBeingKilled()) {
|
||||
// Don't need to unblock the client if it's a worker and have sent kill signal to
|
||||
// it. Because in this case, its task is already cleaned up.
|
||||
RAY_LOG(DEBUG) << "Skip unblocking worker because it's already dead.";
|
||||
} else {
|
||||
// Clean up any open ray.get calls that the worker made.
|
||||
@@ -949,7 +889,7 @@ void NodeManager::ProcessDisconnectClientMessage(
|
||||
// NOTE(swang): HandleTaskUnblocked will modify the worker, so it is
|
||||
// not safe to pass in the iterator directly.
|
||||
const TaskID task_id = *worker->GetBlockedTaskIds().begin();
|
||||
HandleTaskUnblocked(client, task_id);
|
||||
HandleTaskUnblocked(worker_id, task_id);
|
||||
}
|
||||
// Clean up any open ray.wait calls that the worker made.
|
||||
task_dependency_manager_.UnsubscribeWaitDependencies(worker->WorkerId());
|
||||
@@ -957,24 +897,17 @@ void NodeManager::ProcessDisconnectClientMessage(
|
||||
}
|
||||
|
||||
if (is_worker) {
|
||||
// The client is a worker.
|
||||
if (worker->IsDead()) {
|
||||
// If the worker was killed by us because the driver exited,
|
||||
// treat it as intentionally disconnected.
|
||||
intentional_disconnect = true;
|
||||
}
|
||||
|
||||
const ActorID &actor_id = worker->GetActorId();
|
||||
if (!actor_id.IsNil()) {
|
||||
// If the worker was an actor, update actor state, reconstruct the actor if needed,
|
||||
// and clean up actor's tasks if the actor is permanently dead.
|
||||
// If the worker was an actor, update actor state, reconstruct the actor if
|
||||
// needed, and clean up actor's tasks if the actor is permanently dead.
|
||||
HandleDisconnectedActor(actor_id, true, intentional_disconnect);
|
||||
}
|
||||
|
||||
const TaskID &task_id = worker->GetAssignedTaskId();
|
||||
// If the worker was running a task, clean up the task and push an error to
|
||||
// the driver, unless the worker is already dead.
|
||||
if (!task_id.IsNil() && !worker->IsDead()) {
|
||||
// the driver, unless the worker is already being killed.
|
||||
if (!task_id.IsNil() && !worker->IsBeingKilled()) {
|
||||
// If the worker was an actor, the task was already cleaned up in
|
||||
// `HandleDisconnectedActor`.
|
||||
if (actor_id.IsNil()) {
|
||||
@@ -1035,32 +968,34 @@ void NodeManager::ProcessDisconnectClientMessage(
|
||||
<< "job_id: " << job_id;
|
||||
}
|
||||
|
||||
client->Close();
|
||||
|
||||
// TODO(rkn): Tell the object manager that this client has disconnected so
|
||||
// that it can clean up the wait requests for this client. Currently I think
|
||||
// these can be leaked.
|
||||
}
|
||||
|
||||
void NodeManager::ProcessSubmitTaskMessage(const uint8_t *message_data) {
|
||||
// Read the task submitted by the client.
|
||||
auto fbs_message = flatbuffers::GetRoot<protocol::SubmitTaskRequest>(message_data);
|
||||
rpc::Task task_message;
|
||||
RAY_CHECK(task_message.mutable_task_spec()->ParseFromArray(
|
||||
fbs_message->task_spec()->data(), fbs_message->task_spec()->size()));
|
||||
void NodeManager::HandleSubmitTaskRequest(const rpc::SubmitTaskRequest &request,
|
||||
rpc::SubmitTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a SubmitTaskRequest.";
|
||||
|
||||
rpc::Task task;
|
||||
task.mutable_task_spec()->CopyFrom(request.task_spec());
|
||||
|
||||
// Submit the task to the raylet. Since the task was submitted
|
||||
// locally, there is no uncommitted lineage.
|
||||
SubmitTask(Task(task_message), Lineage());
|
||||
SubmitTask(Task(task), Lineage());
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessFetchOrReconstructMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
|
||||
auto message = flatbuffers::GetRoot<protocol::FetchOrReconstruct>(message_data);
|
||||
void NodeManager::HandleFetchOrReconstructRequest(
|
||||
const rpc::FetchOrReconstructRequest &request, rpc::FetchOrReconstructReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a FetchOrReconstructRequest.";
|
||||
const auto &object_ids = request.object_ids();
|
||||
std::vector<ObjectID> required_object_ids;
|
||||
for (size_t i = 0; i < message->object_ids()->size(); ++i) {
|
||||
ObjectID object_id = from_flatbuf<ObjectID>(*message->object_ids()->Get(i));
|
||||
if (message->fetch_only()) {
|
||||
for (size_t i = 0; i < object_ids.size(); ++i) {
|
||||
ObjectID object_id = ObjectID::FromBinary(object_ids[i]);
|
||||
if (request.fetch_only()) {
|
||||
// If only a fetch is required, then do not subscribe to the
|
||||
// dependencies to the task dependency manager.
|
||||
if (!task_dependency_manager_.CheckObjectLocal(object_id)) {
|
||||
@@ -1077,19 +1012,22 @@ void NodeManager::ProcessFetchOrReconstructMessage(
|
||||
}
|
||||
|
||||
if (!required_object_ids.empty()) {
|
||||
const TaskID task_id = from_flatbuf<TaskID>(*message->task_id());
|
||||
HandleTaskBlocked(client, required_object_ids, task_id, /*ray_get=*/true);
|
||||
const TaskID task_id = TaskID::FromBinary(request.task_id());
|
||||
const WorkerID &worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
HandleTaskBlocked(worker_id, required_object_ids, task_id, /*ray_get=*/true);
|
||||
}
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessWaitRequestMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
|
||||
void NodeManager::HandleWaitRequest(const rpc::WaitRequest &request,
|
||||
rpc::WaitReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a WaitRequest.";
|
||||
// Read the data.
|
||||
auto message = flatbuffers::GetRoot<protocol::WaitRequest>(message_data);
|
||||
std::vector<ObjectID> object_ids = from_flatbuf<ObjectID>(*message->object_ids());
|
||||
int64_t wait_ms = message->timeout();
|
||||
uint64_t num_required_objects = static_cast<uint64_t>(message->num_ready_objects());
|
||||
bool wait_local = message->wait_local();
|
||||
std::vector<ObjectID> object_ids = IdVectorFromProtobuf<ObjectID>(request.object_ids());
|
||||
int64_t wait_ms = request.timeout();
|
||||
uint64_t num_required_objects = request.num_ready_objects();
|
||||
bool wait_local = request.wait_local();
|
||||
|
||||
std::vector<ObjectID> required_object_ids;
|
||||
for (auto const &object_id : object_ids) {
|
||||
@@ -1101,63 +1039,56 @@ void NodeManager::ProcessWaitRequestMessage(
|
||||
}
|
||||
}
|
||||
|
||||
const TaskID ¤t_task_id = from_flatbuf<TaskID>(*message->task_id());
|
||||
const TaskID ¤t_task_id = TaskID::FromBinary(request.task_id());
|
||||
const WorkerID &worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
bool client_blocked = !required_object_ids.empty();
|
||||
if (client_blocked) {
|
||||
HandleTaskBlocked(client, required_object_ids, current_task_id, /*ray_get=*/false);
|
||||
HandleTaskBlocked(worker_id, required_object_ids, current_task_id, /*ray_get=*/false);
|
||||
}
|
||||
|
||||
ray::Status status = object_manager_.Wait(
|
||||
object_ids, wait_ms, num_required_objects, wait_local,
|
||||
[this, client_blocked, client, current_task_id](std::vector<ObjectID> found,
|
||||
std::vector<ObjectID> remaining) {
|
||||
// Write the data.
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
flatbuffers::Offset<protocol::WaitReply> wait_reply = protocol::CreateWaitReply(
|
||||
fbb, to_flatbuf(fbb, found), to_flatbuf(fbb, remaining));
|
||||
fbb.Finish(wait_reply);
|
||||
[this, client_blocked, worker_id, current_task_id, reply, send_reply_callback](
|
||||
std::vector<ObjectID> found, std::vector<ObjectID> remaining) {
|
||||
IdVectorToProtobuf<ObjectID, rpc::WaitReply>(found, *reply,
|
||||
&rpc::WaitReply::add_found);
|
||||
IdVectorToProtobuf<ObjectID, rpc::WaitReply>(remaining, *reply,
|
||||
&rpc::WaitReply::add_remaining);
|
||||
|
||||
auto status =
|
||||
client->WriteMessage(static_cast<int64_t>(protocol::MessageType::WaitReply),
|
||||
fbb.GetSize(), fbb.GetBufferPointer());
|
||||
if (status.ok()) {
|
||||
// The client is unblocked now because the wait call has returned.
|
||||
if (client_blocked) {
|
||||
HandleTaskUnblocked(client, current_task_id);
|
||||
}
|
||||
} else {
|
||||
// We failed to write to the client, so disconnect the client.
|
||||
RAY_LOG(WARNING)
|
||||
<< "Failed to send WaitReply to client, so disconnecting client";
|
||||
// We failed to send the reply to the client, so disconnect the worker.
|
||||
ProcessDisconnectClientMessage(client);
|
||||
// Send reply to finish this wait request.
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
if (client_blocked) {
|
||||
HandleTaskUnblocked(worker_id, current_task_id);
|
||||
}
|
||||
});
|
||||
RAY_CHECK_OK(status);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessPushErrorRequestMessage(const uint8_t *message_data) {
|
||||
auto message = flatbuffers::GetRoot<protocol::PushErrorRequest>(message_data);
|
||||
|
||||
JobID job_id = from_flatbuf<JobID>(*message->job_id());
|
||||
auto const &type = string_from_flatbuf(*message->type());
|
||||
auto const &error_message = string_from_flatbuf(*message->error_message());
|
||||
double timestamp = message->timestamp();
|
||||
void NodeManager::HandlePushErrorRequest(const rpc::PushErrorRequest &request,
|
||||
rpc::PushErrorReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
JobID job_id = JobID::FromBinary(request.job_id());
|
||||
const auto &type = request.type();
|
||||
const auto &error_message = request.error_message();
|
||||
double timestamp = request.timestamp();
|
||||
RAY_LOG(DEBUG) << "Handle push error request for job " << job_id << ", type " << type
|
||||
<< " error message " << error_message;
|
||||
|
||||
RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(job_id, type, error_message,
|
||||
timestamp));
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessPrepareActorCheckpointRequest(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
|
||||
auto message =
|
||||
flatbuffers::GetRoot<protocol::PrepareActorCheckpointRequest>(message_data);
|
||||
ActorID actor_id = from_flatbuf<ActorID>(*message->actor_id());
|
||||
void NodeManager::HandlePrepareActorCheckpointRequest(
|
||||
const rpc::PrepareActorCheckpointRequest &request,
|
||||
rpc::PrepareActorCheckpointReply *reply, rpc::SendReplyCallback send_reply_callback) {
|
||||
ActorID actor_id = ActorID::FromBinary(request.actor_id());
|
||||
RAY_LOG(DEBUG) << "Preparing checkpoint for actor " << actor_id;
|
||||
const auto &actor_entry = actor_registry_.find(actor_id);
|
||||
RAY_CHECK(actor_entry != actor_registry_.end());
|
||||
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
|
||||
WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
RAY_CHECK(worker && worker->GetActorId() == actor_id);
|
||||
|
||||
// Find the task that is running on this actor.
|
||||
@@ -1171,44 +1102,34 @@ void NodeManager::ProcessPrepareActorCheckpointRequest(
|
||||
// Write checkpoint data to GCS.
|
||||
RAY_CHECK_OK(gcs_client_->actor_checkpoint_table().Add(
|
||||
JobID::Nil(), checkpoint_id, checkpoint_data,
|
||||
[worker, actor_id, this](ray::gcs::RedisGcsClient *client,
|
||||
const ActorCheckpointID &checkpoint_id,
|
||||
const ActorCheckpointData &data) {
|
||||
[worker, actor_id, reply, send_reply_callback, this](
|
||||
ray::gcs::RedisGcsClient *client, const ActorCheckpointID &checkpoint_id,
|
||||
const ActorCheckpointData &data) {
|
||||
RAY_LOG(DEBUG) << "Checkpoint " << checkpoint_id << " saved for actor "
|
||||
<< worker->GetActorId();
|
||||
// Save this actor-to-checkpoint mapping, and remove old checkpoints associated
|
||||
// with this actor.
|
||||
// Save this actor-to-checkpoint mapping, and remove old checkpoints
|
||||
// associated with this actor.
|
||||
RAY_CHECK_OK(gcs_client_->actor_checkpoint_id_table().AddCheckpointId(
|
||||
JobID::Nil(), actor_id, checkpoint_id));
|
||||
// Send reply to worker.
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto reply = ray::protocol::CreatePrepareActorCheckpointReply(
|
||||
fbb, to_flatbuf(fbb, checkpoint_id));
|
||||
fbb.Finish(reply);
|
||||
worker->Connection()->WriteMessageAsync(
|
||||
static_cast<int64_t>(protocol::MessageType::PrepareActorCheckpointReply),
|
||||
fbb.GetSize(), fbb.GetBufferPointer(), [](const ray::Status &status) {
|
||||
if (!status.ok()) {
|
||||
RAY_LOG(WARNING)
|
||||
<< "Failed to send PrepareActorCheckpointReply to client";
|
||||
}
|
||||
});
|
||||
reply->set_checkpoint_id(checkpoint_id.Binary());
|
||||
send_reply_callback(Status::OK(), nullptr, []() {
|
||||
RAY_LOG(WARNING) << "Failed to send PrepareActorCheckpointReply to client";
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
void NodeManager::ProcessNotifyActorResumedFromCheckpoint(const uint8_t *message_data) {
|
||||
auto message =
|
||||
flatbuffers::GetRoot<protocol::NotifyActorResumedFromCheckpoint>(message_data);
|
||||
ActorID actor_id = from_flatbuf<ActorID>(*message->actor_id());
|
||||
void NodeManager::HandleNotifyActorResumedFromCheckpointRequest(
|
||||
const rpc::NotifyActorResumedFromCheckpointRequest &request,
|
||||
rpc::NotifyActorResumedFromCheckpointReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
ActorID actor_id = ActorID::FromBinary(request.actor_id());
|
||||
ActorCheckpointID checkpoint_id =
|
||||
from_flatbuf<ActorCheckpointID>(*message->checkpoint_id());
|
||||
ActorCheckpointID::FromBinary(request.checkpoint_id());
|
||||
RAY_LOG(DEBUG) << "Actor " << actor_id << " was resumed from checkpoint "
|
||||
<< checkpoint_id;
|
||||
checkpoint_id_to_restore_.emplace(actor_id, checkpoint_id);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessNewNodeManager(TcpClientConnection &node_manager_client) {
|
||||
node_manager_client.ProcessMessages();
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleForwardTask(const rpc::ForwardTaskRequest &request,
|
||||
@@ -1229,16 +1150,14 @@ void NodeManager::HandleForwardTask(const rpc::ForwardTaskRequest &request,
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ProcessSetResourceRequest(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const uint8_t *message_data) {
|
||||
// Read the SetResource message
|
||||
auto message = flatbuffers::GetRoot<protocol::SetResourceRequest>(message_data);
|
||||
|
||||
auto const &resource_name = string_from_flatbuf(*message->resource_name());
|
||||
double const &capacity = message->capacity();
|
||||
void NodeManager::HandleSetResourceRequest(const rpc::SetResourceRequest &request,
|
||||
rpc::SetResourceReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
auto const &resource_name = request.resource_name();
|
||||
double const capacity = request.capacity();
|
||||
bool is_deletion = capacity <= 0;
|
||||
|
||||
ClientID client_id = from_flatbuf<ClientID>(*message->client_id());
|
||||
ClientID client_id = ClientID::FromBinary(request.client_id());
|
||||
|
||||
// If the python arg was null, set client_id to the local client
|
||||
if (client_id.IsNil()) {
|
||||
@@ -1255,7 +1174,7 @@ void NodeManager::ProcessSetResourceRequest(
|
||||
return;
|
||||
}
|
||||
|
||||
// Submit to the client table. This calls the ResourceCreateUpdated or ResourceDeleted
|
||||
// Submit to the resource table. This calls the ResourceCreateUpdated or ResourceDeleted
|
||||
// callback, which updates cluster_resource_map_.
|
||||
if (is_deletion) {
|
||||
RAY_CHECK_OK(gcs_client_->resource_table().RemoveEntries(JobID::Nil(), client_id,
|
||||
@@ -1268,6 +1187,62 @@ void NodeManager::ProcessSetResourceRequest(
|
||||
RAY_CHECK_OK(
|
||||
gcs_client_->resource_table().Update(JobID::Nil(), client_id, data_map, nullptr));
|
||||
}
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleNotifyUnblockedRequest(
|
||||
const rpc::NotifyUnblockedRequest &request, rpc::NotifyUnblockedReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a NotifyUnblockedRequest.";
|
||||
const TaskID current_task_id = TaskID::FromBinary(request.task_id());
|
||||
const WorkerID worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
|
||||
HandleTaskUnblocked(worker_id, current_task_id);
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandlePushProfileEventsRequest(
|
||||
const rpc::PushProfileEventsRequest &request, rpc::PushProfileEventsReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a PushProfileEventsRequest.";
|
||||
const auto &profile_table_data = request.profile_table_data();
|
||||
RAY_CHECK_OK(gcs_client_->profile_table().AddProfileEventBatch(profile_table_data));
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleFreeObjectsInStoreRequest(
|
||||
const rpc::FreeObjectsInStoreRequest &request, rpc::FreeObjectsInStoreReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
RAY_LOG(DEBUG) << "Received a FreeObjectsInStoreRequest.";
|
||||
std::vector<ObjectID> object_ids = IdVectorFromProtobuf<ObjectID>(request.object_ids());
|
||||
object_manager_.FreeObjects(object_ids, request.local_only());
|
||||
if (request.delete_creating_tasks()) {
|
||||
// Clean up their creating tasks from GCS.
|
||||
std::vector<TaskID> creating_task_ids;
|
||||
for (const auto &object_id : object_ids) {
|
||||
creating_task_ids.push_back(object_id.TaskId());
|
||||
}
|
||||
gcs_client_->raylet_task_table().Delete(JobID::Nil(), creating_task_ids);
|
||||
}
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::HandleHeartbeatRequest(const rpc::HeartbeatRequest &request,
|
||||
rpc::HeartbeatReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) {
|
||||
bool is_worker = request.is_worker();
|
||||
const auto worker_id = WorkerID::FromBinary(request.worker_id());
|
||||
|
||||
std::shared_ptr<Worker> worker = nullptr;
|
||||
if (is_worker) {
|
||||
worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
} else {
|
||||
worker = worker_pool_.GetRegisteredDriver(worker_id);
|
||||
}
|
||||
if (worker) {
|
||||
worker->ClearHeartbeat();
|
||||
}
|
||||
send_reply_callback(Status::OK(), nullptr, nullptr);
|
||||
}
|
||||
|
||||
void NodeManager::ScheduleTasks(
|
||||
@@ -1574,7 +1549,8 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag
|
||||
// This is a non-actor task. Queue the task for a placement decision or for dispatch
|
||||
// if the task was forwarded.
|
||||
if (forwarded) {
|
||||
// Check for local dependencies and enqueue as waiting or ready for dispatch.
|
||||
// Check for local dependencies and enqueue as waiting or ready for
|
||||
// dispatch.
|
||||
EnqueuePlaceableTask(task);
|
||||
} else {
|
||||
// (See design_docs/task_states.rst for the state transition diagram.)
|
||||
@@ -1586,10 +1562,10 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag
|
||||
}
|
||||
}
|
||||
|
||||
void NodeManager::HandleTaskBlocked(const std::shared_ptr<LocalClientConnection> &client,
|
||||
void NodeManager::HandleTaskBlocked(const WorkerID &worker_id,
|
||||
const std::vector<ObjectID> &required_object_ids,
|
||||
const TaskID ¤t_task_id, bool ray_get) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
if (worker) {
|
||||
// The client is a worker. If the worker is not already blocked and the
|
||||
// blocked task matches the one assigned to the worker, then mark the
|
||||
@@ -1616,7 +1592,7 @@ void NodeManager::HandleTaskBlocked(const std::shared_ptr<LocalClientConnection>
|
||||
} else {
|
||||
// The client is a driver. Drivers do not hold resources, so we simply mark
|
||||
// the task as blocked.
|
||||
worker = worker_pool_.GetRegisteredDriver(client);
|
||||
worker = worker_pool_.GetRegisteredDriver(worker_id);
|
||||
}
|
||||
|
||||
RAY_CHECK(worker);
|
||||
@@ -1638,9 +1614,9 @@ void NodeManager::HandleTaskBlocked(const std::shared_ptr<LocalClientConnection>
|
||||
}
|
||||
}
|
||||
|
||||
void NodeManager::HandleTaskUnblocked(
|
||||
const std::shared_ptr<LocalClientConnection> &client, const TaskID ¤t_task_id) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
|
||||
void NodeManager::HandleTaskUnblocked(const WorkerID &worker_id,
|
||||
const TaskID ¤t_task_id) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
|
||||
// TODO(swang): Because the object dependencies are tracked in the task
|
||||
// dependency manager, we could actually remove this message entirely and
|
||||
@@ -1670,9 +1646,9 @@ void NodeManager::HandleTaskUnblocked(
|
||||
cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Acquire(
|
||||
cpu_resources);
|
||||
} else {
|
||||
// In this case, we simply don't reacquire the CPU resources for the worker.
|
||||
// The worker can keep running and when the task finishes, it will simply
|
||||
// not have any CPU resources to release.
|
||||
// In this case, we simply don't reacquire the CPU resources for the
|
||||
// worker. The worker can keep running and when the task finishes, it will
|
||||
// simply not have any CPU resources to release.
|
||||
RAY_LOG(WARNING)
|
||||
<< "Resources oversubscribed: "
|
||||
<< cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()]
|
||||
@@ -1684,7 +1660,7 @@ void NodeManager::HandleTaskUnblocked(
|
||||
} else {
|
||||
// The client is a driver. Drivers do not hold resources, so we simply
|
||||
// mark the driver as unblocked.
|
||||
worker = worker_pool_.GetRegisteredDriver(client);
|
||||
worker = worker_pool_.GetRegisteredDriver(worker_id);
|
||||
}
|
||||
|
||||
// Unsubscribe from any `ray.get` objects that the task was blocked on. Any
|
||||
@@ -1719,6 +1695,24 @@ void NodeManager::EnqueuePlaceableTask(const Task &task) {
|
||||
task_dependency_manager_.TaskPending(task);
|
||||
}
|
||||
|
||||
void NodeManager::HandleWorkerAvailable(const WorkerID &worker_id) {
|
||||
std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(worker_id);
|
||||
RAY_CHECK(worker);
|
||||
// If the worker was assigned a task, mark it as finished.
|
||||
if (!worker->GetAssignedTaskId().IsNil()) {
|
||||
FinishAssignedTask(*worker);
|
||||
}
|
||||
|
||||
// Return the worker to the idle pool.
|
||||
worker_pool_.PushWorker(std::move(worker));
|
||||
// Local resource availability changed: invoke scheduling policy for local node.
|
||||
const ClientID &local_client_id = gcs_client_->client_table().GetLocalClientId();
|
||||
cluster_resource_map_[local_client_id].SetLoadResources(
|
||||
local_queues_.GetResourceLoad());
|
||||
// Call task dispatch to assign work to the new worker.
|
||||
DispatchTasks(local_queues_.GetReadyTasksWithResources());
|
||||
}
|
||||
|
||||
bool NodeManager::AssignTask(const Task &task) {
|
||||
const TaskSpecification &spec = task.GetTaskSpecification();
|
||||
|
||||
@@ -1738,12 +1732,12 @@ bool NodeManager::AssignTask(const Task &task) {
|
||||
if (worker == nullptr) {
|
||||
// There are no workers that can execute this task.
|
||||
// We couldn't assign this task, as no worker available.
|
||||
RAY_LOG(DEBUG) << "No idle worker is found to assign task " << spec.TaskId();
|
||||
return false;
|
||||
}
|
||||
|
||||
RAY_LOG(DEBUG) << "Assigning task " << spec.TaskId() << " to worker with pid "
|
||||
<< worker->Pid();
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
RAY_LOG(DEBUG) << "Assigning task " << spec.TaskId() << " to worker "
|
||||
<< worker->GetWorkerId() << " pid " << worker->Pid();
|
||||
|
||||
// Resource accounting: acquire resources for the assigned task.
|
||||
auto acquired_resources =
|
||||
@@ -1760,40 +1754,44 @@ bool NodeManager::AssignTask(const Task &task) {
|
||||
worker->SetTaskResourceIds(acquired_resources);
|
||||
}
|
||||
|
||||
auto task_id = spec.TaskId();
|
||||
auto finish_assign_task_callback = [this, worker, task_id](Status status) {
|
||||
if (worker->UsePush()) {
|
||||
// NOTE: we cannot directly call `FinishAssignTask` here because
|
||||
// it assumes the task is in SWAP queue, thus we need to delay invoking this
|
||||
// function after the assigned tasks are moved from READY queue to SWAP queue
|
||||
// in `DispatchTasks`.
|
||||
// Another option is to move the tasks to SWAP queue here just before calling
|
||||
// `FinishAssignTask` so we can save an io_service post, at the
|
||||
// expense of calling `MoveTask` for each of the assigned tasks.
|
||||
// TODO(zhijunfu): after all workers are fully migrated to push mode, the
|
||||
// `post` below and swap queue can be removed.
|
||||
io_service_.post([this, status, worker, task_id]() {
|
||||
FinishAssignTask(task_id, *worker, status.ok());
|
||||
});
|
||||
} else {
|
||||
FinishAssignTask(task_id, *worker, status.ok());
|
||||
}
|
||||
};
|
||||
|
||||
ResourceIdSet resource_id_set =
|
||||
worker->GetTaskResourceIds().Plus(worker->GetLifetimeResourceIds());
|
||||
|
||||
worker->AssignTask(task, resource_id_set, finish_assign_task_callback);
|
||||
worker->AssignTask(task, resource_id_set);
|
||||
// Actor tasks require extra accounting to track the actor's state.
|
||||
if (spec.IsActorTask()) {
|
||||
auto actor_entry = actor_registry_.find(spec.ActorId());
|
||||
RAY_CHECK(actor_entry != actor_registry_.end());
|
||||
// Process any new actor handles that were created since the
|
||||
// previous task on this handle was executed. The first task
|
||||
// submitted on a new actor handle will depend on the dummy object
|
||||
// returned by the previous task, so the dependency will not be
|
||||
// released until this first task is submitted.
|
||||
for (auto &new_handle_id : spec.NewActorHandles()) {
|
||||
const auto prev_actor_task_id = spec.PreviousActorTaskDummyObjectId();
|
||||
RAY_CHECK(!prev_actor_task_id.IsNil());
|
||||
// Add the new handle and give it a reference to the finished task's
|
||||
// execution dependency.
|
||||
actor_entry->second.AddHandle(new_handle_id, prev_actor_task_id);
|
||||
}
|
||||
|
||||
// TODO(swang): For actors with multiple actor handles, to
|
||||
// guarantee that tasks are replayed in the same order after a
|
||||
// failure, we must update the task's execution dependency to be
|
||||
// the actor's current execution dependency.
|
||||
}
|
||||
|
||||
// Notify the task dependency manager that we no longer need this task's
|
||||
// object dependencies.
|
||||
RAY_CHECK(task_dependency_manager_.UnsubscribeGetDependencies(spec.TaskId()));
|
||||
|
||||
// We assigned this task to a worker.
|
||||
// (Note this means that we sent the task to the worker. The assignment
|
||||
// might still fail if the worker fails in the meantime, for instance.)
|
||||
return true;
|
||||
}
|
||||
|
||||
void NodeManager::FinishAssignedTask(Worker &worker) {
|
||||
TaskID task_id = worker.GetAssignedTaskId();
|
||||
RAY_LOG(DEBUG) << "Finished task " << task_id;
|
||||
RAY_LOG(DEBUG) << "Finished task " << task_id << " from worker " << worker.GetWorkerId()
|
||||
<< " with pid " << worker.Pid();
|
||||
|
||||
// (See design_docs/task_states.rst for the state transition diagram.)
|
||||
Task task;
|
||||
@@ -2202,8 +2200,9 @@ void NodeManager::ForwardTaskOrResubmit(const Task &task,
|
||||
retry_timer->expires_from_now(retry_duration);
|
||||
retry_timer->async_wait(
|
||||
[this, task_id, retry_timer](const boost::system::error_code &error) {
|
||||
// Timer killing will receive the boost::asio::error::operation_aborted,
|
||||
// we only handle the timeout event.
|
||||
// Timer killing will receive the
|
||||
// boost::asio::error::operation_aborted, we only handle the timeout
|
||||
// event.
|
||||
RAY_CHECK(!error);
|
||||
RAY_LOG(INFO) << "Resubmitting task " << task_id
|
||||
<< " because ForwardTask failed.";
|
||||
@@ -2325,63 +2324,6 @@ void NodeManager::ForwardTask(
|
||||
});
|
||||
}
|
||||
|
||||
void NodeManager::FinishAssignTask(const TaskID &task_id, Worker &worker, bool success) {
|
||||
// Remove the ASSIGNED task from the SWAP queue.
|
||||
Task assigned_task;
|
||||
TaskState state;
|
||||
if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) {
|
||||
return;
|
||||
}
|
||||
|
||||
RAY_CHECK(state == TaskState::SWAP);
|
||||
|
||||
if (success) {
|
||||
auto spec = assigned_task.GetTaskSpecification();
|
||||
// We successfully assigned the task to the worker.
|
||||
worker.AssignTaskId(spec.TaskId());
|
||||
worker.AssignJobId(spec.JobId());
|
||||
// Actor tasks require extra accounting to track the actor's state.
|
||||
if (spec.IsActorTask()) {
|
||||
auto actor_entry = actor_registry_.find(spec.ActorId());
|
||||
RAY_CHECK(actor_entry != actor_registry_.end());
|
||||
// Process any new actor handles that were created since the
|
||||
// previous task on this handle was executed. The first task
|
||||
// submitted on a new actor handle will depend on the dummy object
|
||||
// returned by the previous task, so the dependency will not be
|
||||
// released until this first task is submitted.
|
||||
for (auto &new_handle_id : spec.NewActorHandles()) {
|
||||
const auto prev_actor_task_id = spec.PreviousActorTaskDummyObjectId();
|
||||
RAY_CHECK(!prev_actor_task_id.IsNil());
|
||||
// Add the new handle and give it a reference to the finished task's
|
||||
// execution dependency.
|
||||
actor_entry->second.AddHandle(new_handle_id, prev_actor_task_id);
|
||||
}
|
||||
|
||||
// TODO(swang): For actors with multiple actor handles, to
|
||||
// guarantee that tasks are replayed in the same order after a
|
||||
// failure, we must update the task's execution dependency to be
|
||||
// the actor's current execution dependency.
|
||||
}
|
||||
|
||||
// Mark the task as running.
|
||||
// (See design_docs/task_states.rst for the state transition diagram.)
|
||||
local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING);
|
||||
// Notify the task dependency manager that we no longer need this task's
|
||||
// object dependencies.
|
||||
RAY_CHECK(task_dependency_manager_.UnsubscribeGetDependencies(spec.TaskId()));
|
||||
} else {
|
||||
RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client";
|
||||
// We failed to send the task to the worker, so disconnect the worker.
|
||||
ProcessDisconnectClientMessage(worker.Connection());
|
||||
// Queue this task for future assignment. We need to do this since
|
||||
// DispatchTasks() removed it from the ready queue. The task will be
|
||||
// assigned to a worker once one becomes available.
|
||||
// (See design_docs/task_states.rst for the state transition diagram.)
|
||||
local_queues_.QueueTasks({assigned_task}, TaskState::READY);
|
||||
DispatchTasks(MakeTasksWithResources({assigned_task}));
|
||||
}
|
||||
}
|
||||
|
||||
void NodeManager::DumpDebugState() const {
|
||||
std::fstream fs;
|
||||
fs.open(initial_config_.session_dir + "/debug_state.txt",
|
||||
|
||||
@@ -7,6 +7,8 @@
|
||||
#include "ray/rpc/client_call.h"
|
||||
#include "ray/rpc/node_manager/node_manager_server.h"
|
||||
#include "ray/rpc/node_manager/node_manager_client.h"
|
||||
#include "ray/rpc/raylet/raylet_server.h"
|
||||
#include "ray/object_manager/object_manager.h"
|
||||
#include "ray/common/task/task.h"
|
||||
#include "ray/common/client_connection.h"
|
||||
#include "ray/common/task/task_common.h"
|
||||
@@ -64,7 +66,8 @@ struct NodeManagerConfig {
|
||||
std::string session_dir;
|
||||
};
|
||||
|
||||
class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
class NodeManager : public rpc::NodeManagerServiceHandler,
|
||||
public rpc::RayletServiceHandler {
|
||||
public:
|
||||
/// Create a node manager.
|
||||
///
|
||||
@@ -75,29 +78,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
std::shared_ptr<gcs::RedisGcsClient> gcs_client,
|
||||
std::shared_ptr<ObjectDirectoryInterface> object_directory_);
|
||||
|
||||
/// Process a new client connection.
|
||||
///
|
||||
/// \param client The client to process.
|
||||
/// \return Void.
|
||||
void ProcessNewClient(LocalClientConnection &client);
|
||||
|
||||
/// Process a message from a client. This method is responsible for
|
||||
/// explicitly listening for more messages from the client if the client is
|
||||
/// still alive.
|
||||
///
|
||||
/// \param client The client that sent the message.
|
||||
/// \param message_type The message type (e.g., a flatbuffer enum).
|
||||
/// \param message_data A pointer to the message data.
|
||||
/// \return Void.
|
||||
void ProcessClientMessage(const std::shared_ptr<LocalClientConnection> &client,
|
||||
int64_t message_type, const uint8_t *message_data);
|
||||
|
||||
/// Handle a new node manager connection.
|
||||
///
|
||||
/// \param node_manager_client The connection to the remote node manager.
|
||||
/// \return Void.
|
||||
void ProcessNewNodeManager(TcpClientConnection &node_manager_client);
|
||||
|
||||
/// Subscribe to the relevant GCS tables and set up handlers.
|
||||
///
|
||||
/// \return Status indicating whether this was done successfully or not.
|
||||
@@ -111,9 +91,83 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// Record metrics.
|
||||
void RecordMetrics() const;
|
||||
|
||||
public:
|
||||
/// Get the port of the node manager rpc server.
|
||||
int GetServerPort() const { return node_manager_server_.GetPort(); }
|
||||
|
||||
/// Implementation of node manager grpc service.
|
||||
|
||||
/// Handle a `ForwardTask` request.
|
||||
///
|
||||
/// \param request The request.
|
||||
/// \param reply The reply that will be sent to client.
|
||||
/// \param send_reply_callback Invoke this callback to send reply asynchronously.
|
||||
void HandleForwardTask(const rpc::ForwardTaskRequest &request,
|
||||
rpc::ForwardTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
/// Implementation of raylet grpc service handlers, please see definitions
|
||||
/// in `src/ray/protobuf/raylet.proto` for more details.
|
||||
void HandleRegisterClientRequest(const rpc::RegisterClientRequest &request,
|
||||
rpc::RegisterClientReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleSubmitTaskRequest(const rpc::SubmitTaskRequest &request,
|
||||
rpc::SubmitTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleDisconnectClientRequest(const rpc::DisconnectClientRequest &request,
|
||||
rpc::DisconnectClientReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleGetTaskRequest(const rpc::GetTaskRequest &request, rpc::GetTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleTaskDoneRequest(const rpc::TaskDoneRequest &request,
|
||||
rpc::TaskDoneReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleFetchOrReconstructRequest(
|
||||
const rpc::FetchOrReconstructRequest &request, rpc::FetchOrReconstructReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleNotifyUnblockedRequest(const rpc::NotifyUnblockedRequest &request,
|
||||
rpc::NotifyUnblockedReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleWaitRequest(const rpc::WaitRequest &request, rpc::WaitReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandlePushErrorRequest(const rpc::PushErrorRequest &request,
|
||||
rpc::PushErrorReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandlePushProfileEventsRequest(
|
||||
const rpc::PushProfileEventsRequest &request, rpc::PushProfileEventsReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleFreeObjectsInStoreRequest(
|
||||
const rpc::FreeObjectsInStoreRequest &request, rpc::FreeObjectsInStoreReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandlePrepareActorCheckpointRequest(
|
||||
const rpc::PrepareActorCheckpointRequest &request,
|
||||
rpc::PrepareActorCheckpointReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleNotifyActorResumedFromCheckpointRequest(
|
||||
const rpc::NotifyActorResumedFromCheckpointRequest &request,
|
||||
rpc::NotifyActorResumedFromCheckpointReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleSetResourceRequest(const rpc::SetResourceRequest &request,
|
||||
rpc::SetResourceReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
void HandleHeartbeatRequest(const rpc::HeartbeatRequest &request,
|
||||
rpc::HeartbeatReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
private:
|
||||
/// Methods for handling clients.
|
||||
|
||||
@@ -214,6 +268,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// \return Void.
|
||||
void SubmitTask(const Task &task, const Lineage &uncommitted_lineage,
|
||||
bool forwarded = false);
|
||||
/// Handle the case that a worker is available.
|
||||
///
|
||||
/// \param id Id of the worker.
|
||||
/// \return Void.
|
||||
void HandleWorkerAvailable(const WorkerID &worker_id);
|
||||
/// Assign a task. The task is assumed to not be queued in local_queues_.
|
||||
///
|
||||
/// \param task The task in question.
|
||||
@@ -322,7 +381,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// \param ray_get Whether the task is blocked in a `ray.get` call, as
|
||||
/// opposed to a `ray.wait` call.
|
||||
/// \return Void.
|
||||
void HandleTaskBlocked(const std::shared_ptr<LocalClientConnection> &client,
|
||||
void HandleTaskBlocked(const WorkerID &worker_id,
|
||||
const std::vector<ObjectID> &required_object_ids,
|
||||
const TaskID ¤t_task_id, bool ray_get);
|
||||
|
||||
@@ -335,8 +394,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// \param client The client that is executing the unblocked task.
|
||||
/// \param current_task_id The task that is unblocked.
|
||||
/// \return Void.
|
||||
void HandleTaskUnblocked(const std::shared_ptr<LocalClientConnection> &client,
|
||||
const TaskID ¤t_task_id);
|
||||
void HandleTaskUnblocked(const WorkerID &worker_id, const TaskID ¤t_task_id);
|
||||
|
||||
/// Kill a worker.
|
||||
///
|
||||
@@ -410,11 +468,10 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// client.
|
||||
///
|
||||
/// \param client The client that sent the message.
|
||||
/// \param intentional_disconnect Wether the client was intentionally disconnected.
|
||||
/// \param intentional_disconnect Whether the client was intentionally disconnected.
|
||||
/// \return Void.
|
||||
void ProcessDisconnectClientMessage(
|
||||
const std::shared_ptr<LocalClientConnection> &client,
|
||||
bool intentional_disconnect = false);
|
||||
void ProcessDisconnectClientMessage(const WorkerID &worker_id,
|
||||
bool intentional_disconnect = false);
|
||||
|
||||
/// Process client message of SubmitTask
|
||||
///
|
||||
@@ -482,19 +539,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
void HandleDisconnectedActor(const ActorID &actor_id, bool was_local,
|
||||
bool intentional_disconnect);
|
||||
|
||||
/// Finish assigning a task to a worker.
|
||||
///
|
||||
/// \param task_id Id of the task.
|
||||
/// \param worker Worker which the task is assigned to.
|
||||
/// \param success Whether the task is successfully assigned to the worker.
|
||||
/// \return void.
|
||||
void FinishAssignTask(const TaskID &task_id, Worker &worker, bool success);
|
||||
|
||||
/// Handle a `ForwardTask` request.
|
||||
void HandleForwardTask(const rpc::ForwardTaskRequest &request,
|
||||
rpc::ForwardTaskReply *reply,
|
||||
rpc::SendReplyCallback send_reply_callback) override;
|
||||
|
||||
// GCS client ID for this node.
|
||||
ClientID client_id_;
|
||||
boost::asio::io_service &io_service_;
|
||||
@@ -551,9 +595,12 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
|
||||
/// The RPC server.
|
||||
rpc::GrpcServer node_manager_server_;
|
||||
|
||||
/// The RPC service.
|
||||
/// The node manager RPC service.
|
||||
rpc::NodeManagerGrpcService node_manager_service_;
|
||||
|
||||
/// The raylet RPC service.
|
||||
rpc::RayletGrpcService raylet_service_;
|
||||
|
||||
/// The `ClientCallManager` object that is shared by all `NodeManagerClient`s
|
||||
/// as well as all `WorkerTaskClient`s.
|
||||
rpc::ClientCallManager client_call_manager_;
|
||||
|
||||
@@ -7,34 +7,6 @@
|
||||
|
||||
#include "ray/common/status.h"
|
||||
|
||||
namespace {
|
||||
|
||||
const std::vector<std::string> GenerateEnumNames(const char *const *enum_names_ptr,
|
||||
int start_index, int end_index) {
|
||||
std::vector<std::string> enum_names;
|
||||
for (int i = 0; i < start_index; ++i) {
|
||||
enum_names.push_back("EmptyMessageType");
|
||||
}
|
||||
size_t i = 0;
|
||||
while (true) {
|
||||
const char *name = enum_names_ptr[i];
|
||||
if (name == nullptr) {
|
||||
break;
|
||||
}
|
||||
enum_names.push_back(name);
|
||||
i++;
|
||||
}
|
||||
RAY_CHECK(static_cast<size_t>(end_index) == enum_names.size() - 1)
|
||||
<< "Message Type mismatch!";
|
||||
return enum_names;
|
||||
}
|
||||
|
||||
static const std::vector<std::string> node_manager_message_enum =
|
||||
GenerateEnumNames(ray::protocol::EnumNamesMessageType(),
|
||||
static_cast<int>(ray::protocol::MessageType::MIN),
|
||||
static_cast<int>(ray::protocol::MessageType::MAX));
|
||||
} // namespace
|
||||
|
||||
namespace ray {
|
||||
|
||||
namespace raylet {
|
||||
@@ -51,10 +23,10 @@ Raylet::Raylet(boost::asio::io_service &main_service, const std::string &socket_
|
||||
node_manager_(main_service, node_manager_config, object_manager_, gcs_client_,
|
||||
object_directory_),
|
||||
socket_name_(socket_name),
|
||||
acceptor_(main_service, boost::asio::local::stream_protocol::endpoint(socket_name)),
|
||||
socket_(main_service) {
|
||||
// Start listening for clients.
|
||||
DoAccept();
|
||||
raylet_server_("Raylet", socket_name),
|
||||
raylet_service_(main_service, node_manager_) {
|
||||
raylet_server_.RegisterService(raylet_service_);
|
||||
raylet_server_.Run();
|
||||
|
||||
RAY_CHECK_OK(RegisterGcs(
|
||||
node_ip_address, socket_name_, object_manager_config.store_socket_name,
|
||||
@@ -108,31 +80,6 @@ ray::Status Raylet::RegisterGcs(const std::string &node_ip_address,
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
void Raylet::DoAccept() {
|
||||
acceptor_.async_accept(socket_, boost::bind(&Raylet::HandleAccept, this,
|
||||
boost::asio::placeholders::error));
|
||||
}
|
||||
|
||||
void Raylet::HandleAccept(const boost::system::error_code &error) {
|
||||
if (!error) {
|
||||
// TODO: typedef these handlers.
|
||||
ClientHandler<boost::asio::local::stream_protocol> client_handler =
|
||||
[this](LocalClientConnection &client) { node_manager_.ProcessNewClient(client); };
|
||||
MessageHandler<boost::asio::local::stream_protocol> message_handler =
|
||||
[this](std::shared_ptr<LocalClientConnection> client, int64_t message_type,
|
||||
const uint8_t *message) {
|
||||
node_manager_.ProcessClientMessage(client, message_type, message);
|
||||
};
|
||||
// Accept a new local client and dispatch it to the node manager.
|
||||
auto new_connection = LocalClientConnection::Create(
|
||||
client_handler, message_handler, std::move(socket_), "worker",
|
||||
node_manager_message_enum,
|
||||
static_cast<int64_t>(protocol::MessageType::DisconnectClient));
|
||||
}
|
||||
// We're ready to accept another client.
|
||||
DoAccept();
|
||||
}
|
||||
|
||||
} // namespace raylet
|
||||
|
||||
} // namespace ray
|
||||
|
||||
@@ -22,7 +22,7 @@ class NodeManager;
|
||||
|
||||
class Raylet {
|
||||
public:
|
||||
/// Create a node manager server and listen for new clients.
|
||||
/// Create a raylet server and listen for local clients.
|
||||
///
|
||||
/// \param main_service The event loop to run the server on.
|
||||
/// \param object_manager_service The asio io_service tied to the object manager.
|
||||
@@ -75,10 +75,12 @@ class Raylet {
|
||||
/// The name of the socket this raylet listens on.
|
||||
std::string socket_name_;
|
||||
|
||||
/// An acceptor for new clients.
|
||||
boost::asio::local::stream_protocol::acceptor acceptor_;
|
||||
/// The socket to listen on for new clients.
|
||||
boost::asio::local::stream_protocol::socket socket_;
|
||||
/// The gRPC server, listens for local raylet client connections through a unix domain
|
||||
/// socket.
|
||||
rpc::GrpcServer raylet_server_;
|
||||
|
||||
/// The gRPC service.
|
||||
rpc::RayletGrpcService raylet_service_;
|
||||
};
|
||||
|
||||
} // namespace raylet
|
||||
|
||||
@@ -1,396 +0,0 @@
|
||||
#include "raylet_client.h"
|
||||
|
||||
#include <inttypes.h>
|
||||
#include <netdb.h>
|
||||
#include <netinet/in.h>
|
||||
#include <stdarg.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <sys/ioctl.h>
|
||||
#include <sys/socket.h>
|
||||
#include <sys/types.h>
|
||||
#include <sys/un.h>
|
||||
|
||||
#include "ray/common/common_protocol.h"
|
||||
#include "ray/common/ray_config.h"
|
||||
#include "ray/common/task/task_spec.h"
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/util/logging.h"
|
||||
|
||||
using MessageType = ray::protocol::MessageType;
|
||||
|
||||
// TODO(rkn): The io methods below should be removed.
|
||||
int connect_ipc_sock(const std::string &socket_pathname) {
|
||||
struct sockaddr_un socket_address;
|
||||
int socket_fd;
|
||||
|
||||
socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
|
||||
if (socket_fd < 0) {
|
||||
RAY_LOG(ERROR) << "socket() failed for pathname " << socket_pathname;
|
||||
return -1;
|
||||
}
|
||||
|
||||
memset(&socket_address, 0, sizeof(socket_address));
|
||||
socket_address.sun_family = AF_UNIX;
|
||||
if (socket_pathname.length() + 1 > sizeof(socket_address.sun_path)) {
|
||||
RAY_LOG(ERROR) << "Socket pathname is too long.";
|
||||
close(socket_fd);
|
||||
return -1;
|
||||
}
|
||||
strncpy(socket_address.sun_path, socket_pathname.c_str(), socket_pathname.length() + 1);
|
||||
|
||||
if (connect(socket_fd, (struct sockaddr *)&socket_address, sizeof(socket_address)) !=
|
||||
0) {
|
||||
close(socket_fd);
|
||||
return -1;
|
||||
}
|
||||
return socket_fd;
|
||||
}
|
||||
|
||||
int read_bytes(int socket_fd, uint8_t *cursor, size_t length) {
|
||||
ssize_t nbytes = 0;
|
||||
// Termination condition: EOF or read 'length' bytes total.
|
||||
size_t bytesleft = length;
|
||||
size_t offset = 0;
|
||||
while (bytesleft > 0) {
|
||||
nbytes = read(socket_fd, cursor + offset, bytesleft);
|
||||
if (nbytes < 0) {
|
||||
if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
|
||||
continue;
|
||||
}
|
||||
return -1; // Errno will be set.
|
||||
} else if (0 == nbytes) {
|
||||
// Encountered early EOF.
|
||||
return -1;
|
||||
}
|
||||
RAY_CHECK(nbytes > 0);
|
||||
bytesleft -= nbytes;
|
||||
offset += nbytes;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int write_bytes(int socket_fd, uint8_t *cursor, size_t length) {
|
||||
ssize_t nbytes = 0;
|
||||
size_t bytesleft = length;
|
||||
size_t offset = 0;
|
||||
while (bytesleft > 0) {
|
||||
// While we haven't written the whole message, write to the file
|
||||
// descriptor, advance the cursor, and decrease the amount left to write.
|
||||
nbytes = write(socket_fd, cursor + offset, bytesleft);
|
||||
if (nbytes < 0) {
|
||||
if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
|
||||
continue;
|
||||
}
|
||||
return -1; // Errno will be set.
|
||||
} else if (0 == nbytes) {
|
||||
// Encountered early EOF.
|
||||
return -1;
|
||||
}
|
||||
RAY_CHECK(nbytes > 0);
|
||||
bytesleft -= nbytes;
|
||||
offset += nbytes;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
RayletConnection::RayletConnection(const std::string &raylet_socket, int num_retries,
|
||||
int64_t timeout) {
|
||||
// Pick the default values if the user did not specify.
|
||||
if (num_retries < 0) {
|
||||
num_retries = RayConfig::instance().num_connect_attempts();
|
||||
}
|
||||
if (timeout < 0) {
|
||||
timeout = RayConfig::instance().connect_timeout_milliseconds();
|
||||
}
|
||||
RAY_CHECK(!raylet_socket.empty());
|
||||
conn_ = -1;
|
||||
for (int num_attempts = 0; num_attempts < num_retries; ++num_attempts) {
|
||||
conn_ = connect_ipc_sock(raylet_socket);
|
||||
if (conn_ >= 0) break;
|
||||
if (num_attempts > 0) {
|
||||
RAY_LOG(ERROR) << "Retrying to connect to socket for pathname " << raylet_socket
|
||||
<< " (num_attempts = " << num_attempts
|
||||
<< ", num_retries = " << num_retries << ")";
|
||||
}
|
||||
// Sleep for timeout milliseconds.
|
||||
usleep(timeout * 1000);
|
||||
}
|
||||
// If we could not connect to the socket, exit.
|
||||
if (conn_ == -1) {
|
||||
RAY_LOG(FATAL) << "Could not connect to socket " << raylet_socket;
|
||||
}
|
||||
}
|
||||
|
||||
ray::Status RayletConnection::Disconnect() {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateDisconnectClient(fbb);
|
||||
fbb.Finish(message);
|
||||
auto status = WriteMessage(MessageType::IntentionalDisconnectClient, &fbb);
|
||||
// Don't be too strict for disconnection errors.
|
||||
// Just create logs and prevent it from crash.
|
||||
if (!status.ok()) {
|
||||
RAY_LOG(ERROR) << status.ToString()
|
||||
<< " [RayletClient] Failed to disconnect from raylet.";
|
||||
}
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletConnection::ReadMessage(MessageType type,
|
||||
std::unique_ptr<uint8_t[]> &message) {
|
||||
int64_t cookie;
|
||||
int64_t type_field;
|
||||
int64_t length;
|
||||
int closed = read_bytes(conn_, (uint8_t *)&cookie, sizeof(cookie));
|
||||
if (closed) goto disconnected;
|
||||
RAY_CHECK(cookie == RayConfig::instance().ray_cookie());
|
||||
closed = read_bytes(conn_, (uint8_t *)&type_field, sizeof(type_field));
|
||||
if (closed) goto disconnected;
|
||||
closed = read_bytes(conn_, (uint8_t *)&length, sizeof(length));
|
||||
if (closed) goto disconnected;
|
||||
message = std::unique_ptr<uint8_t[]>(new uint8_t[length]);
|
||||
closed = read_bytes(conn_, message.get(), length);
|
||||
if (closed) {
|
||||
// Handle the case in which the socket is closed.
|
||||
message.reset(nullptr);
|
||||
disconnected:
|
||||
message = nullptr;
|
||||
type_field = static_cast<int64_t>(MessageType::DisconnectClient);
|
||||
length = 0;
|
||||
}
|
||||
if (type_field == static_cast<int64_t>(MessageType::DisconnectClient)) {
|
||||
return ray::Status::IOError("[RayletClient] Raylet connection closed.");
|
||||
}
|
||||
if (type_field != static_cast<int64_t>(type)) {
|
||||
return ray::Status::TypeError(
|
||||
std::string("[RayletClient] Raylet connection corrupted. ") +
|
||||
"Expected message type: " + std::to_string(static_cast<int64_t>(type)) +
|
||||
"; got message type: " + std::to_string(type_field) +
|
||||
". Check logs or dmesg for previous errors.");
|
||||
}
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletConnection::WriteMessage(MessageType type,
|
||||
flatbuffers::FlatBufferBuilder *fbb) {
|
||||
std::unique_lock<std::mutex> guard(write_mutex_);
|
||||
int64_t cookie = RayConfig::instance().ray_cookie();
|
||||
int64_t length = fbb ? fbb->GetSize() : 0;
|
||||
uint8_t *bytes = fbb ? fbb->GetBufferPointer() : nullptr;
|
||||
int64_t type_field = static_cast<int64_t>(type);
|
||||
auto io_error = ray::Status::IOError("[RayletClient] Connection closed unexpectedly.");
|
||||
int closed;
|
||||
closed = write_bytes(conn_, (uint8_t *)&cookie, sizeof(cookie));
|
||||
if (closed) return io_error;
|
||||
closed = write_bytes(conn_, (uint8_t *)&type_field, sizeof(type_field));
|
||||
if (closed) return io_error;
|
||||
closed = write_bytes(conn_, (uint8_t *)&length, sizeof(length));
|
||||
if (closed) return io_error;
|
||||
closed = write_bytes(conn_, bytes, length * sizeof(char));
|
||||
if (closed) return io_error;
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletConnection::AtomicRequestReply(
|
||||
MessageType request_type, MessageType reply_type,
|
||||
std::unique_ptr<uint8_t[]> &reply_message, flatbuffers::FlatBufferBuilder *fbb) {
|
||||
std::unique_lock<std::mutex> guard(mutex_);
|
||||
auto status = WriteMessage(request_type, fbb);
|
||||
if (!status.ok()) return status;
|
||||
return ReadMessage(reply_type, reply_message);
|
||||
}
|
||||
|
||||
RayletClient::RayletClient(const std::string &raylet_socket, const ClientID &client_id,
|
||||
bool is_worker, const JobID &job_id, const Language &language,
|
||||
int port)
|
||||
: client_id_(client_id),
|
||||
is_worker_(is_worker),
|
||||
job_id_(job_id),
|
||||
language_(language),
|
||||
port_(port) {
|
||||
// For C++14, we could use std::make_unique
|
||||
conn_ = std::unique_ptr<RayletConnection>(new RayletConnection(raylet_socket, -1, -1));
|
||||
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateRegisterClientRequest(
|
||||
fbb, is_worker, to_flatbuf(fbb, client_id), getpid(), to_flatbuf(fbb, job_id),
|
||||
language, port);
|
||||
fbb.Finish(message);
|
||||
// Register the process ID with the raylet.
|
||||
// NOTE(swang): If raylet exits and we are registered as a worker, we will get killed.
|
||||
auto status = conn_->WriteMessage(MessageType::RegisterClientRequest, &fbb);
|
||||
RAY_CHECK_OK_PREPEND(status, "[RayletClient] Unable to register worker with raylet.");
|
||||
}
|
||||
|
||||
ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateSubmitTaskRequest(
|
||||
fbb, fbb.CreateString(task_spec.Serialize()));
|
||||
fbb.Finish(message);
|
||||
return conn_->WriteMessage(MessageType::SubmitTask, &fbb);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::GetTask(std::unique_ptr<ray::TaskSpecification> *task_spec) {
|
||||
std::unique_ptr<uint8_t[]> reply;
|
||||
// Receive a task from the raylet. This will block until the raylet
|
||||
// gives this client a task.
|
||||
auto status =
|
||||
conn_->AtomicRequestReply(MessageType::GetTask, MessageType::ExecuteTask, reply);
|
||||
if (!status.ok()) return status;
|
||||
// Parse the flatbuffer object.
|
||||
auto reply_message = flatbuffers::GetRoot<ray::protocol::GetTaskReply>(reply.get());
|
||||
// Set the resource IDs for this task.
|
||||
resource_ids_.clear();
|
||||
for (size_t i = 0; i < reply_message->fractional_resource_ids()->size(); ++i) {
|
||||
auto const &fractional_resource_ids =
|
||||
reply_message->fractional_resource_ids()->Get(i);
|
||||
auto &acquired_resources =
|
||||
resource_ids_[string_from_flatbuf(*fractional_resource_ids->resource_name())];
|
||||
|
||||
size_t num_resource_ids = fractional_resource_ids->resource_ids()->size();
|
||||
size_t num_resource_fractions = fractional_resource_ids->resource_fractions()->size();
|
||||
RAY_CHECK(num_resource_ids == num_resource_fractions);
|
||||
RAY_CHECK(num_resource_ids > 0);
|
||||
for (size_t j = 0; j < num_resource_ids; ++j) {
|
||||
int64_t resource_id = fractional_resource_ids->resource_ids()->Get(j);
|
||||
double resource_fraction = fractional_resource_ids->resource_fractions()->Get(j);
|
||||
if (num_resource_ids > 1) {
|
||||
int64_t whole_fraction = resource_fraction;
|
||||
RAY_CHECK(whole_fraction == resource_fraction);
|
||||
}
|
||||
acquired_resources.push_back(std::make_pair(resource_id, resource_fraction));
|
||||
}
|
||||
}
|
||||
|
||||
// Return the copy of the task spec and pass ownership to the caller.
|
||||
task_spec->reset(
|
||||
new ray::TaskSpecification(string_from_flatbuf(*reply_message->task_spec())));
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::TaskDone() {
|
||||
return conn_->WriteMessage(MessageType::TaskDone);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::FetchOrReconstruct(const std::vector<ObjectID> &object_ids,
|
||||
bool fetch_only,
|
||||
const TaskID ¤t_task_id) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto object_ids_message = to_flatbuf(fbb, object_ids);
|
||||
auto message = ray::protocol::CreateFetchOrReconstruct(
|
||||
fbb, object_ids_message, fetch_only, to_flatbuf(fbb, current_task_id));
|
||||
fbb.Finish(message);
|
||||
auto status = conn_->WriteMessage(MessageType::FetchOrReconstruct, &fbb);
|
||||
return status;
|
||||
}
|
||||
|
||||
ray::Status RayletClient::NotifyUnblocked(const TaskID ¤t_task_id) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message =
|
||||
ray::protocol::CreateNotifyUnblocked(fbb, to_flatbuf(fbb, current_task_id));
|
||||
fbb.Finish(message);
|
||||
return conn_->WriteMessage(MessageType::NotifyUnblocked, &fbb);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::Wait(const std::vector<ObjectID> &object_ids, int num_returns,
|
||||
int64_t timeout_milliseconds, bool wait_local,
|
||||
const TaskID ¤t_task_id, WaitResultPair *result) {
|
||||
// Write request.
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateWaitRequest(
|
||||
fbb, to_flatbuf(fbb, object_ids), num_returns, timeout_milliseconds, wait_local,
|
||||
to_flatbuf(fbb, current_task_id));
|
||||
fbb.Finish(message);
|
||||
std::unique_ptr<uint8_t[]> reply;
|
||||
auto status = conn_->AtomicRequestReply(MessageType::WaitRequest,
|
||||
MessageType::WaitReply, reply, &fbb);
|
||||
if (!status.ok()) return status;
|
||||
// Parse the flatbuffer object.
|
||||
auto reply_message = flatbuffers::GetRoot<ray::protocol::WaitReply>(reply.get());
|
||||
auto found = reply_message->found();
|
||||
for (uint i = 0; i < found->size(); i++) {
|
||||
ObjectID object_id = ObjectID::FromBinary(found->Get(i)->str());
|
||||
result->first.push_back(object_id);
|
||||
}
|
||||
auto remaining = reply_message->remaining();
|
||||
for (uint i = 0; i < remaining->size(); i++) {
|
||||
ObjectID object_id = ObjectID::FromBinary(remaining->Get(i)->str());
|
||||
result->second.push_back(object_id);
|
||||
}
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PushError(const ray::JobID &job_id, const std::string &type,
|
||||
const std::string &error_message, double timestamp) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreatePushErrorRequest(
|
||||
fbb, to_flatbuf(fbb, job_id), fbb.CreateString(type),
|
||||
fbb.CreateString(error_message), timestamp);
|
||||
fbb.Finish(message);
|
||||
|
||||
return conn_->WriteMessage(MessageType::PushErrorRequest, &fbb);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PushProfileEvents(const ProfileTableData &profile_events) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = fbb.CreateString(profile_events.SerializeAsString());
|
||||
fbb.Finish(message);
|
||||
|
||||
auto status = conn_->WriteMessage(MessageType::PushProfileEventsRequest, &fbb);
|
||||
// Don't be too strict for profile errors. Just create logs and prevent it from crash.
|
||||
if (!status.ok()) {
|
||||
RAY_LOG(ERROR) << status.ToString()
|
||||
<< " [RayletClient] Failed to push profile events.";
|
||||
}
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::FreeObjects(const std::vector<ray::ObjectID> &object_ids,
|
||||
bool local_only, bool delete_creating_tasks) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateFreeObjectsRequest(
|
||||
fbb, local_only, delete_creating_tasks, to_flatbuf(fbb, object_ids));
|
||||
fbb.Finish(message);
|
||||
|
||||
auto status = conn_->WriteMessage(MessageType::FreeObjectsInObjectStoreRequest, &fbb);
|
||||
return status;
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PrepareActorCheckpoint(const ActorID &actor_id,
|
||||
ActorCheckpointID &checkpoint_id) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message =
|
||||
ray::protocol::CreatePrepareActorCheckpointRequest(fbb, to_flatbuf(fbb, actor_id));
|
||||
fbb.Finish(message);
|
||||
|
||||
std::unique_ptr<uint8_t[]> reply;
|
||||
auto status =
|
||||
conn_->AtomicRequestReply(MessageType::PrepareActorCheckpointRequest,
|
||||
MessageType::PrepareActorCheckpointReply, reply, &fbb);
|
||||
if (!status.ok()) return status;
|
||||
auto reply_message =
|
||||
flatbuffers::GetRoot<ray::protocol::PrepareActorCheckpointReply>(reply.get());
|
||||
checkpoint_id = ActorCheckpointID::FromBinary(reply_message->checkpoint_id()->str());
|
||||
return ray::Status::OK();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::NotifyActorResumedFromCheckpoint(
|
||||
const ActorID &actor_id, const ActorCheckpointID &checkpoint_id) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateNotifyActorResumedFromCheckpoint(
|
||||
fbb, to_flatbuf(fbb, actor_id), to_flatbuf(fbb, checkpoint_id));
|
||||
fbb.Finish(message);
|
||||
|
||||
return conn_->WriteMessage(MessageType::NotifyActorResumedFromCheckpoint, &fbb);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::SetResource(const std::string &resource_name,
|
||||
const double capacity,
|
||||
const ray::ClientID &client_Id) {
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto message = ray::protocol::CreateSetResourceRequest(
|
||||
fbb, fbb.CreateString(resource_name), capacity, to_flatbuf(fbb, client_Id));
|
||||
fbb.Finish(message);
|
||||
return conn_->WriteMessage(MessageType::SetResourceRequest, &fbb);
|
||||
}
|
||||
@@ -5,7 +5,6 @@
|
||||
|
||||
#include <boost/asio.hpp>
|
||||
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/raylet/reconstruction_policy.h"
|
||||
|
||||
#include "ray/object_manager/object_directory.h"
|
||||
@@ -400,6 +399,7 @@ TEST_F(ReconstructionPolicyTest, TestSimultaneousReconstructionSuppressed) {
|
||||
auto task_reconstruction_data = std::make_shared<TaskReconstructionData>();
|
||||
task_reconstruction_data->set_node_manager_id(ClientID::FromRandom().Binary());
|
||||
task_reconstruction_data->set_num_reconstructions(0);
|
||||
|
||||
RAY_CHECK_OK(
|
||||
mock_gcs_.AppendAt(JobID::Nil(), task_id, task_reconstruction_data, nullptr,
|
||||
/*failure_callback=*/
|
||||
|
||||
@@ -340,7 +340,7 @@ void TaskDependencyManager::AcquireTaskLease(const TaskID &task_id) {
|
||||
// Check that we were able to renew the task lease before the previous one
|
||||
// expired.
|
||||
if (now_ms > it->second.expires_at) {
|
||||
RAY_LOG(WARNING) << "Task lease to renew has already expired by "
|
||||
RAY_LOG(WARNING) << "Task " << task_id << " lease to renew has already expired by "
|
||||
<< (it->second.expires_at - now_ms) << "ms";
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
// clang-format off
|
||||
#include "ray/common/id.h"
|
||||
#include "ray/common/ray_config.h"
|
||||
#include "ray/common/task/task.h"
|
||||
#include "ray/object_manager/object_manager.h"
|
||||
#include "ray/raylet/reconstruction_policy.h"
|
||||
|
||||
+31
-26
@@ -2,7 +2,6 @@
|
||||
|
||||
#include <boost/bind.hpp>
|
||||
|
||||
#include "ray/raylet/format/node_manager_generated.h"
|
||||
#include "ray/raylet/raylet.h"
|
||||
|
||||
namespace ray {
|
||||
@@ -11,15 +10,14 @@ namespace raylet {
|
||||
|
||||
/// A constructor responsible for initializing the state of a worker.
|
||||
Worker::Worker(const WorkerID &worker_id, pid_t pid, const Language &language, int port,
|
||||
std::shared_ptr<LocalClientConnection> connection,
|
||||
rpc::ClientCallManager &client_call_manager)
|
||||
: worker_id_(worker_id),
|
||||
pid_(pid),
|
||||
language_(language),
|
||||
port_(port),
|
||||
connection_(connection),
|
||||
dead_(false),
|
||||
language_(language),
|
||||
blocked_(false),
|
||||
num_missed_heartbeats_(0),
|
||||
is_being_killed_(false),
|
||||
client_call_manager_(client_call_manager) {
|
||||
if (port_ > 0) {
|
||||
rpc_client_ = std::unique_ptr<rpc::WorkerTaskClient>(
|
||||
@@ -27,9 +25,9 @@ Worker::Worker(const WorkerID &worker_id, pid_t pid, const Language &language, i
|
||||
}
|
||||
}
|
||||
|
||||
void Worker::MarkDead() { dead_ = true; }
|
||||
void Worker::MarkAsBeingKilled() { is_being_killed_ = true; }
|
||||
|
||||
bool Worker::IsDead() const { return dead_; }
|
||||
bool Worker::IsBeingKilled() const { return is_being_killed_; }
|
||||
|
||||
void Worker::MarkBlocked() { blocked_ = true; }
|
||||
|
||||
@@ -43,6 +41,8 @@ pid_t Worker::Pid() const { return pid_; }
|
||||
|
||||
Language Worker::GetLanguage() const { return language_; }
|
||||
|
||||
const WorkerID &Worker::GetWorkerId() const { return worker_id_; }
|
||||
|
||||
int Worker::Port() const { return port_; }
|
||||
|
||||
void Worker::AssignTaskId(const TaskID &task_id) { assigned_task_id_ = task_id; }
|
||||
@@ -76,10 +76,6 @@ void Worker::AssignActorId(const ActorID &actor_id) {
|
||||
|
||||
const ActorID &Worker::GetActorId() const { return actor_id_; }
|
||||
|
||||
const std::shared_ptr<LocalClientConnection> Worker::Connection() const {
|
||||
return connection_;
|
||||
}
|
||||
|
||||
const ResourceIdSet &Worker::GetLifetimeResourceIds() const {
|
||||
return lifetime_resource_ids_;
|
||||
}
|
||||
@@ -113,10 +109,16 @@ void Worker::AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) {
|
||||
task_resource_ids_.Release(cpu_resources);
|
||||
}
|
||||
|
||||
void Worker::SetGetTaskReplyAndCallback(
|
||||
rpc::GetTaskReply *reply, const rpc::SendReplyCallback &&send_reply_callback) {
|
||||
RAY_CHECK(reply_ == nullptr && send_reply_callback_ == nullptr);
|
||||
reply_ = reply;
|
||||
send_reply_callback_ = std::move(send_reply_callback);
|
||||
}
|
||||
|
||||
bool Worker::UsePush() const { return rpc_client_ != nullptr; }
|
||||
|
||||
void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
|
||||
const std::function<void(Status)> finish_assign_callback) {
|
||||
void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set) {
|
||||
const TaskSpecification &spec = task.GetTaskSpecification();
|
||||
if (rpc_client_ != nullptr) {
|
||||
// Use push mode.
|
||||
@@ -126,29 +128,32 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
|
||||
task.GetTaskSpecification().GetMessage());
|
||||
request.mutable_task()->mutable_task_execution_spec()->CopyFrom(
|
||||
task.GetTaskExecutionSpec().GetMessage());
|
||||
request.set_resource_ids(resource_id_set.Serialize());
|
||||
|
||||
for (const auto &e : resource_id_set.ToProtobuf()) {
|
||||
auto resource = request.add_resource_ids();
|
||||
*resource = e;
|
||||
}
|
||||
auto status = rpc_client_->AssignTask(
|
||||
request, [](Status status, const rpc::AssignTaskReply &reply) {
|
||||
// Worker has finished this task. There's nothing to do here
|
||||
// and assigning new task will be done when raylet receives
|
||||
// `TaskDone` message.
|
||||
});
|
||||
finish_assign_callback(status);
|
||||
} else {
|
||||
// Use pull mode. This corresponds to existing python/java workers that haven't been
|
||||
// migrated to core worker architecture.
|
||||
flatbuffers::FlatBufferBuilder fbb;
|
||||
auto resource_id_set_flatbuf = resource_id_set.ToFlatbuf(fbb);
|
||||
|
||||
auto message =
|
||||
protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()),
|
||||
fbb.CreateVector(resource_id_set_flatbuf));
|
||||
fbb.Finish(message);
|
||||
Connection()->WriteMessageAsync(
|
||||
static_cast<int64_t>(protocol::MessageType::ExecuteTask), fbb.GetSize(),
|
||||
fbb.GetBufferPointer(), finish_assign_callback);
|
||||
RAY_CHECK(reply_ != nullptr && send_reply_callback_ != nullptr);
|
||||
reply_->set_task_spec(task.GetTaskSpecification().Serialize());
|
||||
for (const auto &e : resource_id_set.ToProtobuf()) {
|
||||
auto resource = reply_->add_fractional_resource_ids();
|
||||
*resource = e;
|
||||
}
|
||||
send_reply_callback_(Status::OK(), nullptr, nullptr);
|
||||
reply_ = nullptr;
|
||||
send_reply_callback_ = nullptr;
|
||||
}
|
||||
// The status will be cleared when the worker dies.
|
||||
AssignTaskId(spec.TaskId());
|
||||
AssignJobId(spec.JobId());
|
||||
}
|
||||
|
||||
} // namespace raylet
|
||||
|
||||
+29
-15
@@ -3,12 +3,15 @@
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "ray/common/client_connection.h"
|
||||
#include "ray/common/id.h"
|
||||
#include "ray/common/task/scheduling_resources.h"
|
||||
#include "ray/common/task/task.h"
|
||||
#include "ray/common/task/task_common.h"
|
||||
#include "ray/protobuf/common.pb.h"
|
||||
#include "ray/rpc/worker/worker_client.h"
|
||||
#include "src/ray/protobuf/gcs.pb.h"
|
||||
#include "src/ray/protobuf/raylet.pb.h"
|
||||
#include "src/ray/rpc/server_call.h"
|
||||
|
||||
namespace ray {
|
||||
|
||||
@@ -21,12 +24,11 @@ class Worker {
|
||||
public:
|
||||
/// A constructor that initializes a worker object.
|
||||
Worker(const WorkerID &worker_id, pid_t pid, const Language &language, int port,
|
||||
std::shared_ptr<LocalClientConnection> connection,
|
||||
rpc::ClientCallManager &client_call_manager);
|
||||
/// A destructor responsible for freeing all worker state.
|
||||
~Worker() {}
|
||||
void MarkDead();
|
||||
bool IsDead() const;
|
||||
void MarkAsBeingKilled();
|
||||
bool IsBeingKilled() const;
|
||||
void MarkBlocked();
|
||||
void MarkUnblocked();
|
||||
bool IsBlocked() const;
|
||||
@@ -35,6 +37,7 @@ class Worker {
|
||||
/// Return the worker's PID.
|
||||
pid_t Pid() const;
|
||||
Language GetLanguage() const;
|
||||
const WorkerID &GetWorkerId() const;
|
||||
int Port() const;
|
||||
void AssignTaskId(const TaskID &task_id);
|
||||
const TaskID &GetAssignedTaskId() const;
|
||||
@@ -45,8 +48,6 @@ class Worker {
|
||||
const JobID &GetAssignedJobId() const;
|
||||
void AssignActorId(const ActorID &actor_id);
|
||||
const ActorID &GetActorId() const;
|
||||
/// Return the worker's connection.
|
||||
const std::shared_ptr<LocalClientConnection> Connection() const;
|
||||
|
||||
const ResourceIdSet &GetLifetimeResourceIds() const;
|
||||
void SetLifetimeResourceIds(ResourceIdSet &resource_ids);
|
||||
@@ -58,30 +59,34 @@ class Worker {
|
||||
ResourceIdSet ReleaseTaskCpuResources();
|
||||
void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources);
|
||||
|
||||
int TickHeartbeatTimer() { return ++num_missed_heartbeats_; }
|
||||
void ClearHeartbeat() { num_missed_heartbeats_ = 0; }
|
||||
|
||||
/// When receiving a `GetTask` request from worker, this function should be called to
|
||||
/// pass the reply and callback of the request to the worker. Later, when we actually
|
||||
/// assign a task to the worker, the reply will be filled and the callback will be
|
||||
/// called.
|
||||
void SetGetTaskReplyAndCallback(rpc::GetTaskReply *reply,
|
||||
const rpc::SendReplyCallback &&send_reply_callback);
|
||||
|
||||
bool UsePush() const;
|
||||
void AssignTask(const Task &task, const ResourceIdSet &resource_id_set,
|
||||
const std::function<void(Status)> finish_assign_callback);
|
||||
void AssignTask(const Task &task, const ResourceIdSet &resource_id_set);
|
||||
|
||||
private:
|
||||
/// The worker's ID.
|
||||
WorkerID worker_id_;
|
||||
/// The worker's PID.
|
||||
pid_t pid_;
|
||||
/// The worker port.
|
||||
int port_;
|
||||
/// The language type of this worker.
|
||||
Language language_;
|
||||
/// Port that this worker listens on.
|
||||
/// If port <= 0, this indicates that the worker will not listen to a port.
|
||||
int port_;
|
||||
/// Connection state of a worker.
|
||||
std::shared_ptr<LocalClientConnection> connection_;
|
||||
/// The worker's currently assigned task.
|
||||
TaskID assigned_task_id_;
|
||||
/// Job ID for the worker's current assigned task.
|
||||
JobID assigned_job_id_;
|
||||
/// The worker's actor ID. If this is nil, then the worker is not an actor.
|
||||
ActorID actor_id_;
|
||||
/// Whether the worker is dead.
|
||||
bool dead_;
|
||||
/// Whether the worker is blocked. Workers become blocked in a `ray.get`, if
|
||||
/// they require a data dependency while executing a task.
|
||||
bool blocked_;
|
||||
@@ -92,11 +97,20 @@ class Worker {
|
||||
// of a task.
|
||||
ResourceIdSet task_resource_ids_;
|
||||
std::unordered_set<TaskID> blocked_task_ids_;
|
||||
/// How many heartbeats have been missed for this worker.
|
||||
int num_missed_heartbeats_;
|
||||
/// Indicates we have sent kill signal to the worker if it's true. We cannot treat the
|
||||
/// worker process as really dead until we lost the heartbeats from the worker.
|
||||
bool is_being_killed_;
|
||||
/// The `ClientCallManager` object that is shared by `WorkerTaskClient` from all
|
||||
/// workers.
|
||||
rpc::ClientCallManager &client_call_manager_;
|
||||
/// The rpc client to send tasks to this worker.
|
||||
std::unique_ptr<rpc::WorkerTaskClient> rpc_client_;
|
||||
/// Reply of the `GetTask` request.
|
||||
rpc::GetTaskReply *reply_ = nullptr;
|
||||
/// Callback of the `GetTask` request.
|
||||
rpc::SendReplyCallback send_reply_callback_ = nullptr;
|
||||
};
|
||||
|
||||
} // namespace raylet
|
||||
|
||||
@@ -12,29 +12,6 @@
|
||||
#include "ray/util/logging.h"
|
||||
#include "ray/util/util.h"
|
||||
|
||||
namespace {
|
||||
|
||||
// A helper function to get a worker from a list.
|
||||
std::shared_ptr<ray::raylet::Worker> GetWorker(
|
||||
const std::unordered_set<std::shared_ptr<ray::raylet::Worker>> &worker_pool,
|
||||
const std::shared_ptr<ray::LocalClientConnection> &connection) {
|
||||
for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) {
|
||||
if ((*it)->Connection() == connection) {
|
||||
return (*it);
|
||||
}
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// A helper function to remove a worker from a list. Returns true if the worker
|
||||
// was found and removed.
|
||||
bool RemoveWorker(std::unordered_set<std::shared_ptr<ray::raylet::Worker>> &worker_pool,
|
||||
const std::shared_ptr<ray::raylet::Worker> &worker) {
|
||||
return worker_pool.erase(worker) > 0;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace ray {
|
||||
|
||||
namespace raylet {
|
||||
@@ -73,8 +50,8 @@ WorkerPool::~WorkerPool() {
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
// Kill all registered workers. NOTE(swang): This assumes that the registered
|
||||
// workers were started by the pool.
|
||||
for (const auto &worker : entry.second.registered_workers) {
|
||||
pids_to_kill.insert(worker->Pid());
|
||||
for (const auto &worker_pair : entry.second.registered_workers) {
|
||||
pids_to_kill.insert(worker_pair.second->Pid());
|
||||
}
|
||||
// Kill all the workers that have been started but not registered.
|
||||
for (const auto &starting_worker : entry.second.starting_worker_processes) {
|
||||
@@ -189,12 +166,13 @@ pid_t WorkerPool::StartProcess(const std::vector<std::string> &worker_command_ar
|
||||
return 0;
|
||||
}
|
||||
|
||||
void WorkerPool::RegisterWorker(const std::shared_ptr<Worker> &worker) {
|
||||
void WorkerPool::RegisterWorker(const WorkerID &worker_id,
|
||||
const std::shared_ptr<Worker> &worker) {
|
||||
const auto pid = worker->Pid();
|
||||
const auto port = worker->Port();
|
||||
RAY_LOG(DEBUG) << "Registering worker with pid " << pid << ", port: " << port;
|
||||
auto &state = GetStateForLanguage(worker->GetLanguage());
|
||||
state.registered_workers.insert(std::move(worker));
|
||||
state.registered_workers.emplace(worker_id, std::move(worker));
|
||||
|
||||
auto it = state.starting_worker_processes.find(pid);
|
||||
if (it == state.starting_worker_processes.end()) {
|
||||
@@ -207,29 +185,30 @@ void WorkerPool::RegisterWorker(const std::shared_ptr<Worker> &worker) {
|
||||
}
|
||||
}
|
||||
|
||||
void WorkerPool::RegisterDriver(const std::shared_ptr<Worker> &driver) {
|
||||
void WorkerPool::RegisterDriver(const WorkerID &driver_id,
|
||||
const std::shared_ptr<Worker> &driver) {
|
||||
RAY_CHECK(!driver->GetAssignedTaskId().IsNil());
|
||||
auto &state = GetStateForLanguage(driver->GetLanguage());
|
||||
state.registered_drivers.insert(std::move(driver));
|
||||
state.registered_drivers.emplace(driver_id, std::move(driver));
|
||||
}
|
||||
|
||||
std::shared_ptr<Worker> WorkerPool::GetRegisteredWorker(
|
||||
const std::shared_ptr<LocalClientConnection> &connection) const {
|
||||
std::shared_ptr<Worker> WorkerPool::GetRegisteredWorker(const WorkerID &worker_id) const {
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
auto worker = GetWorker(entry.second.registered_workers, connection);
|
||||
if (worker != nullptr) {
|
||||
return worker;
|
||||
auto ®istered_workers = entry.second.registered_workers;
|
||||
auto it = registered_workers.find(worker_id);
|
||||
if (it != registered_workers.end()) {
|
||||
return it->second;
|
||||
}
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::shared_ptr<Worker> WorkerPool::GetRegisteredDriver(
|
||||
const std::shared_ptr<LocalClientConnection> &connection) const {
|
||||
std::shared_ptr<Worker> WorkerPool::GetRegisteredDriver(const WorkerID &worker_id) const {
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
auto driver = GetWorker(entry.second.registered_drivers, connection);
|
||||
if (driver != nullptr) {
|
||||
return driver;
|
||||
auto ®istered_drivers = entry.second.registered_drivers;
|
||||
auto it = registered_drivers.find(worker_id);
|
||||
if (it != registered_drivers.end()) {
|
||||
return it->second;
|
||||
}
|
||||
}
|
||||
return nullptr;
|
||||
@@ -313,18 +292,20 @@ std::shared_ptr<Worker> WorkerPool::PopWorker(const TaskSpecification &task_spec
|
||||
|
||||
bool WorkerPool::DisconnectWorker(const std::shared_ptr<Worker> &worker) {
|
||||
auto &state = GetStateForLanguage(worker->GetLanguage());
|
||||
RAY_CHECK(RemoveWorker(state.registered_workers, worker));
|
||||
RAY_CHECK(state.registered_workers.erase(worker->GetWorkerId()));
|
||||
|
||||
stats::CurrentWorker().Record(
|
||||
0, {{stats::LanguageKey, Language_Name(worker->GetLanguage())},
|
||||
{stats::WorkerPidKey, std::to_string(worker->Pid())}});
|
||||
|
||||
return RemoveWorker(state.idle, worker);
|
||||
// Indicates that we disconnect a idle worker successfully.
|
||||
return (state.idle.erase(worker) > 0);
|
||||
}
|
||||
|
||||
void WorkerPool::DisconnectDriver(const std::shared_ptr<Worker> &driver) {
|
||||
auto &state = GetStateForLanguage(driver->GetLanguage());
|
||||
RAY_CHECK(RemoveWorker(state.registered_drivers, driver));
|
||||
RAY_CHECK(state.registered_drivers.erase(driver->GetWorkerId()));
|
||||
|
||||
stats::CurrentDriver().Record(
|
||||
0, {{stats::LanguageKey, Language_Name(driver->GetLanguage())},
|
||||
{stats::WorkerPidKey, std::to_string(driver->Pid())}});
|
||||
@@ -341,7 +322,8 @@ std::vector<std::shared_ptr<Worker>> WorkerPool::GetWorkersRunningTasksForJob(
|
||||
std::vector<std::shared_ptr<Worker>> workers;
|
||||
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
for (const auto &worker : entry.second.registered_workers) {
|
||||
for (const auto &worker_pair : entry.second.registered_workers) {
|
||||
auto &worker = worker_pair.second;
|
||||
if (worker->GetAssignedJobId() == job_id) {
|
||||
workers.push_back(worker);
|
||||
}
|
||||
@@ -396,14 +378,16 @@ std::string WorkerPool::DebugString() const {
|
||||
void WorkerPool::RecordMetrics() const {
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
// Record worker.
|
||||
for (auto worker : entry.second.registered_workers) {
|
||||
for (auto worker_pair : entry.second.registered_workers) {
|
||||
auto &worker = worker_pair.second;
|
||||
stats::CurrentWorker().Record(
|
||||
worker->Pid(), {{stats::LanguageKey, Language_Name(worker->GetLanguage())},
|
||||
{stats::WorkerPidKey, std::to_string(worker->Pid())}});
|
||||
}
|
||||
|
||||
// Record driver.
|
||||
for (auto driver : entry.second.registered_drivers) {
|
||||
for (auto driver_pair : entry.second.registered_drivers) {
|
||||
auto &driver = driver_pair.second;
|
||||
stats::CurrentDriver().Record(
|
||||
driver->Pid(), {{stats::LanguageKey, Language_Name(driver->GetLanguage())},
|
||||
{stats::WorkerPidKey, std::to_string(driver->Pid())}});
|
||||
@@ -411,6 +395,27 @@ void WorkerPool::RecordMetrics() const {
|
||||
}
|
||||
}
|
||||
|
||||
void WorkerPool::TickHeartbeatTimer(int max_missed_heartbeats,
|
||||
std::vector<std::shared_ptr<Worker>> *dead_workers) {
|
||||
for (const auto &entry : states_by_lang_) {
|
||||
// Worker heartbeat.
|
||||
for (const auto &worker_pair : entry.second.registered_workers) {
|
||||
auto &worker = worker_pair.second;
|
||||
if (worker->TickHeartbeatTimer() >= max_missed_heartbeats) {
|
||||
dead_workers->emplace_back(worker);
|
||||
}
|
||||
}
|
||||
|
||||
// Driver heartbeat.
|
||||
for (const auto &driver_pair : entry.second.registered_drivers) {
|
||||
auto &driver = driver_pair.second;
|
||||
if (driver->TickHeartbeatTimer() >= max_missed_heartbeats) {
|
||||
dead_workers->emplace_back(driver);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace raylet
|
||||
|
||||
} // namespace ray
|
||||
|
||||
@@ -51,28 +51,27 @@ class WorkerPool {
|
||||
/// pool after it becomes idle (e.g., requests a work assignment).
|
||||
///
|
||||
/// \param The Worker to be registered.
|
||||
void RegisterWorker(const std::shared_ptr<Worker> &worker);
|
||||
void RegisterWorker(const WorkerID &worker_id, const std::shared_ptr<Worker> &worker);
|
||||
|
||||
/// Register a new driver.
|
||||
/// Driver is a treated as a special worker, so use WorkerID as key here.
|
||||
///
|
||||
/// \param The driver to be registered.
|
||||
void RegisterDriver(const std::shared_ptr<Worker> &worker);
|
||||
void RegisterDriver(const WorkerID &driver_id, const std::shared_ptr<Worker> &worker);
|
||||
|
||||
/// Get the client connection's registered worker.
|
||||
///
|
||||
/// \param The client connection owned by a registered worker.
|
||||
/// \return The Worker that owns the given client connection. Returns nullptr
|
||||
/// if the client has not registered a worker yet.
|
||||
std::shared_ptr<Worker> GetRegisteredWorker(
|
||||
const std::shared_ptr<LocalClientConnection> &connection) const;
|
||||
std::shared_ptr<Worker> GetRegisteredWorker(const WorkerID &worker_id) const;
|
||||
|
||||
/// Get the client connection's registered driver.
|
||||
///
|
||||
/// \param The client connection owned by a registered driver.
|
||||
/// \return The Worker that owns the given client connection. Returns nullptr
|
||||
/// if the client has not registered a driver.
|
||||
std::shared_ptr<Worker> GetRegisteredDriver(
|
||||
const std::shared_ptr<LocalClientConnection> &connection) const;
|
||||
std::shared_ptr<Worker> GetRegisteredDriver(const WorkerID &driver_id) const;
|
||||
|
||||
/// Disconnect a registered worker.
|
||||
///
|
||||
@@ -129,6 +128,16 @@ class WorkerPool {
|
||||
/// Record metrics.
|
||||
void RecordMetrics() const;
|
||||
|
||||
/// Tick the heartbeat timer and get the workers that have timed out.
|
||||
/// A worker which has missed `max_missed_heartbeats` times would be treated as a
|
||||
/// dead process or the network to it has been down.
|
||||
///
|
||||
/// \param[in] max_missed_heartbeats The maximum number of heartbeats that can be
|
||||
/// missed before a worker times out.
|
||||
/// \param[out] dead_workers Workers whose processes have been dead.
|
||||
void TickHeartbeatTimer(int max_missed_heartbeats,
|
||||
std::vector<std::shared_ptr<Worker>> *dead_workers);
|
||||
|
||||
protected:
|
||||
/// Asynchronously start a new worker process. Once the worker process has
|
||||
/// registered with an external server, the process should create and
|
||||
@@ -166,9 +175,9 @@ class WorkerPool {
|
||||
std::unordered_map<ActorID, std::shared_ptr<Worker>> idle_actor;
|
||||
/// All workers that have registered and are still connected, including both
|
||||
/// idle and executing.
|
||||
std::unordered_set<std::shared_ptr<Worker>> registered_workers;
|
||||
std::unordered_map<WorkerID, std::shared_ptr<Worker>> registered_workers;
|
||||
/// All drivers that have registered and are still connected.
|
||||
std::unordered_set<std::shared_ptr<Worker>> registered_drivers;
|
||||
std::unordered_map<WorkerID, std::shared_ptr<Worker>> registered_drivers;
|
||||
/// A map from the pids of starting worker processes
|
||||
/// to the number of their unregistered workers.
|
||||
std::unordered_map<pid_t, int> starting_worker_processes;
|
||||
|
||||
@@ -71,19 +71,9 @@ class WorkerPoolTest : public ::testing::Test {
|
||||
|
||||
std::shared_ptr<Worker> CreateWorker(pid_t pid,
|
||||
const Language &language = Language::PYTHON) {
|
||||
std::function<void(LocalClientConnection &)> client_handler =
|
||||
[this](LocalClientConnection &client) { HandleNewClient(client); };
|
||||
std::function<void(std::shared_ptr<LocalClientConnection>, int64_t, const uint8_t *)>
|
||||
message_handler = [this](std::shared_ptr<LocalClientConnection> client,
|
||||
int64_t message_type, const uint8_t *message) {
|
||||
HandleMessage(client, message_type, message);
|
||||
};
|
||||
boost::asio::local::stream_protocol::socket socket(io_service_);
|
||||
auto client =
|
||||
LocalClientConnection::Create(client_handler, message_handler, std::move(socket),
|
||||
"worker", {}, error_message_type_);
|
||||
return std::shared_ptr<Worker>(new Worker(WorkerID::FromRandom(), pid, language, -1,
|
||||
client, client_call_manager_));
|
||||
WorkerID worker_id = WorkerID::FromRandom();
|
||||
return std::shared_ptr<Worker>(new Worker(
|
||||
worker_id, pid, language, /* listening port */ -1, client_call_manager_));
|
||||
}
|
||||
|
||||
void SetWorkerCommands(const WorkerCommandMap &worker_commands) {
|
||||
@@ -96,10 +86,6 @@ class WorkerPoolTest : public ::testing::Test {
|
||||
boost::asio::io_service io_service_;
|
||||
int64_t error_message_type_;
|
||||
rpc::ClientCallManager client_call_manager_;
|
||||
|
||||
private:
|
||||
void HandleNewClient(LocalClientConnection &){};
|
||||
void HandleMessage(std::shared_ptr<LocalClientConnection>, int64_t, const uint8_t *){};
|
||||
};
|
||||
|
||||
static inline TaskSpecification ExampleTaskSpec(
|
||||
@@ -131,21 +117,23 @@ TEST_F(WorkerPoolTest, HandleWorkerRegistration) {
|
||||
workers.push_back(CreateWorker(pid));
|
||||
}
|
||||
for (const auto &worker : workers) {
|
||||
WorkerID worker_id = worker->GetWorkerId();
|
||||
// Check that there's still a starting worker process
|
||||
// before all workers have been registered
|
||||
ASSERT_EQ(worker_pool_.NumWorkerProcessesStarting(), 1);
|
||||
// Check that we cannot lookup the worker before it's registered.
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker->Connection()), nullptr);
|
||||
worker_pool_.RegisterWorker(worker);
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker_id), nullptr);
|
||||
worker_pool_.RegisterWorker(worker_id, worker);
|
||||
// Check that we can lookup the worker after it's registered.
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker->Connection()), worker);
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker_id), worker);
|
||||
}
|
||||
// Check that there's no starting worker process
|
||||
ASSERT_EQ(worker_pool_.NumWorkerProcessesStarting(), 0);
|
||||
for (const auto &worker : workers) {
|
||||
WorkerID worker_id = worker->GetWorkerId();
|
||||
worker_pool_.DisconnectWorker(worker);
|
||||
// Check that we cannot lookup the worker after it's disconnected.
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker->Connection()), nullptr);
|
||||
ASSERT_EQ(worker_pool_.GetRegisteredWorker(worker_id), nullptr);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -126,11 +126,14 @@ class ClientCallManager {
|
||||
explicit ClientCallManager(boost::asio::io_service &main_service)
|
||||
: main_service_(main_service) {
|
||||
// Start the polling thread.
|
||||
std::thread polling_thread(&ClientCallManager::PollEventsFromCompletionQueue, this);
|
||||
polling_thread.detach();
|
||||
polling_thread_ =
|
||||
std::thread(&ClientCallManager::PollEventsFromCompletionQueue, this);
|
||||
}
|
||||
|
||||
~ClientCallManager() { cq_.Shutdown(); }
|
||||
~ClientCallManager() {
|
||||
cq_.Shutdown();
|
||||
polling_thread_.join();
|
||||
}
|
||||
|
||||
/// Create a new `ClientCall` and send request.
|
||||
///
|
||||
@@ -177,7 +180,7 @@ class ClientCallManager {
|
||||
// Keep reading events from the `CompletionQueue` until it's shutdown.
|
||||
while (cq_.Next(&got_tag, &ok)) {
|
||||
auto tag = reinterpret_cast<ClientCallTag *>(got_tag);
|
||||
if (ok) {
|
||||
if (ok && !main_service_.stopped()) {
|
||||
// Post the callback to the main event loop.
|
||||
main_service_.post([tag]() {
|
||||
tag->GetCall()->OnReplyReceived();
|
||||
@@ -195,6 +198,9 @@ class ClientCallManager {
|
||||
|
||||
/// The gRPC `CompletionQueue` object used to poll events.
|
||||
grpc::CompletionQueue cq_;
|
||||
|
||||
/// Polling thread to check the completion queue.
|
||||
std::thread polling_thread_;
|
||||
};
|
||||
|
||||
} // namespace rpc
|
||||
|
||||
@@ -6,14 +6,20 @@ namespace ray {
|
||||
namespace rpc {
|
||||
|
||||
void GrpcServer::Run() {
|
||||
std::string server_address("0.0.0.0:" + std::to_string(port_));
|
||||
std::string server_address;
|
||||
// Set unix domain socket or tcp address.
|
||||
if (!unix_socket_path_.empty()) {
|
||||
server_address = "unix://" + unix_socket_path_;
|
||||
} else {
|
||||
server_address = "0.0.0.0:" + std::to_string(port_);
|
||||
}
|
||||
|
||||
grpc::ServerBuilder builder;
|
||||
// TODO(hchen): Add options for authentication.
|
||||
builder.AddListeningPort(server_address, grpc::InsecureServerCredentials(), &port_);
|
||||
// Register all the services to this server.
|
||||
if (services_.empty()) {
|
||||
RAY_LOG(WARNING) << "No service is found when start grpc server " << name_;
|
||||
RAY_LOG(WARNING) << "No service found when start grpc server " << name_;
|
||||
}
|
||||
for (auto &entry : services_) {
|
||||
builder.RegisterService(&entry.get());
|
||||
@@ -23,7 +29,11 @@ void GrpcServer::Run() {
|
||||
cq_ = builder.AddCompletionQueue();
|
||||
// Build and start server.
|
||||
server_ = builder.BuildAndStart();
|
||||
RAY_LOG(INFO) << name_ << " server started, listening on port " << port_ << ".";
|
||||
if (unix_socket_path_.empty()) {
|
||||
// For a TCP-based server, the actual port is decided after `AddListeningPort`.
|
||||
server_address = "0.0.0.0:" + std::to_string(port_);
|
||||
}
|
||||
RAY_LOG(INFO) << name_ << " server started, listening on " << server_address;
|
||||
|
||||
// Create calls for all the server call factories.
|
||||
for (auto &entry : server_call_factories_and_concurrencies_) {
|
||||
|
||||
@@ -26,15 +26,22 @@ class GrpcService;
|
||||
/// which kinds of requests this server should accept.
|
||||
class GrpcServer {
|
||||
public:
|
||||
/// Constructor.
|
||||
/// Construct a gRPC server that listens on a TCP port.
|
||||
///
|
||||
/// \param[in] name Name of this server, used for logging and debugging purpose.
|
||||
/// \param[in] port The port to bind this server to. If it's 0, a random available port
|
||||
/// will be chosen.
|
||||
/// \param[in] main_service The main event loop, to which service handler functions
|
||||
/// will be posted.
|
||||
GrpcServer(std::string name, const uint32_t port)
|
||||
: name_(std::move(name)), port_(port), is_closed_(true) {}
|
||||
: name_(std::move(name)), port_(port), is_closed_(false) {}
|
||||
|
||||
/// Construct a gRPC server that listens on unix domain socket.
|
||||
///
|
||||
/// \param[in] name Name of this server, used for logging and debugging purpose.
|
||||
/// \param[in] unix_socket_path Unix domain socket full path.
|
||||
GrpcServer(std::string name, const std::string &unix_socket_path)
|
||||
: GrpcServer(std::move(name), 0) {
|
||||
unix_socket_path_ = unix_socket_path;
|
||||
}
|
||||
|
||||
/// Destruct this gRPC server.
|
||||
~GrpcServer() { Shutdown(); }
|
||||
@@ -73,6 +80,10 @@ class GrpcServer {
|
||||
const std::string name_;
|
||||
/// Port of this server.
|
||||
int port_;
|
||||
/// Indicates whether this server has been closed.
|
||||
bool is_closed_;
|
||||
/// Unix domain socket path.
|
||||
std::string unix_socket_path_;
|
||||
/// The `grpc::Service` objects which should be registered to `ServerBuilder`.
|
||||
std::vector<std::reference_wrapper<grpc::Service>> services_;
|
||||
/// The `ServerCallFactory` objects, and the maximum number of concurrent requests that
|
||||
@@ -83,10 +94,8 @@ class GrpcServer {
|
||||
std::unique_ptr<grpc::ServerCompletionQueue> cq_;
|
||||
/// The `Server` object.
|
||||
std::unique_ptr<grpc::Server> server_;
|
||||
/// The polling thread used to check the completion queue
|
||||
/// The polling thread used to check the completion queue.
|
||||
std::thread polling_thread_;
|
||||
/// Flag indicates whether this server has closed
|
||||
bool is_closed_;
|
||||
};
|
||||
|
||||
/// Base class that represents an abstract gRPC service.
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
#include <grpcpp/grpcpp.h>
|
||||
|
||||
#include "ray/common/grpc_util.h"
|
||||
#include "ray/common/status.h"
|
||||
#include "ray/rpc/client_call.h"
|
||||
#include "ray/util/logging.h"
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
#include <grpcpp/grpcpp.h>
|
||||
|
||||
#include "ray/common/grpc_util.h"
|
||||
#include "ray/common/status.h"
|
||||
#include "ray/util/logging.h"
|
||||
#include "src/ray/protobuf/object_manager.grpc.pb.h"
|
||||
|
||||
@@ -0,0 +1,411 @@
|
||||
|
||||
#include "src/ray/rpc/raylet/raylet_client.h"
|
||||
|
||||
namespace ray {
|
||||
namespace rpc {
|
||||
|
||||
#define RETURN_IF_DISCONNECTED(connected) \
|
||||
do { \
|
||||
if (!(connected)) { \
|
||||
return Status::Invalid("Raylet connection is closed."); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
RayletClient::RayletClient(const std::string &raylet_socket, const WorkerID &worker_id,
|
||||
bool is_worker, const JobID &job_id, const Language &language,
|
||||
int port)
|
||||
: worker_id_(worker_id),
|
||||
is_worker_(is_worker),
|
||||
job_id_(job_id),
|
||||
language_(language),
|
||||
port_(port),
|
||||
main_service_(),
|
||||
work_(main_service_),
|
||||
client_call_manager_(main_service_),
|
||||
heartbeat_timer_(main_service_),
|
||||
is_connected_(false) {
|
||||
std::shared_ptr<grpc::Channel> channel =
|
||||
grpc::CreateChannel("unix://" + raylet_socket, grpc::InsecureChannelCredentials());
|
||||
stub_ = RayletService::NewStub(channel);
|
||||
|
||||
rpc_thread_ = std::thread([this]() { main_service_.run(); });
|
||||
RAY_LOG(DEBUG) << "Connecting to unix socket: "
|
||||
<< "unix://" + raylet_socket
|
||||
<< ", is worker: " << (is_worker_ ? "true" : "false")
|
||||
<< ", worker id: " << worker_id;
|
||||
// Try to register client `num_raylet_client_retry_times` times.
|
||||
TryRegisterClient(RayConfig::instance().num_raylet_client_retry_times());
|
||||
}
|
||||
|
||||
void RayletClient::TryRegisterClient(int retry_times) {
|
||||
// We should block here until register succeeds.
|
||||
for (int i = 0; i < retry_times; i++) {
|
||||
auto st = RegisterClient();
|
||||
if (st.ok()) {
|
||||
is_connected_ = true;
|
||||
Heartbeat();
|
||||
return;
|
||||
}
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(500));
|
||||
}
|
||||
RAY_LOG(FATAL) << "Failed to register to raylet server, worker id: " << worker_id_
|
||||
<< ", pid: " << static_cast<int>(getpid())
|
||||
<< ", is worker: " << is_worker_;
|
||||
}
|
||||
|
||||
RayletClient::~RayletClient() {
|
||||
is_connected_ = false;
|
||||
main_service_.stop();
|
||||
rpc_thread_.join();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::Disconnect() {
|
||||
DisconnectClientRequest disconnect_client_request;
|
||||
disconnect_client_request.set_worker_id(worker_id_.Binary());
|
||||
|
||||
DisconnectClientReply reply;
|
||||
grpc::ClientContext context;
|
||||
auto status = stub_->DisconnectClient(&context, disconnect_client_request, &reply);
|
||||
if (!status.ok()) {
|
||||
RAY_LOG(ERROR) << "Failed to disconnect from raylet, msg: " << status.error_message();
|
||||
}
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
SubmitTaskRequest submit_task_request;
|
||||
submit_task_request.mutable_task_spec()->CopyFrom(task_spec.GetMessage());
|
||||
|
||||
grpc::ClientContext context;
|
||||
SubmitTaskReply reply;
|
||||
auto status = stub_->SubmitTask(&context, submit_task_request, &reply);
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::GetTask(std::unique_ptr<ray::TaskSpecification> *task_spec) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
GetTaskRequest get_task_request;
|
||||
get_task_request.set_worker_id(worker_id_.Binary());
|
||||
|
||||
grpc::ClientContext context;
|
||||
GetTaskReply reply;
|
||||
// The actual RPC.
|
||||
auto status = stub_->GetTask(&context, get_task_request, &reply);
|
||||
|
||||
if (status.ok()) {
|
||||
resource_ids_.clear();
|
||||
// Parse resources that would be used by this assigned task.
|
||||
for (size_t i = 0; i < reply.fractional_resource_ids().size(); ++i) {
|
||||
auto const &fractional_resource_ids = reply.fractional_resource_ids()[i];
|
||||
auto &acquired_resources = resource_ids_[fractional_resource_ids.resource_name()];
|
||||
|
||||
// Each resource includes a series of resource IDs (e.g., GPU 0) and corresponding
|
||||
// amount for that resource ID. If the resource amount is fractional, then there
|
||||
// should only be one resource ID.
|
||||
size_t num_resource_ids = fractional_resource_ids.resource_ids().size();
|
||||
size_t num_resource_fractions = fractional_resource_ids.resource_fractions().size();
|
||||
RAY_CHECK(num_resource_ids == num_resource_fractions);
|
||||
RAY_CHECK(num_resource_ids > 0);
|
||||
for (size_t j = 0; j < num_resource_ids; ++j) {
|
||||
int64_t resource_id = fractional_resource_ids.resource_ids()[j];
|
||||
double resource_fraction = fractional_resource_ids.resource_fractions()[j];
|
||||
if (num_resource_ids > 1) {
|
||||
int64_t whole_fraction = resource_fraction;
|
||||
RAY_CHECK(whole_fraction == resource_fraction);
|
||||
}
|
||||
acquired_resources.emplace_back(resource_id, resource_fraction);
|
||||
}
|
||||
}
|
||||
task_spec->reset(new ray::TaskSpecification(reply.task_spec()));
|
||||
} else {
|
||||
*task_spec = nullptr;
|
||||
RAY_LOG(INFO) << "Failed to get task, msg: " << status.error_message();
|
||||
}
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::TaskDone() {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
TaskDoneRequest task_done_request;
|
||||
task_done_request.set_worker_id(worker_id_.Binary());
|
||||
|
||||
auto callback = [this](const Status &status, const TaskDoneReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send TaskDoneRequest, msg: " << status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_.CreateCall<RayletService, TaskDoneRequest, TaskDoneReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncTaskDone, task_done_request,
|
||||
callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::FetchOrReconstruct(const std::vector<ObjectID> &object_ids,
|
||||
bool fetch_only,
|
||||
const TaskID ¤t_task_id) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
FetchOrReconstructRequest fetch_or_reconstruct_request;
|
||||
fetch_or_reconstruct_request.set_fetch_only(fetch_only);
|
||||
fetch_or_reconstruct_request.set_task_id(current_task_id.Binary());
|
||||
fetch_or_reconstruct_request.set_worker_id(worker_id_.Binary());
|
||||
IdVectorToProtobuf<ObjectID, FetchOrReconstructRequest>(
|
||||
object_ids, fetch_or_reconstruct_request,
|
||||
&FetchOrReconstructRequest::add_object_ids);
|
||||
|
||||
// Callback to deal with reply.
|
||||
auto callback = [this](const Status &status, const FetchOrReconstructReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send FetchOrReconstructRequest, msg: "
|
||||
<< status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_
|
||||
.CreateCall<RayletService, FetchOrReconstructRequest, FetchOrReconstructReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncFetchOrReconstruct,
|
||||
fetch_or_reconstruct_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::NotifyUnblocked(const TaskID ¤t_task_id) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
NotifyUnblockedRequest notify_unblocked_request;
|
||||
notify_unblocked_request.set_worker_id(worker_id_.Binary());
|
||||
notify_unblocked_request.set_task_id(current_task_id.Binary());
|
||||
|
||||
auto callback = [this](const Status &status, const NotifyUnblockedReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send NotifyUnblockedRequest, msg: " << status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_
|
||||
.CreateCall<RayletService, NotifyUnblockedRequest, NotifyUnblockedReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncNotifyUnblocked,
|
||||
notify_unblocked_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::Wait(const std::vector<ObjectID> &object_ids, int num_returns,
|
||||
int64_t timeout_milliseconds, bool wait_local,
|
||||
const TaskID ¤t_task_id, WaitResultPair *result) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
WaitRequest wait_request;
|
||||
wait_request.set_worker_id(worker_id_.Binary());
|
||||
wait_request.set_timeout(timeout_milliseconds);
|
||||
wait_request.set_wait_local(wait_local);
|
||||
wait_request.set_task_id(current_task_id.Binary());
|
||||
wait_request.set_num_ready_objects(num_returns);
|
||||
IdVectorToProtobuf<ObjectID, WaitRequest>(object_ids, wait_request,
|
||||
&WaitRequest::add_object_ids);
|
||||
|
||||
grpc::ClientContext context;
|
||||
WaitReply reply;
|
||||
auto status = stub_->Wait(&context, wait_request, &reply);
|
||||
|
||||
if (status.ok()) {
|
||||
result->first = IdVectorFromProtobuf<ObjectID>(reply.found());
|
||||
result->second = IdVectorFromProtobuf<ObjectID>(reply.remaining());
|
||||
} else {
|
||||
RAY_LOG(INFO) << "Failed to send WaitRequest, msg: " << status.error_message();
|
||||
}
|
||||
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PushError(const ray::JobID &job_id, const std::string &type,
|
||||
const std::string &error_message, double timestamp) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
PushErrorRequest push_error_request;
|
||||
push_error_request.set_job_id(job_id.Binary());
|
||||
push_error_request.set_type(type);
|
||||
push_error_request.set_error_message(error_message);
|
||||
push_error_request.set_timestamp(timestamp);
|
||||
|
||||
auto callback = [this](const Status &status, const PushErrorReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send PushErrorRequest, msg: " << status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_.CreateCall<RayletService, PushErrorRequest, PushErrorReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncPushError, push_error_request,
|
||||
callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PushProfileEvents(const ProfileTableData &profile_events) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
PushProfileEventsRequest push_profile_events_request;
|
||||
push_profile_events_request.mutable_profile_table_data()->CopyFrom(profile_events);
|
||||
|
||||
auto callback = [this](const Status &status, const PushProfileEventsReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send PushProfileEventsRequest, msg: "
|
||||
<< status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_
|
||||
.CreateCall<RayletService, PushProfileEventsRequest, PushProfileEventsReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncPushProfileEvents,
|
||||
push_profile_events_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::FreeObjects(const std::vector<ray::ObjectID> &object_ids,
|
||||
bool local_only, bool delete_creating_tasks) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
FreeObjectsInStoreRequest free_objects_request;
|
||||
free_objects_request.set_local_only(local_only);
|
||||
free_objects_request.set_delete_creating_tasks(delete_creating_tasks);
|
||||
IdVectorToProtobuf<ray::ObjectID, FreeObjectsInStoreRequest>(
|
||||
object_ids, free_objects_request, &FreeObjectsInStoreRequest::add_object_ids);
|
||||
|
||||
auto callback = [this](const Status &status, const FreeObjectsInStoreReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Failed to send FreeObjectsInStoreRequest, msg: "
|
||||
<< status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_
|
||||
.CreateCall<RayletService, FreeObjectsInStoreRequest, FreeObjectsInStoreReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncFreeObjectsInStore,
|
||||
free_objects_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::PrepareActorCheckpoint(const ActorID &actor_id,
|
||||
ActorCheckpointID &checkpoint_id) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
PrepareActorCheckpointRequest prepare_actor_checkpoint_request;
|
||||
prepare_actor_checkpoint_request.set_actor_id(actor_id.Binary());
|
||||
prepare_actor_checkpoint_request.set_worker_id(worker_id_.Binary());
|
||||
|
||||
grpc::ClientContext context;
|
||||
PrepareActorCheckpointReply reply;
|
||||
auto status =
|
||||
stub_->PrepareActorCheckpoint(&context, prepare_actor_checkpoint_request, &reply);
|
||||
|
||||
if (status.ok()) {
|
||||
checkpoint_id = ActorCheckpointID::FromBinary(reply.checkpoint_id());
|
||||
} else {
|
||||
RAY_LOG(INFO) << "Failed to send PrepareActorCheckpointRequest, msg: "
|
||||
<< status.error_message();
|
||||
}
|
||||
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
ray::Status RayletClient::NotifyActorResumedFromCheckpoint(
|
||||
const ActorID &actor_id, const ActorCheckpointID &checkpoint_id) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
NotifyActorResumedFromCheckpointRequest notify_actor_resumed_from_checkpoint_request;
|
||||
notify_actor_resumed_from_checkpoint_request.set_actor_id(actor_id.Binary());
|
||||
notify_actor_resumed_from_checkpoint_request.set_checkpoint_id(checkpoint_id.Binary());
|
||||
|
||||
auto callback = [this](const Status &status,
|
||||
const NotifyActorResumedFromCheckpointReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "NotifyActorResumedFromCheckpoint failed, msg: "
|
||||
<< status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call =
|
||||
client_call_manager_
|
||||
.CreateCall<RayletService, NotifyActorResumedFromCheckpointRequest,
|
||||
NotifyActorResumedFromCheckpointReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncNotifyActorResumedFromCheckpoint,
|
||||
notify_actor_resumed_from_checkpoint_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::SetResource(const std::string &resource_name,
|
||||
const double capacity,
|
||||
const ray::ClientID &client_id) {
|
||||
RETURN_IF_DISCONNECTED(is_connected_);
|
||||
SetResourceRequest set_resource_request;
|
||||
set_resource_request.set_resource_name(resource_name);
|
||||
set_resource_request.set_capacity(capacity);
|
||||
set_resource_request.set_client_id(client_id.Binary());
|
||||
|
||||
auto callback = [this](const Status &status, const SetResourceReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "SetResource failed, msg: " << status.message();
|
||||
}
|
||||
};
|
||||
|
||||
auto call = client_call_manager_
|
||||
.CreateCall<RayletService, SetResourceRequest, SetResourceReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncSetResource,
|
||||
set_resource_request, callback);
|
||||
return call->GetStatus();
|
||||
}
|
||||
|
||||
ray::Status RayletClient::RegisterClient() {
|
||||
RegisterClientRequest register_client_request;
|
||||
register_client_request.set_is_worker(is_worker_);
|
||||
register_client_request.set_worker_id(worker_id_.Binary());
|
||||
register_client_request.set_worker_pid(getpid());
|
||||
register_client_request.set_job_id(job_id_.Binary());
|
||||
register_client_request.set_language(language_);
|
||||
register_client_request.set_port(port_);
|
||||
|
||||
grpc::ClientContext context;
|
||||
RegisterClientReply reply;
|
||||
auto status = stub_->RegisterClient(&context, register_client_request, &reply);
|
||||
|
||||
if (!status.ok()) {
|
||||
RAY_LOG(DEBUG) << "Failed to register client, msg: " << status.error_message();
|
||||
}
|
||||
|
||||
return GrpcStatusToRayStatus(status);
|
||||
}
|
||||
|
||||
void RayletClient::Heartbeat() {
|
||||
if (!is_connected_) {
|
||||
return;
|
||||
}
|
||||
HeartbeatRequest heartbeat_request;
|
||||
heartbeat_request.set_is_worker(is_worker_);
|
||||
heartbeat_request.set_worker_id(worker_id_.Binary());
|
||||
|
||||
auto callback = [this](const Status &status, const HeartbeatReply &reply) {
|
||||
if (!status.ok() && is_connected_) {
|
||||
is_connected_ = false;
|
||||
RAY_LOG(INFO) << "Heartbeat failed, msg: " << status.message();
|
||||
}
|
||||
};
|
||||
auto call =
|
||||
client_call_manager_.CreateCall<RayletService, HeartbeatRequest, HeartbeatReply>(
|
||||
*stub_, &RayletService::Stub::PrepareAsyncHeartbeat, heartbeat_request,
|
||||
callback);
|
||||
|
||||
heartbeat_timer_.expires_from_now(boost::posix_time::milliseconds(
|
||||
RayConfig::instance().heartbeat_timeout_milliseconds()));
|
||||
heartbeat_timer_.async_wait([this](const boost::system::error_code &error) {
|
||||
RAY_CHECK(!error);
|
||||
Heartbeat();
|
||||
});
|
||||
}
|
||||
|
||||
} // namespace rpc
|
||||
} // namespace ray
|
||||
@@ -1,84 +1,66 @@
|
||||
#ifndef RAYLET_CLIENT_H
|
||||
#define RAYLET_CLIENT_H
|
||||
#ifndef RAY_RPC_RAYLET_CLIENT_H
|
||||
#define RAY_RPC_RAYLET_CLIENT_H
|
||||
|
||||
#include <ray/common/ray_config.h>
|
||||
#include <ray/protobuf/gcs.pb.h>
|
||||
#include <unistd.h>
|
||||
#include <future>
|
||||
#include <mutex>
|
||||
#include <thread>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include <grpcpp/grpcpp.h>
|
||||
|
||||
#include "ray/common/status.h"
|
||||
#include "ray/common/task/task_spec.h"
|
||||
#include "src/ray/common/status.h"
|
||||
#include "src/ray/protobuf/raylet.grpc.pb.h"
|
||||
#include "src/ray/protobuf/raylet.pb.h"
|
||||
#include "src/ray/rpc/client_call.h"
|
||||
|
||||
using ray::ActorCheckpointID;
|
||||
using ray::ActorID;
|
||||
using ray::ClientID;
|
||||
using ray::JobID;
|
||||
using ray::ObjectID;
|
||||
using ray::TaskID;
|
||||
using ray::UniqueID;
|
||||
using ray::WorkerID;
|
||||
|
||||
using ray::Language;
|
||||
using ray::rpc::ProfileTableData;
|
||||
|
||||
using MessageType = ray::protocol::MessageType;
|
||||
using ResourceMappingType =
|
||||
std::unordered_map<std::string, std::vector<std::pair<int64_t, double>>>;
|
||||
using WaitResultPair = std::pair<std::vector<ObjectID>, std::vector<ObjectID>>;
|
||||
|
||||
class RayletConnection {
|
||||
public:
|
||||
/// Connect to the raylet.
|
||||
///
|
||||
/// \param raylet_socket The name of the socket to use to connect to the raylet.
|
||||
/// \param worker_id A unique ID to represent the worker.
|
||||
/// \param is_worker Whether this client is a worker. If it is a worker, an
|
||||
/// additional message will be sent to register as one.
|
||||
/// \param job_id The ID of the driver. This is non-nil if the client is a
|
||||
/// driver.
|
||||
/// \return The connection information.
|
||||
RayletConnection(const std::string &raylet_socket, int num_retries, int64_t timeout);
|
||||
namespace ray {
|
||||
namespace rpc {
|
||||
|
||||
~RayletConnection() { close(conn_); }
|
||||
/// Notify the raylet that this client is disconnecting gracefully. This
|
||||
/// is used by actors to exit gracefully so that the raylet doesn't
|
||||
/// propagate an error message to the driver.
|
||||
///
|
||||
/// \return ray::Status.
|
||||
ray::Status Disconnect();
|
||||
ray::Status ReadMessage(MessageType type, std::unique_ptr<uint8_t[]> &message);
|
||||
ray::Status WriteMessage(MessageType type,
|
||||
flatbuffers::FlatBufferBuilder *fbb = nullptr);
|
||||
ray::Status AtomicRequestReply(MessageType request_type, MessageType reply_type,
|
||||
std::unique_ptr<uint8_t[]> &reply_message,
|
||||
flatbuffers::FlatBufferBuilder *fbb = nullptr);
|
||||
|
||||
private:
|
||||
/// File descriptor of the Unix domain socket that connects to raylet.
|
||||
int conn_;
|
||||
/// A mutex to protect stateful operations of the raylet client.
|
||||
std::mutex mutex_;
|
||||
/// A mutex to protect write operations of the raylet client.
|
||||
std::mutex write_mutex_;
|
||||
};
|
||||
using ResourceMappingType =
|
||||
std::unordered_map<std::string, std::vector<std::pair<int64_t, double>>>;
|
||||
|
||||
/// Client used for communicating with the raylet.
|
||||
class RayletClient {
|
||||
public:
|
||||
/// Connect to the raylet.
|
||||
/// Constructor for the raylet client.
|
||||
/// TODO(jzh): At present, client call manager and reply handler service are generated
|
||||
/// in raylet client. Instead, we should add parameters to the constructor and pass them
|
||||
/// in. Change them as input parameters once we changed the worker into a server.
|
||||
///
|
||||
/// \param raylet_socket The name of the socket to use to connect to the raylet.
|
||||
/// \param worker_id A unique ID to represent the worker.
|
||||
/// \param is_worker Whether this client is a worker. If it is a worker, an
|
||||
/// additional message will be sent to register as one.
|
||||
/// \param job_id The ID of the driver. This is non-nil if the client is a driver.
|
||||
/// \return The connection information.
|
||||
RayletClient(const std::string &raylet_socket, const ClientID &client_id,
|
||||
/// \param[in] raylet_socket Unix domain socket of the raylet server.
|
||||
/// \param[in] worker_id The worker id.
|
||||
/// \param[in] is_worker Indicates whether a worker or a driver.
|
||||
/// \param[in] job_id The job id that this raylet client belongs to.
|
||||
/// \param[in] language The language type, python or java.
|
||||
/// \param[in] port The listening port of the worker server, -1 means that the worker
|
||||
/// does not have a server.
|
||||
RayletClient(const std::string &raylet_socket, const WorkerID &worker_id,
|
||||
bool is_worker, const JobID &job_id, const Language &language,
|
||||
int port = -1);
|
||||
|
||||
ray::Status Disconnect() { return conn_->Disconnect(); };
|
||||
~RayletClient();
|
||||
|
||||
/// Submit a task using the raylet code path.
|
||||
/// Send disconnect request to local raylet.
|
||||
ray::Status Disconnect();
|
||||
|
||||
/// Submit a task to the local raylet.
|
||||
///
|
||||
/// \param The task specification.
|
||||
/// \return ray::Status.
|
||||
@@ -178,7 +160,7 @@ class RayletClient {
|
||||
|
||||
Language GetLanguage() const { return language_; }
|
||||
|
||||
ClientID GetClientID() const { return client_id_; }
|
||||
WorkerID GetWorkerId() const { return worker_id_; }
|
||||
|
||||
JobID GetJobID() const { return job_id_; }
|
||||
|
||||
@@ -187,7 +169,21 @@ class RayletClient {
|
||||
const ResourceMappingType &GetResourceIDs() const { return resource_ids_; }
|
||||
|
||||
private:
|
||||
const ClientID client_id_;
|
||||
/// Try to register client in raylet, we would retry serveral time to
|
||||
/// reconnect if failed. We need this because raylet client may start before raylet
|
||||
/// server.
|
||||
///
|
||||
/// \param times Number of times to retry.
|
||||
void TryRegisterClient(int retry_times);
|
||||
|
||||
ray::Status RegisterClient();
|
||||
|
||||
/// Send heartbeat requests to the raylet server.
|
||||
void Heartbeat();
|
||||
/// Id of the worker to which this raylet client belongs.
|
||||
const WorkerID worker_id_;
|
||||
/// Indicates whether this worker is a driver worker.
|
||||
/// Driver is treated as a special worker.
|
||||
const bool is_worker_;
|
||||
const JobID job_id_;
|
||||
const Language language_;
|
||||
@@ -196,8 +192,30 @@ class RayletClient {
|
||||
/// for this worker. Each pair consists of the resource ID and the fraction
|
||||
/// of that resource allocated for this worker.
|
||||
ResourceMappingType resource_ids_;
|
||||
/// The connection to the raylet server.
|
||||
std::unique_ptr<RayletConnection> conn_;
|
||||
|
||||
/// The gRPC-generated stub.
|
||||
std::unique_ptr<RayletService::Stub> stub_;
|
||||
|
||||
/// Service for handling reply.
|
||||
boost::asio::io_service main_service_;
|
||||
|
||||
/// Asio work for main service.
|
||||
boost::asio::io_service::work work_;
|
||||
|
||||
/// The `ClientCallManager` used for managing requests.
|
||||
ClientCallManager client_call_manager_;
|
||||
|
||||
/// The thread used to handle reply.
|
||||
std::thread rpc_thread_;
|
||||
|
||||
/// Heartbeat timer.
|
||||
boost::asio::deadline_timer heartbeat_timer_;
|
||||
|
||||
/// Indicates whether the connection has been closed.
|
||||
bool is_connected_;
|
||||
};
|
||||
|
||||
#endif
|
||||
} // namespace rpc
|
||||
} // namespace ray
|
||||
|
||||
#endif // RAY_RPC_RAYLET_CLIENT_H
|
||||
@@ -0,0 +1,258 @@
|
||||
#ifndef RAY_RPC_RAYLET_SERVER_H
|
||||
#define RAY_RPC_RAYLET_SERVER_H
|
||||
|
||||
#include "src/ray/rpc/grpc_server.h"
|
||||
#include "src/ray/rpc/server_call.h"
|
||||
|
||||
#include "src/ray/protobuf/raylet.grpc.pb.h"
|
||||
#include "src/ray/protobuf/raylet.pb.h"
|
||||
|
||||
namespace ray {
|
||||
namespace rpc {
|
||||
|
||||
/// Implementations of the `RayletService`, check interface in
|
||||
/// `src/ray/protobuf/raylet.proto`.
|
||||
class RayletServiceHandler {
|
||||
public:
|
||||
virtual void HandleRegisterClientRequest(const RegisterClientRequest &request,
|
||||
RegisterClientReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `SubmitTask` request.
|
||||
/// The implementation can handle this request asynchronously. When handling is done,
|
||||
/// the `send_reply_callback` should be called.
|
||||
///
|
||||
/// \param[in] request The request message.
|
||||
/// \param[out] reply The reply message.
|
||||
/// \param[in] send_reply_callback The callback to be called when the request is done.
|
||||
virtual void HandleSubmitTaskRequest(const SubmitTaskRequest &request,
|
||||
SubmitTaskReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `DisconnectClient` request.
|
||||
virtual void HandleDisconnectClientRequest(const DisconnectClientRequest &request,
|
||||
DisconnectClientReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `GetTask` request.
|
||||
virtual void HandleGetTaskRequest(const GetTaskRequest &request, GetTaskReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `TaskDone` request.
|
||||
virtual void HandleTaskDoneRequest(const TaskDoneRequest &request, TaskDoneReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `HandleFetchOrReconstruct` request.
|
||||
virtual void HandleFetchOrReconstructRequest(const FetchOrReconstructRequest &request,
|
||||
FetchOrReconstructReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `HandleNotifyUnblocked` request.
|
||||
virtual void HandleNotifyUnblockedRequest(const NotifyUnblockedRequest &request,
|
||||
NotifyUnblockedReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `Wait` request.
|
||||
virtual void HandleWaitRequest(const WaitRequest &request, WaitReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `PushError` request.
|
||||
virtual void HandlePushErrorRequest(const PushErrorRequest &request,
|
||||
PushErrorReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `PushProfileEvents` request.
|
||||
virtual void HandlePushProfileEventsRequest(const PushProfileEventsRequest &request,
|
||||
PushProfileEventsReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `FreeObjectsInStoreInObjectStore` request.
|
||||
virtual void HandleFreeObjectsInStoreRequest(const FreeObjectsInStoreRequest &request,
|
||||
FreeObjectsInStoreReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `PrepareActorCheckpoint` request.
|
||||
virtual void HandlePrepareActorCheckpointRequest(
|
||||
const PrepareActorCheckpointRequest &request, PrepareActorCheckpointReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `NotifyActorResumedFromCheckpoint` request.
|
||||
virtual void HandleNotifyActorResumedFromCheckpointRequest(
|
||||
const NotifyActorResumedFromCheckpointRequest &request,
|
||||
NotifyActorResumedFromCheckpointReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `SetResource` request.
|
||||
virtual void HandleSetResourceRequest(const SetResourceRequest &request,
|
||||
SetResourceReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
/// Handle a `SetResourceReply` request.
|
||||
virtual void HandleHeartbeatRequest(const HeartbeatRequest &request,
|
||||
HeartbeatReply *reply,
|
||||
SendReplyCallback send_reply_callback) = 0;
|
||||
};
|
||||
|
||||
/// The `GrpcService` for `RayletGrpcService`.
|
||||
class RayletGrpcService : public GrpcService {
|
||||
public:
|
||||
/// Construct a `RayletGrpcService`.
|
||||
///
|
||||
/// \param[in] io_service Service used to handle incoming requests
|
||||
/// \param[in] handler The service handler that actually handle the requests.
|
||||
RayletGrpcService(boost::asio::io_service &io_service,
|
||||
RayletServiceHandler &service_handler)
|
||||
: GrpcService(io_service), service_handler_(service_handler){};
|
||||
|
||||
protected:
|
||||
grpc::Service &GetGrpcService() override { return service_; }
|
||||
|
||||
void InitServerCallFactories(
|
||||
const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
|
||||
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
|
||||
*server_call_factories_and_concurrencies) override {
|
||||
// Initialize the factory for `RegisterClient` requests.
|
||||
std::unique_ptr<ServerCallFactory> register_client_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
RegisterClientRequest, RegisterClientReply>(
|
||||
service_, &RayletService::AsyncService::RequestRegisterClient,
|
||||
service_handler_, &RayletServiceHandler::HandleRegisterClientRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(register_client_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `SubmitTask` requests.
|
||||
std::unique_ptr<ServerCallFactory> submit_task_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, SubmitTaskRequest,
|
||||
SubmitTaskReply>(
|
||||
service_, &RayletService::AsyncService::RequestSubmitTask, service_handler_,
|
||||
&RayletServiceHandler::HandleSubmitTaskRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(submit_task_call_factory), 20);
|
||||
|
||||
// Initialize the factory for `DisconnectClient` requests.
|
||||
std::unique_ptr<ServerCallFactory> disconnect_client_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
DisconnectClientRequest, DisconnectClientReply>(
|
||||
service_, &RayletService::AsyncService::RequestDisconnectClient,
|
||||
service_handler_, &RayletServiceHandler::HandleDisconnectClientRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(disconnect_client_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `GetTask` requests.
|
||||
std::unique_ptr<ServerCallFactory> get_task_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, GetTaskRequest,
|
||||
GetTaskReply>(
|
||||
service_, &RayletService::AsyncService::RequestGetTask, service_handler_,
|
||||
&RayletServiceHandler::HandleGetTaskRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(get_task_call_factory), 20);
|
||||
|
||||
// Initialize the factory for `TaskDone` requests.
|
||||
std::unique_ptr<ServerCallFactory> task_done_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, TaskDoneRequest,
|
||||
TaskDoneReply>(
|
||||
service_, &RayletService::AsyncService::RequestTaskDone, service_handler_,
|
||||
&RayletServiceHandler::HandleTaskDoneRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(task_done_call_factory), 20);
|
||||
|
||||
// Initialize the factory for `FetchOrReconstruct` requests.
|
||||
std::unique_ptr<ServerCallFactory> fetch_or_reconstruct_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
FetchOrReconstructRequest, FetchOrReconstructReply>(
|
||||
service_, &RayletService::AsyncService::RequestFetchOrReconstruct,
|
||||
service_handler_, &RayletServiceHandler::HandleFetchOrReconstructRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(fetch_or_reconstruct_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `NotifyUnblocked` requests.
|
||||
std::unique_ptr<ServerCallFactory> notify_unblocked_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
NotifyUnblockedRequest, NotifyUnblockedReply>(
|
||||
service_, &RayletService::AsyncService::RequestNotifyUnblocked,
|
||||
service_handler_, &RayletServiceHandler::HandleNotifyUnblockedRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(notify_unblocked_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `Wait` requests.
|
||||
std::unique_ptr<ServerCallFactory> wait_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, WaitRequest,
|
||||
WaitReply>(
|
||||
service_, &RayletService::AsyncService::RequestWait, service_handler_,
|
||||
&RayletServiceHandler::HandleWaitRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(std::move(wait_call_factory),
|
||||
20);
|
||||
|
||||
// Initialize the factory for `PushError` requests.
|
||||
std::unique_ptr<ServerCallFactory> push_error_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, PushErrorRequest,
|
||||
PushErrorReply>(
|
||||
service_, &RayletService::AsyncService::RequestPushError, service_handler_,
|
||||
&RayletServiceHandler::HandlePushErrorRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(push_error_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `PushProfileEvents` requests.
|
||||
std::unique_ptr<ServerCallFactory> push_profile_events_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
PushProfileEventsRequest, PushProfileEventsReply>(
|
||||
service_, &RayletService::AsyncService::RequestPushProfileEvents,
|
||||
service_handler_, &RayletServiceHandler::HandlePushProfileEventsRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(push_profile_events_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `FreeObjectsInStore` requests.
|
||||
std::unique_ptr<ServerCallFactory> free_objects_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
FreeObjectsInStoreRequest, FreeObjectsInStoreReply>(
|
||||
service_, &RayletService::AsyncService::RequestFreeObjectsInStore,
|
||||
service_handler_, &RayletServiceHandler::HandleFreeObjectsInStoreRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(free_objects_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `PrepareActorCheckpoint` requests.
|
||||
std::unique_ptr<ServerCallFactory> prepare_actor_checkpoint_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
PrepareActorCheckpointRequest,
|
||||
PrepareActorCheckpointReply>(
|
||||
service_, &RayletService::AsyncService::RequestPrepareActorCheckpoint,
|
||||
service_handler_, &RayletServiceHandler::HandlePrepareActorCheckpointRequest,
|
||||
cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(prepare_actor_checkpoint_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `NotifyActorResumedFromCheckpoint` requests.
|
||||
std::unique_ptr<ServerCallFactory> notify_actor_resumed_from_checkpoint_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler,
|
||||
NotifyActorResumedFromCheckpointRequest,
|
||||
NotifyActorResumedFromCheckpointReply>(
|
||||
service_,
|
||||
&RayletService::AsyncService::RequestNotifyActorResumedFromCheckpoint,
|
||||
service_handler_,
|
||||
&RayletServiceHandler::HandleNotifyActorResumedFromCheckpointRequest, cq,
|
||||
main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(notify_actor_resumed_from_checkpoint_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `SetResource` requests.
|
||||
std::unique_ptr<ServerCallFactory> set_resource_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, SetResourceRequest,
|
||||
SetResourceReply>(
|
||||
service_, &RayletService::AsyncService::RequestSetResource, service_handler_,
|
||||
&RayletServiceHandler::HandleSetResourceRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(set_resource_call_factory), 10);
|
||||
|
||||
// Initialize the factory for `Heartbeat` requests.
|
||||
std::unique_ptr<ServerCallFactory> heartbeat_call_factory(
|
||||
new ServerCallFactoryImpl<RayletService, RayletServiceHandler, HeartbeatRequest,
|
||||
HeartbeatReply>(
|
||||
service_, &RayletService::AsyncService::RequestHeartbeat, service_handler_,
|
||||
&RayletServiceHandler::HandleHeartbeatRequest, cq, main_service_));
|
||||
server_call_factories_and_concurrencies->emplace_back(
|
||||
std::move(heartbeat_call_factory), 10);
|
||||
}
|
||||
|
||||
private:
|
||||
/// The grpc async service object.
|
||||
RayletService::AsyncService service_;
|
||||
/// The service handler that actually handle the requests.
|
||||
RayletServiceHandler &service_handler_;
|
||||
};
|
||||
|
||||
} // namespace rpc
|
||||
} // namespace ray
|
||||
|
||||
#endif
|
||||
@@ -2,6 +2,7 @@
|
||||
#define RAY_RPC_SERVER_CALL_H
|
||||
|
||||
#include <grpcpp/grpcpp.h>
|
||||
#include <boost/asio.hpp>
|
||||
|
||||
#include "ray/common/grpc_util.h"
|
||||
#include "ray/common/status.h"
|
||||
|
||||
Reference in New Issue
Block a user