wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-26 13:37:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

[GCS]Add gcs table storage interface (#7949)

Browse Source
This commit is contained in:
fangfengbin
2020-04-15 10:48:12 +08:00
committed by GitHub
parent b4656ca244
commit c17404918c
12 changed files with 874 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BUILD.bazel
+53 -5
View File
@@ -714,7 +714,10 @@ cc_test(
cc_test(
name = "gcs_server_rpc_test",
srcs = ["src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc"],
srcs = [
"src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc",
"src/ray/gcs/test/gcs_test_util.h",
],
args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"],
copts = COPTS,
data = [
@@ -732,7 +735,7 @@ cc_test(
name = "gcs_node_manager_test",
srcs = [
"src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc",
"src/ray/gcs/gcs_server/test/gcs_test_util.h",
"src/ray/gcs/test/gcs_test_util.h",
],
copts = COPTS,
deps = [
@@ -745,7 +748,7 @@ cc_test(
name = "gcs_actor_scheduler_test",
srcs = [
"src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc",
"src/ray/gcs/gcs_server/test/gcs_test_util.h",
"src/ray/gcs/test/gcs_test_util.h",
],
copts = COPTS,
deps = [
@@ -758,7 +761,7 @@ cc_test(
name = "gcs_actor_manager_test",
srcs = [
"src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc",
"src/ray/gcs/gcs_server/test/gcs_test_util.h",
"src/ray/gcs/test/gcs_test_util.h",
],
copts = COPTS,
deps = [
@@ -767,6 +770,47 @@ cc_test(
],
)
cc_library(
name = "gcs_table_storage_lib",
srcs = glob(
[
"src/ray/gcs/gcs_server/gcs_table_storage.cc",
],
),
hdrs = glob(
[
"src/ray/gcs/gcs_server/gcs_table_storage.h",
],
),
copts = COPTS,
deps = [
":gcs",
":ray_common",
":redis_store_client",
],
)
cc_test(
name = "redis_gcs_table_storage_test",
srcs = [
"src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc",
"src/ray/gcs/test/gcs_test_util.h",
],
args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"],
copts = COPTS,
data = [
"//:libray_redis_module.so",
"//:redis-cli",
"//:redis-server",
],
deps = [
":gcs_server_lib",
":gcs_table_storage_lib",
":store_client_test_lib",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "service_based_gcs_client_lib",
srcs = glob(
@@ -783,12 +827,16 @@ cc_library(
deps = [
":gcs",
":gcs_service_rpc",
":redis_store_client",
],
)
cc_test(
name = "gcs_server_test",
srcs = ["src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc"],
srcs = [
"src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc",
"src/ray/gcs/test/gcs_test_util.h",
],
args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"],
copts = COPTS,
data = [
src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
+20 -87
View File
@@ -17,6 +17,7 @@
#include "ray/common/test_util.h"
#include "ray/gcs/gcs_client/service_based_accessor.h"
#include "ray/gcs/gcs_server/gcs_server.h"
#include "ray/gcs/test/gcs_test_util.h"
#include "ray/rpc/gcs_server/gcs_rpc_client.h"
namespace ray {
@@ -452,74 +453,6 @@ class ServiceBasedGcsClientTest : public RedisServiceManagerForTest {
EXPECT_TRUE(WaitForCondition(condition, timeout_ms_.count()));
}
std::shared_ptr<rpc::JobTableData> GenJobTableData(JobID job_id) {
auto job_table_data = std::make_shared<rpc::JobTableData>();
job_table_data->set_job_id(job_id.Binary());
job_table_data->set_is_dead(false);
job_table_data->set_timestamp(std::time(nullptr));
job_table_data->set_node_manager_address("127.0.0.1");
job_table_data->set_driver_pid(5667L);
return job_table_data;
}
std::shared_ptr<rpc::ActorTableData> GenActorTableData(const JobID &job_id) {
auto actor_table_data = std::make_shared<rpc::ActorTableData>();
ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0);
actor_table_data->set_actor_id(actor_id.Binary());
actor_table_data->set_job_id(job_id.Binary());
actor_table_data->set_state(
rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE);
actor_table_data->set_max_reconstructions(1);
actor_table_data->set_remaining_reconstructions(1);
return actor_table_data;
}
rpc::GcsNodeInfo GenGcsNodeInfo(const std::string &node_id) {
rpc::GcsNodeInfo gcs_node_info;
gcs_node_info.set_node_id(node_id);
gcs_node_info.set_state(rpc::GcsNodeInfo_GcsNodeState_ALIVE);
gcs_node_info.set_node_manager_address("127.0.0.1");
return gcs_node_info;
}
std::shared_ptr<rpc::TaskTableData> GenTaskTableData(const std::string &job_id,
const std::string &task_id) {
auto task_table_data = std::make_shared<rpc::TaskTableData>();
rpc::Task task;
rpc::TaskSpec task_spec;
task_spec.set_job_id(job_id);
task_spec.set_task_id(task_id);
task.mutable_task_spec()->CopyFrom(task_spec);
task_table_data->mutable_task()->CopyFrom(task);
return task_table_data;
}
std::shared_ptr<rpc::TaskLeaseData> GenTaskLeaseData(const std::string &task_id,
const std::string &node_id) {
auto task_lease_data = std::make_shared<rpc::TaskLeaseData>();
task_lease_data->set_task_id(task_id);
task_lease_data->set_node_manager_id(node_id);
return task_lease_data;
}
std::shared_ptr<rpc::ProfileTableData> GenProfileTableData(const ClientID &node_id) {
auto profile_table_data = std::make_shared<rpc::ProfileTableData>();
profile_table_data->set_component_id(node_id.Binary());
return profile_table_data;
}
std::shared_ptr<rpc::ErrorTableData> GenErrorTableData(const JobID &job_id) {
auto error_table_data = std::make_shared<rpc::ErrorTableData>();
error_table_data->set_job_id(job_id.Binary());
return error_table_data;
}
std::shared_ptr<rpc::WorkerFailureData> GenWorkerFailureData() {
auto worker_failure_data = std::make_shared<rpc::WorkerFailureData>();
worker_failure_data->set_timestamp(std::time(nullptr));
return worker_failure_data;
}
// GCS server.
gcs::GcsServerConfig config;
std::unique_ptr<gcs::GcsServer> gcs_server_;
@@ -537,7 +470,7 @@ class ServiceBasedGcsClientTest : public RedisServiceManagerForTest {
TEST_F(ServiceBasedGcsClientTest, TestJobInfo) {
// Create job table data.
JobID add_job_id = JobID::FromInt(1);
auto job_table_data = GenJobTableData(add_job_id);
auto job_table_data = Mocker::GenJobTableData(add_job_id);
// Subscribe to finished jobs.
std::atomic<int> finished_job_count(0);
@@ -555,7 +488,7 @@ TEST_F(ServiceBasedGcsClientTest, TestJobInfo) {
TEST_F(ServiceBasedGcsClientTest, TestActorInfo) {
// Create actor table data.
JobID job_id = JobID::FromInt(1);
auto actor_table_data = GenActorTableData(job_id);
auto actor_table_data = Mocker::GenActorTableData(job_id);
ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
// Subscribe to any update operations of an actor.
@@ -586,7 +519,7 @@ TEST_F(ServiceBasedGcsClientTest, TestActorInfo) {
TEST_F(ServiceBasedGcsClientTest, TestActorCheckpoint) {
// Create actor checkpoint data.
JobID job_id = JobID::FromInt(1);
auto actor_table_data = GenActorTableData(job_id);
auto actor_table_data = Mocker::GenActorTableData(job_id);
ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom();
@@ -611,8 +544,8 @@ TEST_F(ServiceBasedGcsClientTest, TestActorCheckpoint) {
TEST_F(ServiceBasedGcsClientTest, TestActorSubscribeAll) {
// Create actor table data.
JobID job_id = JobID::FromInt(1);
auto actor_table_data1 = GenActorTableData(job_id);
auto actor_table_data2 = GenActorTableData(job_id);
auto actor_table_data1 = Mocker::GenActorTableData(job_id);
auto actor_table_data2 = Mocker::GenActorTableData(job_id);
// Subscribe to any register or update operations of actors.
std::atomic<int> actor_update_count(0);
@@ -630,8 +563,8 @@ TEST_F(ServiceBasedGcsClientTest, TestActorSubscribeAll) {
TEST_F(ServiceBasedGcsClientTest, TestNodeInfo) {
// Create gcs node info.
ClientID node1_id = ClientID::FromRandom();
auto gcs_node1_info = GenGcsNodeInfo(node1_id.Binary());
auto gcs_node1_info = Mocker::GenNodeInfo();
ClientID node1_id = ClientID::FromBinary(gcs_node1_info->node_id());
// Subscribe to node addition and removal events from GCS.
std::atomic<int> register_count(0);
@@ -647,16 +580,16 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeInfo) {
ASSERT_TRUE(SubscribeToNodeChange(on_subscribe));
// Register local node to GCS.
ASSERT_TRUE(RegisterSelf(gcs_node1_info));
ASSERT_TRUE(RegisterSelf(*gcs_node1_info));
sleep(1);
EXPECT_EQ(gcs_client_->Nodes().GetSelfId(), node1_id);
EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().node_id(), gcs_node1_info.node_id());
EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().state(), gcs_node1_info.state());
EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().node_id(), gcs_node1_info->node_id());
EXPECT_EQ(gcs_client_->Nodes().GetSelfInfo().state(), gcs_node1_info->state());
// Register a node to GCS.
ClientID node2_id = ClientID::FromRandom();
auto gcs_node2_info = GenGcsNodeInfo(node2_id.Binary());
ASSERT_TRUE(RegisterNode(gcs_node2_info));
auto gcs_node2_info = Mocker::GenNodeInfo();
ClientID node2_id = ClientID::FromBinary(gcs_node2_info->node_id());
ASSERT_TRUE(RegisterNode(*gcs_node2_info));
WaitPendingDone(register_count, 2);
// Get information of all nodes from GCS.
@@ -734,7 +667,7 @@ TEST_F(ServiceBasedGcsClientTest, TestNodeHeartbeat) {
TEST_F(ServiceBasedGcsClientTest, TestTaskInfo) {
JobID job_id = JobID::FromInt(1);
TaskID task_id = TaskID::ForDriverTask(job_id);
auto task_table_data = GenTaskTableData(job_id.Binary(), task_id.Binary());
auto task_table_data = Mocker::GenTaskTableData(job_id.Binary(), task_id.Binary());
// Subscribe to the event that the given task is added in GCS.
std::atomic<int> task_count(0);
@@ -773,7 +706,7 @@ TEST_F(ServiceBasedGcsClientTest, TestTaskInfo) {
// Add a task lease to GCS.
ClientID node_id = ClientID::FromRandom();
auto task_lease = GenTaskLeaseData(task_id.Binary(), node_id.Binary());
auto task_lease = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary());
ASSERT_TRUE(AddTaskLease(task_lease));
WaitPendingDone(task_lease_count, 2);
@@ -842,7 +775,7 @@ TEST_F(ServiceBasedGcsClientTest, TestObjectInfo) {
TEST_F(ServiceBasedGcsClientTest, TestStats) {
// Add profile data to GCS.
ClientID node_id = ClientID::FromRandom();
auto profile_table_data = GenProfileTableData(node_id);
auto profile_table_data = Mocker::GenProfileTableData(node_id);
ASSERT_TRUE(AddProfileData(profile_table_data));
}
@@ -856,7 +789,7 @@ TEST_F(ServiceBasedGcsClientTest, TestWorkerInfo) {
ASSERT_TRUE(SubscribeToWorkerFailures(on_subscribe));
// Report a worker failure to GCS.
auto worker_failure_data = GenWorkerFailureData();
auto worker_failure_data = Mocker::GenWorkerFailureData();
ASSERT_TRUE(ReportWorkerFailure(worker_failure_data));
WaitPendingDone(worker_failure_count, 1);
}
@@ -864,14 +797,14 @@ TEST_F(ServiceBasedGcsClientTest, TestWorkerInfo) {
TEST_F(ServiceBasedGcsClientTest, TestErrorInfo) {
// Report a job error to GCS.
JobID job_id = JobID::FromInt(1);
auto error_table_data = GenErrorTableData(job_id);
auto error_table_data = Mocker::GenErrorTableData(job_id);
ASSERT_TRUE(ReportJobError(error_table_data));
}
TEST_F(ServiceBasedGcsClientTest, TestDetectGcsAvailability) {
// Create job table data.
JobID add_job_id = JobID::FromInt(1);
auto job_table_data = GenJobTableData(add_job_id);
auto job_table_data = Mocker::GenJobTableData(add_job_id);
RAY_LOG(INFO) << "Initializing GCS service, port = " << gcs_server_->GetPort();
gcs_server_->Stop();
src/ray/gcs/gcs_server/gcs_table_storage.cc
+108
View File
@@ -0,0 +1,108 @@
// Copyright 2017 The Ray Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "gcs_table_storage.h"
#include "ray/common/id.h"
#include "ray/common/status.h"
#include "ray/gcs/callback.h"
namespace ray {
namespace gcs {
template <typename Key, typename Data>
Status GcsTable<Key, Data>::Put(const Key &key, const Data &value,
const StatusCallback &callback) {
return store_client_->AsyncPut(table_name_, key, value, callback);
}
template <typename Key, typename Data>
Status GcsTable<Key, Data>::Get(const Key &key,
const OptionalItemCallback<Data> &callback) {
return store_client_->AsyncGet(table_name_, key, callback);
}
template <typename Key, typename Data>
Status GcsTable<Key, Data>::GetAll(
const SegmentedCallback<std::pair<Key, Data>> &callback) {
return store_client_->AsyncGetAll(table_name_, callback);
}
template <typename Key, typename Data>
Status GcsTable<Key, Data>::Delete(const Key &key, const StatusCallback &callback) {
return store_client_->AsyncDelete(table_name_, key, callback);
}
template <typename Key, typename Data>
Status GcsTable<Key, Data>::BatchDelete(const std::vector<Key> &keys,
const StatusCallback &callback) {
// TODO(ffbin): We will use redis store client batch delete interface directly later.
auto finished_count = std::make_shared<int>(0);
int size = keys.size();
for (Key key : keys) {
auto done = [finished_count, size, callback](Status status) {
++(*finished_count);
if (*finished_count == size) {
callback(Status::OK());
}
};
RAY_CHECK_OK(store_client_->AsyncDelete(table_name_, key, done));
}
return Status::OK();
}
template <typename Key, typename Data>
Status GcsTableWithJobId<Key, Data>::Put(const Key &key, const Data &value,
const StatusCallback &callback) {
return this->store_client_->AsyncPutWithIndex(this->table_name_, key,
GetJobIdFromKey(key), value, callback);
}
template <typename Key, typename Data>
Status GcsTableWithJobId<Key, Data>::GetByJobId(
const JobID &job_id, const SegmentedCallback<std::pair<Key, Data>> &callback) {
// TODO(ffbin): We will add this function after redis store client support
// AsyncGetByIndex interface.
return Status::NotImplemented("GetByJobId not implemented");
}
template <typename Key, typename Data>
Status GcsTableWithJobId<Key, Data>::DeleteByJobId(const JobID &job_id,
const StatusCallback &callback) {
return this->store_client_->AsyncDeleteByIndex(this->table_name_, job_id, callback);
}
template class GcsTable<JobID, JobTableData>;
template class GcsTable<ClientID, GcsNodeInfo>;
template class GcsTable<ClientID, ResourceMap>;
template class GcsTable<ClientID, HeartbeatTableData>;
template class GcsTable<ClientID, HeartbeatBatchTableData>;
template class GcsTable<JobID, ErrorTableData>;
template class GcsTable<UniqueID, ProfileTableData>;
template class GcsTable<WorkerID, WorkerFailureData>;
template class GcsTable<ActorID, ActorTableData>;
template class GcsTable<ActorCheckpointID, ActorCheckpointData>;
template class GcsTable<ActorID, ActorCheckpointIdData>;
template class GcsTable<TaskID, TaskTableData>;
template class GcsTable<TaskID, TaskLeaseData>;
template class GcsTable<TaskID, TaskReconstructionData>;
template class GcsTable<ObjectID, ObjectTableDataList>;
template class GcsTableWithJobId<ActorID, ActorTableData>;
template class GcsTableWithJobId<ActorID, ActorCheckpointIdData>;
template class GcsTableWithJobId<TaskID, TaskTableData>;
template class GcsTableWithJobId<TaskID, TaskLeaseData>;
template class GcsTableWithJobId<TaskID, TaskReconstructionData>;
template class GcsTableWithJobId<ObjectID, ObjectTableDataList>;
} // namespace gcs
} // namespace ray
src/ray/gcs/gcs_server/gcs_table_storage.h
+444
View File
@@ -0,0 +1,444 @@
// Copyright 2017 The Ray Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef RAY_GCS_GCS_TABLE_STORAGE_H_
#define RAY_GCS_GCS_TABLE_STORAGE_H_
#include <utility>
#include "ray/gcs/store_client/redis_store_client.h"
#include "ray/protobuf/gcs.pb.h"
namespace ray {
namespace gcs {
using rpc::ActorCheckpointData;
using rpc::ActorCheckpointIdData;
using rpc::ActorTableData;
using rpc::ErrorTableData;
using rpc::GcsNodeInfo;
using rpc::HeartbeatBatchTableData;
using rpc::HeartbeatTableData;
using rpc::JobTableData;
using rpc::ObjectTableData;
using rpc::ObjectTableDataList;
using rpc::ProfileTableData;
using rpc::ResourceMap;
using rpc::ResourceTableData;
using rpc::TaskLeaseData;
using rpc::TaskReconstructionData;
using rpc::TaskTableData;
using rpc::WorkerFailureData;
/// \class GcsTable
///
/// GcsTable is the storage interface for all GCS tables whose data do not belong to
/// specific jobs. This class is not meant to be used directly. All gcs table classes
/// without job id should derive from this class and override the table_name_ member with
/// a unique value for that table.
template <typename Key, typename Data>
class GcsTable {
public:
explicit GcsTable(std::shared_ptr<StoreClient<Key, Data, JobID>> store_client)
: store_client_(store_client) {}
virtual ~GcsTable() = default;
/// Write data to the table asynchronously.
///
/// \param key The key that will be written to the table.
/// \param value The value of the key that will be written to the table.
/// \param callback Callback that will be called after write finishes.
/// \return Status
virtual Status Put(const Key &key, const Data &value, const StatusCallback &callback);
/// Get data from the table asynchronously.
///
/// \param key The key to lookup from the table.
/// \param callback Callback that will be called after read finishes.
/// \return Status
Status Get(const Key &key, const OptionalItemCallback<Data> &callback);
/// Get all data from the table asynchronously.
///
/// \param callback Callback that will be called after data has been received.
/// If the callback return `has_more == true` mean there's more data to be received.
/// \return Status
Status GetAll(const SegmentedCallback<std::pair<Key, Data>> &callback);
/// Delete data from the table asynchronously.
///
/// \param key The key that will be deleted from the table.
/// \param callback Callback that will be called after delete finishes.
/// \return Status
Status Delete(const Key &key, const StatusCallback &callback);
/// Delete a batch of data from the table asynchronously.
///
/// \param keys The batch key that will be deleted from the table.
/// \param callback Callback that will be called after delete finishes.
/// \return Status
Status BatchDelete(const std::vector<Key> &keys, const StatusCallback &callback);
protected:
std::string table_name_;
std::shared_ptr<StoreClient<Key, Data, JobID>> store_client_;
};
/// \class GcsTableWithJobId
///
/// GcsTableWithJobId is the storage interface for all GCS tables whose data belongs to
/// specific jobs. This class is not meant to be used directly. All gcs table classes with
/// job id should derive from this class and override the table_name_ member with a unique
/// value for that table.
template <typename Key, typename Data>
class GcsTableWithJobId : public GcsTable<Key, Data> {
public:
explicit GcsTableWithJobId(std::shared_ptr<StoreClient<Key, Data, JobID>> store_client)
: GcsTable<Key, Data>(store_client) {}
/// Write data to the table asynchronously.
///
/// \param key The key that will be written to the table. The job id can be obtained
/// from the key. \param value The value of the key that will be written to the table.
/// \param callback Callback that will be called after write finishes.
/// \return Status
Status Put(const Key &key, const Data &value, const StatusCallback &callback) override;
/// Get all the data of the specified job id from the table asynchronously.
///
/// \param job_id The key to lookup from the table.
/// \param callback Callback that will be called after read finishes.
/// \return Status
Status GetByJobId(const JobID &job_id,
const SegmentedCallback<std::pair<Key, Data>> &callback);
/// Delete all the data of the specified job id from the table asynchronously.
///
/// \param job_id The key that will be deleted from the table.
/// \param callback Callback that will be called after delete finishes.
/// \return Status
Status DeleteByJobId(const JobID &job_id, const StatusCallback &callback);
protected:
virtual JobID GetJobIdFromKey(const Key &key) = 0;
};
class GcsJobTable : public GcsTable<JobID, JobTableData> {
public:
explicit GcsJobTable(
std::shared_ptr<StoreClient<JobID, JobTableData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::JOB);
}
};
class GcsActorTable : public GcsTableWithJobId<ActorID, ActorTableData> {
public:
explicit GcsActorTable(
std::shared_ptr<StoreClient<ActorID, ActorTableData, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::ACTOR);
}
private:
JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); }
};
class GcsActorCheckpointTable : public GcsTable<ActorCheckpointID, ActorCheckpointData> {
public:
explicit GcsActorCheckpointTable(
std::shared_ptr<StoreClient<ActorCheckpointID, ActorCheckpointData, JobID>>
store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT);
}
};
class GcsActorCheckpointIdTable
: public GcsTableWithJobId<ActorID, ActorCheckpointIdData> {
public:
explicit GcsActorCheckpointIdTable(
std::shared_ptr<StoreClient<ActorID, ActorCheckpointIdData, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::ACTOR_CHECKPOINT_ID);
}
private:
JobID GetJobIdFromKey(const ActorID &key) { return key.JobId(); }
};
class GcsTaskTable : public GcsTableWithJobId<TaskID, TaskTableData> {
public:
explicit GcsTaskTable(
std::shared_ptr<StoreClient<TaskID, TaskTableData, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::TASK);
}
private:
JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); }
};
class GcsTaskLeaseTable : public GcsTableWithJobId<TaskID, TaskLeaseData> {
public:
explicit GcsTaskLeaseTable(
std::shared_ptr<StoreClient<TaskID, TaskLeaseData, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::TASK_LEASE);
}
private:
JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); }
};
class GcsTaskReconstructionTable
: public GcsTableWithJobId<TaskID, TaskReconstructionData> {
public:
explicit GcsTaskReconstructionTable(
std::shared_ptr<StoreClient<TaskID, TaskReconstructionData, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::TASK_RECONSTRUCTION);
}
private:
JobID GetJobIdFromKey(const TaskID &key) { return key.ActorId().JobId(); }
};
class GcsObjectTable : public GcsTableWithJobId<ObjectID, ObjectTableDataList> {
public:
explicit GcsObjectTable(
std::shared_ptr<StoreClient<ObjectID, ObjectTableDataList, JobID>> store_client)
: GcsTableWithJobId(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::OBJECT);
}
private:
JobID GetJobIdFromKey(const ObjectID &key) { return key.TaskId().JobId(); }
};
class GcsNodeTable : public GcsTable<ClientID, GcsNodeInfo> {
public:
explicit GcsNodeTable(
std::shared_ptr<StoreClient<ClientID, GcsNodeInfo, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::CLIENT);
}
};
class GcsNodeResourceTable : public GcsTable<ClientID, ResourceMap> {
public:
explicit GcsNodeResourceTable(
std::shared_ptr<StoreClient<ClientID, ResourceMap, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::NODE_RESOURCE);
}
};
class GcsHeartbeatTable : public GcsTable<ClientID, HeartbeatTableData> {
public:
explicit GcsHeartbeatTable(
std::shared_ptr<StoreClient<ClientID, HeartbeatTableData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::HEARTBEAT);
}
};
class GcsHeartbeatBatchTable : public GcsTable<ClientID, HeartbeatBatchTableData> {
public:
explicit GcsHeartbeatBatchTable(
std::shared_ptr<StoreClient<ClientID, HeartbeatBatchTableData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::HEARTBEAT_BATCH);
}
};
class GcsErrorInfoTable : public GcsTable<JobID, ErrorTableData> {
public:
explicit GcsErrorInfoTable(
std::shared_ptr<StoreClient<JobID, ErrorTableData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::ERROR_INFO);
}
};
class GcsProfileTable : public GcsTable<UniqueID, ProfileTableData> {
public:
explicit GcsProfileTable(
std::shared_ptr<StoreClient<UniqueID, ProfileTableData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::PROFILE);
}
};
class GcsWorkerFailureTable : public GcsTable<WorkerID, WorkerFailureData> {
public:
explicit GcsWorkerFailureTable(
std::shared_ptr<StoreClient<WorkerID, WorkerFailureData, JobID>> store_client)
: GcsTable(std::move(store_client)) {
table_name_ = TablePrefix_Name(TablePrefix::WORKER_FAILURE);
}
};
/// \class GcsTableStorage
///
/// This class is not meant to be used directly. All gcs table storage classes should
/// derive from this class and override class member variables.
class GcsTableStorage {
public:
GcsJobTable &JobTable() {
RAY_CHECK(job_table_ != nullptr);
return *job_table_;
}
GcsActorTable &ActorTable() {
RAY_CHECK(actor_table_ != nullptr);
return *actor_table_;
}
GcsActorCheckpointTable &ActorCheckpointTable() {
RAY_CHECK(actor_checkpoint_table_ != nullptr);
return *actor_checkpoint_table_;
}
GcsActorCheckpointIdTable &ActorCheckpointIdTable() {
RAY_CHECK(actor_checkpoint_id_table_ != nullptr);
return *actor_checkpoint_id_table_;
}
GcsTaskTable &TaskTable() {
RAY_CHECK(task_table_ != nullptr);
return *task_table_;
}
GcsTaskLeaseTable &TaskLeaseTable() {
RAY_CHECK(task_lease_table_ != nullptr);
return *task_lease_table_;
}
GcsTaskReconstructionTable &TaskReconstructionTable() {
RAY_CHECK(task_reconstruction_table_ != nullptr);
return *task_reconstruction_table_;
}
GcsObjectTable &ObjectTable() {
RAY_CHECK(object_table_ != nullptr);
return *object_table_;
}
GcsNodeTable &NodeTable() {
RAY_CHECK(node_table_ != nullptr);
return *node_table_;
}
GcsNodeResourceTable &NodeResourceTable() {
RAY_CHECK(node_resource_table_ != nullptr);
return *node_resource_table_;
}
GcsHeartbeatTable &HeartbeatTable() {
RAY_CHECK(heartbeat_table_ != nullptr);
return *heartbeat_table_;
}
GcsHeartbeatBatchTable &HeartbeatBatchTable() {
RAY_CHECK(heartbeat_batch_table_ != nullptr);
return *heartbeat_batch_table_;
}
GcsErrorInfoTable &ErrorInfoTable() {
RAY_CHECK(error_info_table_ != nullptr);
return *error_info_table_;
}
GcsProfileTable &ProfileTable() {
RAY_CHECK(profile_table_ != nullptr);
return *profile_table_;
}
GcsWorkerFailureTable &WorkerFailureTable() {
RAY_CHECK(worker_failure_table_ != nullptr);
return *worker_failure_table_;
}
protected:
std::unique_ptr<GcsJobTable> job_table_;
std::unique_ptr<GcsActorTable> actor_table_;
std::unique_ptr<GcsActorCheckpointTable> actor_checkpoint_table_;
std::unique_ptr<GcsActorCheckpointIdTable> actor_checkpoint_id_table_;
std::unique_ptr<GcsTaskTable> task_table_;
std::unique_ptr<GcsTaskLeaseTable> task_lease_table_;
std::unique_ptr<GcsTaskReconstructionTable> task_reconstruction_table_;
std::unique_ptr<GcsObjectTable> object_table_;
std::unique_ptr<GcsNodeTable> node_table_;
std::unique_ptr<GcsNodeResourceTable> node_resource_table_;
std::unique_ptr<GcsHeartbeatTable> heartbeat_table_;
std::unique_ptr<GcsHeartbeatBatchTable> heartbeat_batch_table_;
std::unique_ptr<GcsErrorInfoTable> error_info_table_;
std::unique_ptr<GcsProfileTable> profile_table_;
std::unique_ptr<GcsWorkerFailureTable> worker_failure_table_;
};
/// \class RedisGcsTableStorage
/// RedisGcsTableStorage is an implementation of `GcsTableStorage`
/// that uses redis as storage.
class RedisGcsTableStorage : public GcsTableStorage {
public:
explicit RedisGcsTableStorage(std::shared_ptr<RedisClient> redis_client) {
job_table_.reset(new GcsJobTable(
std::make_shared<RedisStoreClient<JobID, JobTableData, JobID>>(redis_client)));
actor_table_.reset(new GcsActorTable(
std::make_shared<RedisStoreClient<ActorID, ActorTableData, JobID>>(
redis_client)));
actor_checkpoint_table_.reset(new GcsActorCheckpointTable(
std::make_shared<RedisStoreClient<ActorCheckpointID, ActorCheckpointData, JobID>>(
redis_client)));
actor_checkpoint_id_table_.reset(new GcsActorCheckpointIdTable(
std::make_shared<RedisStoreClient<ActorID, ActorCheckpointIdData, JobID>>(
redis_client)));
task_table_.reset(new GcsTaskTable(
std::make_shared<RedisStoreClient<TaskID, TaskTableData, JobID>>(redis_client)));
task_lease_table_.reset(new GcsTaskLeaseTable(
std::make_shared<RedisStoreClient<TaskID, TaskLeaseData, JobID>>(redis_client)));
task_reconstruction_table_.reset(new GcsTaskReconstructionTable(
std::make_shared<RedisStoreClient<TaskID, TaskReconstructionData, JobID>>(
redis_client)));
object_table_.reset(new GcsObjectTable(
std::make_shared<RedisStoreClient<ObjectID, ObjectTableDataList, JobID>>(
redis_client)));
node_table_.reset(new GcsNodeTable(
std::make_shared<RedisStoreClient<ClientID, GcsNodeInfo, JobID>>(redis_client)));
node_resource_table_.reset(new GcsNodeResourceTable(
std::make_shared<RedisStoreClient<ClientID, ResourceMap, JobID>>(redis_client)));
heartbeat_table_.reset(new GcsHeartbeatTable(
std::make_shared<RedisStoreClient<ClientID, HeartbeatTableData, JobID>>(
redis_client)));
heartbeat_batch_table_.reset(new GcsHeartbeatBatchTable(
std::make_shared<RedisStoreClient<ClientID, HeartbeatBatchTableData, JobID>>(
redis_client)));
error_info_table_.reset(new GcsErrorInfoTable(
std::make_shared<RedisStoreClient<JobID, ErrorTableData, JobID>>(redis_client)));
profile_table_.reset(new GcsProfileTable(
std::make_shared<RedisStoreClient<UniqueID, ProfileTableData, JobID>>(
redis_client)));
worker_failure_table_.reset(new GcsWorkerFailureTable(
std::make_shared<RedisStoreClient<WorkerID, WorkerFailureData, JobID>>(
redis_client)));
}
};
} // namespace gcs
} // namespace ray
#endif // RAY_GCS_GCS_TABLE_STORAGE_H_
src/ray/gcs/gcs_server/test/gcs_actor_manager_test.cc
+1 -1
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ray/gcs/gcs_server/test/gcs_test_util.h>
#include <ray/gcs/test/gcs_test_util.h>
#include <memory>
#include "gtest/gtest.h"
src/ray/gcs/gcs_server/test/gcs_actor_scheduler_test.cc
+1 -1
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ray/gcs/gcs_server/test/gcs_test_util.h>
#include <ray/gcs/test/gcs_test_util.h>
#include <memory>
#include "gmock/gmock.h"
src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc
+1 -1
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <ray/gcs/gcs_server/test/gcs_test_util.h>
#include <ray/gcs/test/gcs_test_util.h>
#include <memory>
#include "gtest/gtest.h"
src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc
+28 -80
View File
@@ -15,6 +15,7 @@
#include "gtest/gtest.h"
#include "ray/common/test_util.h"
#include "ray/gcs/gcs_server/gcs_server.h"
#include "ray/gcs/test/gcs_test_util.h"
#include "ray/rpc/gcs_server/gcs_rpc_client.h"
namespace ray {
@@ -382,56 +383,6 @@ class GcsServerTest : public RedisServiceManagerForTest {
return status == std::future_status::ready;
}
rpc::JobTableData GenJobTableData(JobID job_id) {
rpc::JobTableData job_table_data;
job_table_data.set_job_id(job_id.Binary());
job_table_data.set_is_dead(false);
job_table_data.set_timestamp(std::time(nullptr));
job_table_data.set_node_manager_address("127.0.0.1");
job_table_data.set_driver_pid(5667L);
return job_table_data;
}
rpc::ActorTableData GenActorTableData(const ActorID &actor_id) {
rpc::ActorTableData actor_table_data;
actor_table_data.set_actor_id(actor_id.Binary());
actor_table_data.set_job_id(actor_id.JobId().Binary());
actor_table_data.set_state(
rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE);
actor_table_data.set_max_reconstructions(1);
actor_table_data.set_remaining_reconstructions(1);
return actor_table_data;
}
rpc::GcsNodeInfo GenGcsNodeInfo(const std::string &node_id) {
rpc::GcsNodeInfo gcs_node_info;
gcs_node_info.set_node_id(node_id);
gcs_node_info.set_node_manager_address("127.0.0.1");
gcs_node_info.set_node_manager_port(6000);
gcs_node_info.set_state(rpc::GcsNodeInfo_GcsNodeState_ALIVE);
return gcs_node_info;
}
rpc::TaskTableData GenTaskTableData(const std::string &job_id,
const std::string &task_id) {
rpc::TaskTableData task_table_data;
rpc::Task task;
rpc::TaskSpec task_spec;
task_spec.set_job_id(job_id);
task_spec.set_task_id(task_id);
task.mutable_task_spec()->CopyFrom(task_spec);
task_table_data.mutable_task()->CopyFrom(task);
return task_table_data;
}
rpc::TaskLeaseData GenTaskLeaseData(const std::string &task_id,
const std::string &node_id) {
rpc::TaskLeaseData task_lease_data;
task_lease_data.set_task_id(task_id);
task_lease_data.set_node_manager_id(node_id);
return task_lease_data;
}
protected:
// Gcs server
std::unique_ptr<gcs::GcsServer> gcs_server_;
@@ -450,32 +401,31 @@ class GcsServerTest : public RedisServiceManagerForTest {
TEST_F(GcsServerTest, TestActorInfo) {
// Create actor_table_data
JobID job_id = JobID::FromInt(1);
ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0);
rpc::ActorTableData actor_table_data = GenActorTableData(actor_id);
auto actor_table_data = Mocker::GenActorTableData(job_id);
// Register actor
rpc::RegisterActorInfoRequest register_actor_info_request;
register_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data);
register_actor_info_request.mutable_actor_table_data()->CopyFrom(*actor_table_data);
ASSERT_TRUE(RegisterActorInfo(register_actor_info_request));
rpc::ActorTableData result = GetActorInfo(actor_table_data.actor_id());
rpc::ActorTableData result = GetActorInfo(actor_table_data->actor_id());
ASSERT_TRUE(result.state() ==
rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE);
// Update actor state
rpc::UpdateActorInfoRequest update_actor_info_request;
actor_table_data.set_state(
actor_table_data->set_state(
rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD);
update_actor_info_request.set_actor_id(actor_table_data.actor_id());
update_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data);
update_actor_info_request.set_actor_id(actor_table_data->actor_id());
update_actor_info_request.mutable_actor_table_data()->CopyFrom(*actor_table_data);
ASSERT_TRUE(UpdateActorInfo(update_actor_info_request));
result = GetActorInfo(actor_table_data.actor_id());
result = GetActorInfo(actor_table_data->actor_id());
ASSERT_TRUE(result.state() ==
rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD);
// Add actor checkpoint
ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom();
rpc::ActorCheckpointData checkpoint;
checkpoint.set_actor_id(actor_table_data.actor_id());
checkpoint.set_actor_id(actor_table_data->actor_id());
checkpoint.set_checkpoint_id(checkpoint_id.Binary());
checkpoint.set_execution_dependency(checkpoint_id.Binary());
@@ -483,39 +433,38 @@ TEST_F(GcsServerTest, TestActorInfo) {
add_actor_checkpoint_request.mutable_checkpoint_data()->CopyFrom(checkpoint);
ASSERT_TRUE(AddActorCheckpoint(add_actor_checkpoint_request));
rpc::ActorCheckpointData checkpoint_result =
GetActorCheckpoint(actor_id.Binary(), checkpoint_id.Binary());
ASSERT_TRUE(checkpoint_result.actor_id() == actor_table_data.actor_id());
GetActorCheckpoint(actor_table_data->actor_id(), checkpoint_id.Binary());
ASSERT_TRUE(checkpoint_result.actor_id() == actor_table_data->actor_id());
ASSERT_TRUE(checkpoint_result.checkpoint_id() == checkpoint_id.Binary());
rpc::ActorCheckpointIdData checkpoint_id_result =
GetActorCheckpointID(actor_table_data.actor_id());
ASSERT_TRUE(checkpoint_id_result.actor_id() == actor_table_data.actor_id());
GetActorCheckpointID(actor_table_data->actor_id());
ASSERT_TRUE(checkpoint_id_result.actor_id() == actor_table_data->actor_id());
ASSERT_TRUE(checkpoint_id_result.checkpoint_ids_size() == 1);
}
TEST_F(GcsServerTest, TestJobInfo) {
// Create job_table_data
JobID job_id = JobID::FromInt(1);
rpc::JobTableData job_table_data = GenJobTableData(job_id);
auto job_table_data = Mocker::GenJobTableData(job_id);
// Add job
rpc::AddJobRequest add_job_request;
add_job_request.mutable_data()->CopyFrom(job_table_data);
add_job_request.mutable_data()->CopyFrom(*job_table_data);
ASSERT_TRUE(AddJob(add_job_request));
// Mark job finished
rpc::MarkJobFinishedRequest mark_job_finished_request;
mark_job_finished_request.set_job_id(job_table_data.job_id());
mark_job_finished_request.set_job_id(job_table_data->job_id());
ASSERT_TRUE(MarkJobFinished(mark_job_finished_request));
}
TEST_F(GcsServerTest, TestNodeInfo) {
// Create gcs node info
ClientID node_id = ClientID::FromRandom();
rpc::GcsNodeInfo gcs_node_info = GenGcsNodeInfo(node_id.Binary());
auto gcs_node_info = Mocker::GenNodeInfo();
// Register node info
rpc::RegisterNodeRequest register_node_info_request;
register_node_info_request.mutable_node_info()->CopyFrom(gcs_node_info);
register_node_info_request.mutable_node_info()->CopyFrom(*gcs_node_info);
ASSERT_TRUE(RegisterNode(register_node_info_request));
std::vector<rpc::GcsNodeInfo> node_info_list = GetAllNodeInfo();
ASSERT_TRUE(node_info_list.size() == 1);
@@ -524,12 +473,12 @@ TEST_F(GcsServerTest, TestNodeInfo) {
// Report heartbeat
rpc::ReportHeartbeatRequest report_heartbeat_request;
report_heartbeat_request.mutable_heartbeat()->set_client_id(node_id.Binary());
report_heartbeat_request.mutable_heartbeat()->set_client_id(gcs_node_info->node_id());
ASSERT_TRUE(ReportHeartbeat(report_heartbeat_request));
// Unregister node info
rpc::UnregisterNodeRequest unregister_node_info_request;
unregister_node_info_request.set_node_id(node_id.Binary());
unregister_node_info_request.set_node_id(gcs_node_info->node_id());
ASSERT_TRUE(UnregisterNode(unregister_node_info_request));
node_info_list = GetAllNodeInfo();
ASSERT_TRUE(node_info_list.size() == 1);
@@ -538,21 +487,21 @@ TEST_F(GcsServerTest, TestNodeInfo) {
// Update node resources
rpc::UpdateResourcesRequest update_resources_request;
update_resources_request.set_node_id(node_id.Binary());
update_resources_request.set_node_id(gcs_node_info->node_id());
rpc::ResourceTableData resource_table_data;
resource_table_data.set_resource_capacity(1.0);
std::string resource_name = "CPU";
(*update_resources_request.mutable_resources())[resource_name] = resource_table_data;
ASSERT_TRUE(UpdateResources(update_resources_request));
auto resources = GetResources(node_id.Binary());
auto resources = GetResources(gcs_node_info->node_id());
ASSERT_TRUE(resources.size() == 1);
// Delete node resources
rpc::DeleteResourcesRequest delete_resources_request;
delete_resources_request.set_node_id(node_id.Binary());
delete_resources_request.set_node_id(gcs_node_info->node_id());
delete_resources_request.add_resource_name_list(resource_name);
ASSERT_TRUE(DeleteResources(delete_resources_request));
resources = GetResources(node_id.Binary());
resources = GetResources(gcs_node_info->node_id());
ASSERT_TRUE(resources.empty());
}
@@ -591,11 +540,11 @@ TEST_F(GcsServerTest, TestTaskInfo) {
// Create task_table_data
JobID job_id = JobID::FromInt(1);
TaskID task_id = TaskID::ForDriverTask(job_id);
rpc::TaskTableData job_table_data = GenTaskTableData(job_id.Binary(), task_id.Binary());
auto job_table_data = Mocker::GenTaskTableData(job_id.Binary(), task_id.Binary());
// Add task
rpc::AddTaskRequest add_task_request;
add_task_request.mutable_task_data()->CopyFrom(job_table_data);
add_task_request.mutable_task_data()->CopyFrom(*job_table_data);
ASSERT_TRUE(AddTask(add_task_request));
rpc::TaskTableData result = GetTask(task_id.Binary());
ASSERT_TRUE(result.task().task_spec().job_id() == job_id.Binary());
@@ -609,10 +558,9 @@ TEST_F(GcsServerTest, TestTaskInfo) {
// Add task lease
ClientID node_id = ClientID::FromRandom();
rpc::TaskLeaseData task_lease_data =
GenTaskLeaseData(task_id.Binary(), node_id.Binary());
auto task_lease_data = Mocker::GenTaskLeaseData(task_id.Binary(), node_id.Binary());
rpc::AddTaskLeaseRequest add_task_lease_request;
add_task_lease_request.mutable_task_lease_data()->CopyFrom(task_lease_data);
add_task_lease_request.mutable_task_lease_data()->CopyFrom(*task_lease_data);
ASSERT_TRUE(AddTaskLease(add_task_lease_request));
// Attempt task reconstruction
src/ray/gcs/gcs_server/test/redis_gcs_table_storage_test.cc
+129
View File
@@ -0,0 +1,129 @@
// Copyright 2017 The Ray Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "gtest/gtest.h"
#include "ray/common/test_util.h"
#include "ray/gcs/gcs_server/gcs_table_storage.h"
#include "ray/gcs/store_client/redis_store_client.h"
#include "ray/gcs/store_client/test/store_client_test_base.h"
#include "ray/gcs/test/gcs_test_util.h"
namespace ray {
class GcsTableStorageTest : public gcs::StoreClientTestBase {
public:
GcsTableStorageTest() {}
virtual ~GcsTableStorageTest() {}
void InitStoreClient() override {
gcs::RedisClientOptions options("127.0.0.1", REDIS_SERVER_PORT, "", true);
redis_client_ = std::make_shared<gcs::RedisClient>(options);
RAY_CHECK_OK(redis_client_->Connect(io_service_pool_->GetAll()));
gcs_table_storage_ = std::make_shared<gcs::RedisGcsTableStorage>(redis_client_);
}
void DisconnectStoreClient() override { redis_client_->Disconnect(); }
protected:
template <typename TABLE, typename KEY, typename VALUE>
void Put(TABLE &table, const KEY &key, const VALUE &value) {
auto on_done = [this](Status status) { --pending_count_; };
++pending_count_;
RAY_CHECK_OK(table.Put(key, value, on_done));
WaitPendingDone();
}
template <typename TABLE, typename KEY, typename VALUE>
int Get(TABLE &table, const KEY &key, std::vector<VALUE> &values) {
auto on_done = [this, &values](Status status, const boost::optional<VALUE> &result) {
RAY_CHECK_OK(status);
--pending_count_;
values.clear();
if (result) {
values.push_back(*result);
}
};
++pending_count_;
RAY_CHECK_OK(table.Get(key, on_done));
WaitPendingDone();
return values.size();
}
template <typename TABLE, typename KEY>
void Delete(TABLE &table, const KEY &key) {
auto on_done = [this](Status status) {
RAY_CHECK_OK(status);
--pending_count_;
};
++pending_count_;
RAY_CHECK_OK(table.Delete(key, on_done));
WaitPendingDone();
}
std::shared_ptr<gcs::RedisClient> redis_client_;
std::shared_ptr<gcs::GcsTableStorage> gcs_table_storage_;
};
TEST_F(GcsTableStorageTest, TestGcsTableApi) {
auto table = gcs_table_storage_->JobTable();
JobID job1_id = JobID::FromInt(1);
JobID job2_id = JobID::FromInt(2);
auto job1_table_data = Mocker::GenJobTableData(job1_id);
auto job2_table_data = Mocker::GenJobTableData(job2_id);
// Put.
Put(table, job1_id, *job1_table_data);
Put(table, job2_id, *job2_table_data);
// Get.
std::vector<rpc::JobTableData> values;
ASSERT_EQ(Get(table, job2_id, values), 1);
ASSERT_EQ(Get(table, job2_id, values), 1);
// Delete.
Delete(table, job1_id);
ASSERT_EQ(Get(table, job1_id, values), 0);
ASSERT_EQ(Get(table, job2_id, values), 1);
}
TEST_F(GcsTableStorageTest, TestGcsTableWithJobIdApi) {
auto table = gcs_table_storage_->ActorTable();
JobID job_id = JobID::FromInt(3);
auto actor_table_data = Mocker::GenActorTableData(job_id);
ActorID actor_id = ActorID::FromBinary(actor_table_data->actor_id());
// Put.
Put(table, actor_id, *actor_table_data);
// Get.
std::vector<rpc::ActorTableData> values;
ASSERT_EQ(Get(table, actor_id, values), 1);
// Delete.
Delete(table, actor_id);
ASSERT_EQ(Get(table, actor_id, values), 0);
}
} // namespace ray
int main(int argc, char **argv) {
::testing::InitGoogleTest(&argc, argv);
RAY_CHECK(argc == 4);
ray::REDIS_SERVER_EXEC_PATH = argv[1];
ray::REDIS_CLIENT_EXEC_PATH = argv[2];
ray::REDIS_MODULE_LIBRARY_PATH = argv[3];
return RUN_ALL_TESTS();
}
src/ray/gcs/store_client/redis_store_client.cc
+14
View File
@@ -145,6 +145,20 @@ Status RedisStoreClient<Key, Data, IndexKey>::AsyncDeleteByIndex(
}
template class RedisStoreClient<ActorID, rpc::ActorTableData, JobID>;
template class RedisStoreClient<JobID, rpc::JobTableData, JobID>;
template class RedisStoreClient<ActorCheckpointID, rpc::ActorCheckpointData, JobID>;
template class RedisStoreClient<ActorID, rpc::ActorCheckpointIdData, JobID>;
template class RedisStoreClient<TaskID, rpc::TaskTableData, JobID>;
template class RedisStoreClient<TaskID, rpc::TaskLeaseData, JobID>;
template class RedisStoreClient<TaskID, rpc::TaskReconstructionData, JobID>;
template class RedisStoreClient<ObjectID, rpc::ObjectTableDataList, JobID>;
template class RedisStoreClient<ClientID, rpc::GcsNodeInfo, JobID>;
template class RedisStoreClient<ClientID, rpc::ResourceMap, JobID>;
template class RedisStoreClient<ClientID, rpc::HeartbeatTableData, JobID>;
template class RedisStoreClient<ClientID, rpc::HeartbeatBatchTableData, JobID>;
template class RedisStoreClient<JobID, rpc::ErrorTableData, JobID>;
template class RedisStoreClient<UniqueID, rpc::ProfileTableData, JobID>;
template class RedisStoreClient<WorkerID, rpc::WorkerFailureData, JobID>;
} // namespace gcs
src/ray/gcs/gcs_server/test/gcs_test_util.h → src/ray/gcs/test/gcs_test_util.h
+68 -7
View File
@@ -15,13 +15,13 @@
#ifndef RAY_GCS_TEST_UTIL_H
#define RAY_GCS_TEST_UTIL_H
#include <ray/common/task/task.h>
#include <ray/common/task/task_util.h>
#include <ray/common/test_util.h>
#include <ray/gcs/gcs_server/gcs_actor_manager.h>
#include <ray/gcs/gcs_server/gcs_actor_scheduler.h>
#include <ray/gcs/gcs_server/gcs_node_manager.h>
#include <ray/util/asio_util.h>
#include "src/ray/common/task/task.h"
#include "src/ray/common/task/task_util.h"
#include "src/ray/common/test_util.h"
#include "src/ray/gcs/gcs_server/gcs_actor_manager.h"
#include "src/ray/gcs/gcs_server/gcs_actor_scheduler.h"
#include "src/ray/gcs/gcs_server/gcs_node_manager.h"
#include "src/ray/util/asio_util.h"
#include <memory>
#include <utility>
@@ -59,6 +59,67 @@ struct Mocker {
return node;
}
static std::shared_ptr<rpc::JobTableData> GenJobTableData(JobID job_id) {
auto job_table_data = std::make_shared<rpc::JobTableData>();
job_table_data->set_job_id(job_id.Binary());
job_table_data->set_is_dead(false);
job_table_data->set_timestamp(std::time(nullptr));
job_table_data->set_node_manager_address("127.0.0.1");
job_table_data->set_driver_pid(5667L);
return job_table_data;
}
static std::shared_ptr<rpc::ActorTableData> GenActorTableData(const JobID &job_id) {
auto actor_table_data = std::make_shared<rpc::ActorTableData>();
ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0);
actor_table_data->set_actor_id(actor_id.Binary());
actor_table_data->set_job_id(job_id.Binary());
actor_table_data->set_state(
rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE);
actor_table_data->set_max_reconstructions(1);
actor_table_data->set_remaining_reconstructions(1);
return actor_table_data;
}
static std::shared_ptr<rpc::TaskTableData> GenTaskTableData(
const std::string &job_id, const std::string &task_id) {
auto task_table_data = std::make_shared<rpc::TaskTableData>();
rpc::Task task;
rpc::TaskSpec task_spec;
task_spec.set_job_id(job_id);
task_spec.set_task_id(task_id);
task.mutable_task_spec()->CopyFrom(task_spec);
task_table_data->mutable_task()->CopyFrom(task);
return task_table_data;
}
static std::shared_ptr<rpc::TaskLeaseData> GenTaskLeaseData(
const std::string &task_id, const std::string &node_id) {
auto task_lease_data = std::make_shared<rpc::TaskLeaseData>();
task_lease_data->set_task_id(task_id);
task_lease_data->set_node_manager_id(node_id);
return task_lease_data;
}
static std::shared_ptr<rpc::ProfileTableData> GenProfileTableData(
const ClientID &node_id) {
auto profile_table_data = std::make_shared<rpc::ProfileTableData>();
profile_table_data->set_component_id(node_id.Binary());
return profile_table_data;
}
static std::shared_ptr<rpc::ErrorTableData> GenErrorTableData(const JobID &job_id) {
auto error_table_data = std::make_shared<rpc::ErrorTableData>();
error_table_data->set_job_id(job_id.Binary());
return error_table_data;
}
static std::shared_ptr<rpc::WorkerFailureData> GenWorkerFailureData() {
auto worker_failure_data = std::make_shared<rpc::WorkerFailureData>();
worker_failure_data->set_timestamp(std::time(nullptr));
return worker_failure_data;
}
class MockWorkerClient : public rpc::CoreWorkerClientInterface {
public:
ray::Status PushNormalTask(
src/ray/protobuf/gcs.proto
+7
View File
@@ -303,6 +303,13 @@ message WorkerFailureData {
bool intentional_disconnect = 4;
}
message ResourceMap {
map<string, ResourceTableData> items = 1;
}
message ObjectTableDataList {
repeated ObjectTableData items = 1;
}
// This enum type is used as object's metadata to indicate the object's creating
// task has failed because of a certain error.
// TODO(hchen): We may want to make these errors more specific. E.g., we may want