From 3c080f4baa273f844f61c309f706a8c9bd6ae699 Mon Sep 17 00:00:00 2001 From: Melih Elibol Date: Thu, 15 Mar 2018 22:25:01 -0700 Subject: [PATCH] Add a callback for gcs table lookup failures. (#1702) * Add callback to gcs client for table lookup failures. * update plasma_manager reflecting changes to gcs callback. --- src/common/redis_module/ray_redis_module.cc | 3 +++ src/plasma/plasma_manager.cc | 3 +++ src/ray/gcs/client_test.cc | 24 +++++++++++++++--- src/ray/gcs/redis_context.cc | 2 ++ src/ray/gcs/tables.h | 27 ++++++++++++++------- 5 files changed, 46 insertions(+), 13 deletions(-) diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index 7b5469665..7a870ed74 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -524,6 +524,9 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleString *id = argv[2]; RedisModuleKey *key = OpenPrefixedKey(ctx, "T:", id, REDISMODULE_READ); + if (key == nullptr) { + return RedisModule_ReplyWithNull(ctx); + } size_t len = 0; const char *buf = RedisModule_StringDMA(key, &len, REDISMODULE_READ); diff --git a/src/plasma/plasma_manager.cc b/src/plasma/plasma_manager.cc index 39c977a8e..a557fdeb2 100644 --- a/src/plasma/plasma_manager.cc +++ b/src/plasma/plasma_manager.cc @@ -1335,6 +1335,9 @@ void log_object_hash_mismatch_error_result_callback(ObjectID object_id, DBClientID::from_binary(t->scheduler_id), std::vector()); log_object_hash_mismatch_error_task_callback(task, user_context); Task_free(task); + }, + [user_context](gcs::AsyncGcsClient *, const TaskID &) { + // TODO(pcmoritz): Handle failure. })); #endif } diff --git a/src/ray/gcs/client_test.cc b/src/ray/gcs/client_test.cc index bff022076..075142ff0 100644 --- a/src/ray/gcs/client_test.cc +++ b/src/ray/gcs/client_test.cc @@ -98,13 +98,19 @@ void Lookup(gcs::AsyncGcsClient *client, const UniqueID &id, test->Stop(); } +void LookupFailed(gcs::AsyncGcsClient *client, const UniqueID &id) { + // Object entry failed. + RAY_CHECK(false); + test->Stop(); +} + void TestObjectTable(const JobID &job_id, std::shared_ptr client) { auto data = std::make_shared(); data->managers.push_back("A"); data->managers.push_back("B"); ObjectID object_id = ObjectID::from_random(); RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded)); - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, &Lookup)); + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, &Lookup, &LookupFailed)); // Run the event loop. The loop will only stop if the Lookup callback is // called (or an assertion failure). test->Start(); @@ -130,16 +136,25 @@ void TaskLookup(gcs::AsyncGcsClient *client, const TaskID &id, ASSERT_EQ(data->scheduling_state, SchedulingState_SCHEDULED); } +void TaskLookupFailure(gcs::AsyncGcsClient *client, const TaskID &id) { + RAY_CHECK(false); +} + void TaskLookupAfterUpdate(gcs::AsyncGcsClient *client, const TaskID &id, std::shared_ptr data) { ASSERT_EQ(data->scheduling_state, SchedulingState_LOST); test->Stop(); } +void TaskLookupAfterUpdateFailure(gcs::AsyncGcsClient *client, const TaskID &id) { + RAY_CHECK(false); + test->Stop(); +} + void TaskUpdateCallback(gcs::AsyncGcsClient *client, const TaskID &task_id, const TaskTableDataT &task, bool updated) { - RAY_CHECK_OK( - client->task_table().Lookup(DriverID::nil(), task_id, &TaskLookupAfterUpdate)); + RAY_CHECK_OK(client->task_table().Lookup( + DriverID::nil(), task_id, &TaskLookupAfterUpdate, &TaskLookupAfterUpdateFailure)); } void TestTaskTable(const JobID &job_id, std::shared_ptr client) { @@ -149,7 +164,8 @@ void TestTaskTable(const JobID &job_id, std::shared_ptr cli data->scheduler_id = local_scheduler_id.binary(); TaskID task_id = TaskID::from_random(); RAY_CHECK_OK(client->task_table().Add(job_id, task_id, data, &TaskAdded)); - RAY_CHECK_OK(client->task_table().Lookup(job_id, task_id, &TaskLookup)); + RAY_CHECK_OK( + client->task_table().Lookup(job_id, task_id, &TaskLookup, &TaskLookupFailure)); auto update = std::make_shared(); update->test_scheduler_id = local_scheduler_id.binary(); update->test_state_bitmask = SchedulingState_SCHEDULED; diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc index aa3cf02d4..d1b3dc946 100644 --- a/src/ray/gcs/redis_context.cc +++ b/src/ray/gcs/redis_context.cc @@ -26,6 +26,7 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) { redisReply *reply = reinterpret_cast(r); std::string data = ""; if (reply->type == REDIS_REPLY_NIL) { + // Respond with blank string, which triggers a failure callback for lookups. } else if (reply->type == REDIS_REPLY_STRING) { data = std::string(reply->str, reply->len); } else if (reply->type == REDIS_REPLY_ARRAY) { @@ -84,6 +85,7 @@ int64_t RedisCallbackManager::add(const RedisCallback &function) { RedisCallbackManager::RedisCallback &RedisCallbackManager::get( int64_t callback_index) { + RAY_CHECK(callbacks_.find(callback_index) != callbacks_.end()); return *callbacks_[callback_index]; } diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 74d8712da..448aa3b8f 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -32,11 +32,13 @@ class Table { using DataT = typename Data::NativeTableType; using Callback = std::function data)>; + using FailureCallback = std::function; struct CallbackData { ID id; std::shared_ptr data; Callback callback; + FailureCallback failure; // An optional callback to call for subscription operations, where the // first message is a notification of subscription success. Callback subscription_callback; @@ -58,7 +60,7 @@ class Table { Status Add(const JobID &job_id, const ID &id, std::shared_ptr data, const Callback &done) { auto d = std::shared_ptr( - new CallbackData({id, data, done, nullptr, this, client_})); + new CallbackData({id, data, done, nullptr, nullptr, this, client_})); int64_t callback_index = RedisCallbackManager::instance().add([d](const std::string &data) { if (d->callback != nullptr) { @@ -79,16 +81,23 @@ class Table { /// \param id The ID of the data that is looked up in the GCS. /// \param lookup Callback that is called after lookup. /// \return Status - Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup) { + Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup, + const FailureCallback &failure) { auto d = std::shared_ptr( - new CallbackData({id, nullptr, lookup, nullptr, this, client_})); + new CallbackData({id, nullptr, lookup, failure, nullptr, this, client_})); int64_t callback_index = RedisCallbackManager::instance().add([d](const std::string &data) { - auto result = std::make_shared(); - auto root = flatbuffers::GetRoot(data.data()); - root->UnPackTo(result.get()); - if (d->callback != nullptr) { - (d->callback)(d->client, d->id, result); + if (data.empty()) { + if (d->failure != nullptr) { + (d->failure)(d->client, d->id); + } + } else { + auto result = std::make_shared(); + auto root = flatbuffers::GetRoot(data.data()); + root->UnPackTo(result.get()); + if (d->callback != nullptr) { + (d->callback)(d->client, d->id, result); + } } }); std::vector nil; @@ -110,7 +119,7 @@ class Table { Status Subscribe(const JobID &job_id, const ClientID &client_id, const Callback &subscribe, const Callback &done) { auto d = std::shared_ptr( - new CallbackData({client_id, nullptr, subscribe, done, this, client_})); + new CallbackData({client_id, nullptr, subscribe, nullptr, done, this, client_})); int64_t callback_index = RedisCallbackManager::instance().add([d](const std::string &data) { if (data.empty()) {