From 985c4241720bf2c15a61bf605b68e2d9b69783b2 Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Sun, 25 Dec 2016 23:57:05 -0800 Subject: [PATCH] Use redismodules for task table and result table. (#156) * Switch to using redis modules for task table. * Switch to using redis modules for the task table. * Fix some tests. * Fix naming and remove code duplication. * Remove duplication in redis modules and add more cleanups. * Address comments. --- src/common/redis_module/ray_redis_module.c | 225 +++++---- src/common/redis_module/runtest.py | 23 +- src/common/state/redis.c | 526 +++++++++------------ src/common/state/task_table.c | 8 +- src/common/state/task_table.h | 42 +- src/common/task.c | 24 +- src/common/task.h | 51 +- src/common/test/db_tests.c | 17 +- src/common/test/object_table_tests.c | 2 +- src/common/test/task_table_tests.c | 12 +- src/global_scheduler/global_scheduler.c | 8 +- src/global_scheduler/global_scheduler.h | 2 +- src/global_scheduler/test/test.py | 7 +- src/photon/photon_scheduler.c | 9 +- 14 files changed, 456 insertions(+), 500 deletions(-) diff --git a/src/common/redis_module/ray_redis_module.c b/src/common/redis_module/ray_redis_module.c index 508c5daa7..831848282 100644 --- a/src/common/redis_module/ray_redis_module.c +++ b/src/common/redis_module/ray_redis_module.c @@ -201,8 +201,7 @@ int ObjectTableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleKey *key = OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, argv[1], REDISMODULE_READ); - int keytype = RedisModule_KeyType(key); - if (keytype == REDISMODULE_KEYTYPE_EMPTY || + if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY || RedisModule_ValueLength(key) == 0) { return RedisModule_ReplyWithArray(ctx, 0); } @@ -330,15 +329,20 @@ int ObjectTableAdd_RedisCommand(RedisModuleCtx *ctx, key = OpenPrefixedKey(ctx, OBJECT_INFO_PREFIX, object_id, REDISMODULE_READ | REDISMODULE_WRITE); - int keytype = RedisModule_KeyType(key); /* Check if this object was already registered and if the hashes agree. */ - if (keytype != REDISMODULE_KEYTYPE_EMPTY) { + if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_EMPTY) { RedisModuleString *existing_hash; RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "hash", &existing_hash, NULL); - if (RedisModule_StringCompare(existing_hash, new_hash) != 0) { - RedisModule_CloseKey(key); - return RedisModule_ReplyWithError(ctx, "hash mismatch"); + /* The existing hash may be NULL even if the key is present because a call + * to RAY.RESULT_TABLE_ADD may have already created the key. */ + if (existing_hash != NULL) { + if (RedisModule_StringCompare(existing_hash, new_hash) != 0) { + RedisModule_CloseKey(key); + RedisModule_FreeString(ctx, existing_hash); + return RedisModule_ReplyWithError(ctx, "hash mismatch"); + } + RedisModule_FreeString(ctx, existing_hash); } } @@ -371,9 +375,8 @@ int ObjectTableAdd_RedisCommand(RedisModuleCtx *ctx, OpenPrefixedKey(ctx, OBJECT_NOTIFICATION_PREFIX, object_id, REDISMODULE_READ | REDISMODULE_WRITE); /* If the zset exists, initialize the key to iterate over the zset. */ - int object_notification_keytype = - RedisModule_KeyType(object_notification_key); - if (object_notification_keytype != REDISMODULE_KEYTYPE_EMPTY) { + if (RedisModule_KeyType(object_notification_key) != + REDISMODULE_KEYTYPE_EMPTY) { CHECK_ERROR(RedisModule_ZsetFirstInScoreRange( object_notification_key, REDISMODULE_NEGATIVE_INFINITE, REDISMODULE_POSITIVE_INFINITE, 1, 1), @@ -437,8 +440,7 @@ int ObjectTableRemove_RedisCommand(RedisModuleCtx *ctx, RedisModuleKey *table_key; table_key = OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, object_id, REDISMODULE_READ | REDISMODULE_WRITE); - int keytype = RedisModule_KeyType(table_key); - if (keytype == REDISMODULE_KEYTYPE_EMPTY) { + if (RedisModule_KeyType(table_key) == REDISMODULE_KEYTYPE_EMPTY) { RedisModule_CloseKey(table_key); return RedisModule_ReplyWithError(ctx, "object not found"); } @@ -485,8 +487,7 @@ int ObjectTableRequestNotifications_RedisCommand(RedisModuleCtx *ctx, RedisModuleString *object_id = argv[i]; RedisModuleKey *key = OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, object_id, REDISMODULE_READ); - int keytype = RedisModule_KeyType(key); - if (keytype == REDISMODULE_KEYTYPE_EMPTY || + if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY || RedisModule_ValueLength(key) == 0) { /* This object ID is currently not present, so make a note that this * client should be notified when this object ID becomes available. */ @@ -505,8 +506,7 @@ int ObjectTableRequestNotifications_RedisCommand(RedisModuleCtx *ctx, RedisModuleKey *object_info_key; object_info_key = OpenPrefixedKey(ctx, OBJECT_INFO_PREFIX, object_id, REDISMODULE_READ); - int keytype = RedisModule_KeyType(key); - if (keytype == REDISMODULE_KEYTYPE_EMPTY) { + if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) { RedisModule_CloseKey(object_info_key); RedisModule_CloseKey(key); return RedisModule_ReplyWithError(ctx, "requested object not found"); @@ -515,9 +515,15 @@ int ObjectTableRequestNotifications_RedisCommand(RedisModuleCtx *ctx, RedisModule_HashGet(object_info_key, REDISMODULE_HASH_CFIELDS, "data_size", &existing_data_size, NULL); RedisModule_CloseKey(object_info_key); /* No longer needed. */ + if (existing_data_size == NULL) { + RedisModule_CloseKey(key); + return RedisModule_ReplyWithError(ctx, + "no data_size field in object info"); + } bool success = PublishObjectNotification(ctx, client_id, object_id, existing_data_size, key); + RedisModule_FreeString(ctx, existing_data_size); if (!success) { /* The publish failed somehow. */ RedisModule_CloseKey(key); @@ -575,6 +581,66 @@ int ResultTableAdd_RedisCommand(RedisModuleCtx *ctx, return REDISMODULE_OK; } +/** + * Reply with information about a task ID. This is used by + * RAY.RESULT_TABLE_LOOKUP and RAY.TASK_TABLE_GET. + * + * @param task_id The task ID of the task to reply about. + * @return NIL if the task ID is not in the task table. An error if the task ID + * is in the task table but the appropriate fields are not there, and + * an array of the task scheduling state, the local scheduler ID, and + * the task spec for the task otherwise. + */ +int ReplyWithTask(RedisModuleCtx *ctx, RedisModuleString *task_id) { + RedisModuleKey *key = + OpenPrefixedKey(ctx, TASK_PREFIX, task_id, REDISMODULE_READ); + + if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_EMPTY) { + /* If the key exists, look up the fields and return them in an array. */ + RedisModuleString *state = NULL; + RedisModuleString *local_scheduler_id = NULL; + RedisModuleString *task_spec = NULL; + RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state, "node", + &local_scheduler_id, "task_spec", &task_spec, NULL); + if (state == NULL || local_scheduler_id == NULL || task_spec == NULL) { + /* We must have either all fields or no fields. */ + RedisModule_CloseKey(key); + return RedisModule_ReplyWithError( + ctx, "Missing fields in the task table entry"); + } + + size_t state_length; + const char *state_string = RedisModule_StringPtrLen(state, &state_length); + int state_integer; + int scanned = sscanf(state_string, "%2d", &state_integer); + if (scanned != 1 || state_length != 2) { + RedisModule_CloseKey(key); + RedisModule_FreeString(ctx, state); + RedisModule_FreeString(ctx, local_scheduler_id); + RedisModule_FreeString(ctx, task_spec); + return RedisModule_ReplyWithError(ctx, + "Found invalid scheduling state (must " + "be an integer of width 2"); + } + + RedisModule_ReplyWithArray(ctx, 3); + RedisModule_ReplyWithLongLong(ctx, state_integer); + RedisModule_ReplyWithString(ctx, local_scheduler_id); + RedisModule_ReplyWithString(ctx, task_spec); + + RedisModule_FreeString(ctx, state); + RedisModule_FreeString(ctx, local_scheduler_id); + RedisModule_FreeString(ctx, task_spec); + } else { + /* If the key does not exist, return nil. */ + RedisModule_ReplyWithNull(ctx); + } + + RedisModule_CloseKey(key); + + return REDISMODULE_OK; +} + /** * Lookup an entry in the result table. * @@ -583,8 +649,10 @@ int ResultTableAdd_RedisCommand(RedisModuleCtx *ctx, * RAY.RESULT_TABLE_LOOKUP * * @param object_id A string representing the object ID. - * @return An empty string if the object ID is not in the result table and the - * task ID of the task that created the object ID otherwise. + * @return NIL if the object ID is not in the result table or if the + * corresponding task ID is not in the task table. Otherwise, this + * returns an array of the scheduling state, the local scheduler ID, and + * the task spec for the task corresponding to this object ID. */ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, @@ -599,23 +667,24 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleKey *key; key = OpenPrefixedKey(ctx, OBJECT_INFO_PREFIX, object_id, REDISMODULE_READ); - int keytype = RedisModule_KeyType(key); - if (keytype == REDISMODULE_KEYTYPE_EMPTY) { - return RedisModule_ReplyWithStringBuffer(ctx, "", 0); + if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) { + return RedisModule_ReplyWithNull(ctx); } RedisModuleString *task_id; RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task", &task_id, NULL); if (task_id == NULL) { - return RedisModule_ReplyWithStringBuffer(ctx, "", 0); + return RedisModule_ReplyWithNull(ctx); } - RedisModule_ReplyWithString(ctx, task_id); + + /* Construct a reply by getting the task from the task ID. */ + int status = ReplyWithTask(ctx, task_id); /* Clean up. */ RedisModule_FreeString(ctx, task_id); RedisModule_CloseKey(key); - return REDISMODULE_OK; + return status; } int TaskTableWrite(RedisModuleCtx *ctx, @@ -641,14 +710,16 @@ int TaskTableWrite(RedisModuleCtx *ctx, /* Add the task to the task table. If no spec was provided, get the existing * spec out of the task table so we can publish it. */ + RedisModuleString *existing_task_spec = NULL; RedisModuleKey *key = OpenPrefixedKey(ctx, TASK_PREFIX, task_id, REDISMODULE_WRITE); if (task_spec == NULL) { RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node", node_id, NULL); - RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task_spec", &task_spec, - NULL); - if (task_spec == NULL) { + RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task_spec", + &existing_task_spec, NULL); + if (existing_task_spec == NULL) { + RedisModule_CloseKey(key); return RedisModule_ReplyWithError( ctx, "Cannot update a task that doesn't exist yet"); } @@ -664,18 +735,29 @@ int TaskTableWrite(RedisModuleCtx *ctx, * specification>". */ RedisModuleString *publish_topic = RedisString_Format(ctx, "%s%S:%S", TASK_PREFIX, node_id, state); - RedisModuleString *publish_message = RedisString_Format( - ctx, "%S %S %S %S", task_id, state, node_id, task_spec); + RedisModuleString *publish_message; + if (task_spec != NULL) { + publish_message = RedisString_Format(ctx, "%S %S %S %S", task_id, state, + node_id, task_spec); + } else { + publish_message = RedisString_Format(ctx, "%S %S %S %S", task_id, state, + node_id, existing_task_spec); + } RedisModuleCallReply *reply = RedisModule_Call(ctx, "PUBLISH", "ss", publish_topic, publish_message); + + RedisModule_FreeString(ctx, publish_message); + RedisModule_FreeString(ctx, publish_topic); + if (existing_task_spec != NULL) { + RedisModule_FreeString(ctx, existing_task_spec); + } + if (reply == NULL) { return RedisModule_ReplyWithError(ctx, "PUBLISH unsuccessful"); } - RedisModule_FreeString(ctx, publish_message); - RedisModule_FreeString(ctx, publish_topic); - RedisModule_ReplyWithSimpleString(ctx, "ok"); + RedisModule_ReplyWithSimpleString(ctx, "OK"); return REDISMODULE_OK; } @@ -686,7 +768,7 @@ int TaskTableWrite(RedisModuleCtx *ctx, * * This is called from a client with the command: * - * RAY.task_table_add + * RAY.TASK_TABLE_ADD * * @param task_id A string that is the ID of the task. * @param state A string that is the current scheduling state (a @@ -694,7 +776,8 @@ int TaskTableWrite(RedisModuleCtx *ctx, * nonnegative integer less than 100, so that it has width at most 2. If * less than 2, the value will be left-padded with spaces to a width of * 2. - * @param node_id A string that is the ID of the associated node, if any. + * @param local_scheduler_id A string that is the ray client ID of the + * associated local scheduler, if any. * @param task_spec A string that is the specification of the task, which can * be cast to a `task_spec`. * @return OK if the operation was successful. @@ -715,7 +798,7 @@ int TaskTableAddTask_RedisCommand(RedisModuleCtx *ctx, * * This is called from a client with the command: * - * RAY.task_table_update_task + * RAY.TASK_TABLE_UPDATE * * @param task_id A string that is the ID of the task. * @param state A string that is the current scheduling state (a @@ -723,7 +806,8 @@ int TaskTableAddTask_RedisCommand(RedisModuleCtx *ctx, * nonnegative integer less than 100, so that it has width at most 2. If * less than 2, the value will be left-padded with spaces to a width of * 2. - * @param node_id A string that is the ID of the associated node, if any. + * @param ray_client_id A string that is the ray client ID of the associated + * local scheduler, if any. * @return OK if the operation was successful. */ int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx, @@ -741,7 +825,7 @@ int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx, * * This is called from a client with the command: * - * RAY.task_table_get + * RAY.TASK_TABLE_GET * * @param task_id A string of the task ID to look up. * @return An array of strings representing the task fields in the following @@ -749,64 +833,15 @@ int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx, * if any 3) (string) the task specification, which can be casted to a * task_spec. If the task ID is not in the table, returns nil. */ -int TaskTableGetTask_RedisCommand(RedisModuleCtx *ctx, - RedisModuleString **argv, - int argc) { +int TaskTableGet_RedisCommand(RedisModuleCtx *ctx, + RedisModuleString **argv, + int argc) { if (argc != 2) { return RedisModule_WrongArity(ctx); } - RedisModuleKey *key = - OpenPrefixedKey(ctx, TASK_PREFIX, argv[1], REDISMODULE_READ); - - int keytype = RedisModule_KeyType(key); - if (keytype != REDISMODULE_KEYTYPE_EMPTY) { - /* If the key exists, look up the fields and return them in an array. */ - RedisModuleString *state = NULL, *node = NULL, *task_spec = NULL; - RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state, "node", - &node, "task_spec", &task_spec, NULL); - if (state == NULL || node == NULL || task_spec == NULL) { - /* We must have either all fields or no fields. */ - return RedisModule_ReplyWithError( - ctx, "Missing fields in the task table entry"); - } - - size_t state_length; - const char *state_string = RedisModule_StringPtrLen(state, &state_length); - int state_integer; - int scanned = sscanf(state_string, "%2d", &state_integer); - if (scanned != 1 || state_length != 2) { - return RedisModule_ReplyWithError(ctx, - "Found invalid scheduling state (must " - "be an integer of width 2"); - } - - RedisModule_ReplyWithArray(ctx, 3); - RedisModule_ReplyWithLongLong(ctx, state_integer); - RedisModule_ReplyWithString(ctx, node); - RedisModule_ReplyWithString(ctx, task_spec); - - RedisModule_FreeString(ctx, task_spec); - RedisModule_FreeString(ctx, node); - RedisModule_FreeString(ctx, state); - } else { - /* If the key does not exist, return nil. */ - RedisModule_ReplyWithNull(ctx); - } - - RedisModule_CloseKey(key); - - return REDISMODULE_OK; -} - -int TaskTableSubscribe_RedisCommand(RedisModuleCtx *ctx, - RedisModuleString **argv, - int argc) { - /* TODO(swang): Implement this. */ - REDISMODULE_NOT_USED(ctx); - REDISMODULE_NOT_USED(argv); - REDISMODULE_NOT_USED(argc); - return REDISMODULE_OK; + /* Construct a reply by getting the task from the task ID. */ + return ReplyWithTask(ctx, argv[1]); } /* This function must be present on each Redis module. It is used in order to @@ -888,13 +923,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, } if (RedisModule_CreateCommand(ctx, "ray.task_table_get", - TaskTableGetTask_RedisCommand, "readonly", 0, 0, - 0) == REDISMODULE_ERR) { - return REDISMODULE_ERR; - } - - if (RedisModule_CreateCommand(ctx, "ray.task_table_subscribe", - TaskTableSubscribe_RedisCommand, "pubsub", 0, 0, + TaskTableGet_RedisCommand, "readonly", 0, 0, 0) == REDISMODULE_ERR) { return REDISMODULE_ERR; } diff --git a/src/common/redis_module/runtest.py b/src/common/redis_module/runtest.py index 3d9544116..156bab161 100644 --- a/src/common/redis_module/runtest.py +++ b/src/common/redis_module/runtest.py @@ -45,6 +45,7 @@ class TestGlobalStateStore(unittest.TestCase): redis_port = random.randint(2000, 50000) self.redis_process = subprocess.Popen([redis_path, "--port", str(redis_port), + "--loglevel", "warning", "--loadmodule", module_path]) time.sleep(1.5) self.redis = redis.StrictRedis(host="localhost", port=redis_port, db=0) @@ -165,17 +166,31 @@ class TestGlobalStateStore(unittest.TestCase): %integerToAsciiHex(data_size, 8)) def testResultTableAddAndLookup(self): + # Try looking up something in the result table before anything is added. response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id1") - self.assertEqual(set(response), set([])) + self.assertIsNone(response) + # Adding the object to the object table should have no effect. self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id1", 1, "hash1", "manager_id1") response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id1") - self.assertEqual(set(response), set([])) + self.assertIsNone(response) + # Add the result to the result table. This is necessary, but not sufficient + # because the task is still not in the task table. self.redis.execute_command("RAY.RESULT_TABLE_ADD", "object_id1", "task_id1") response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id1") - self.assertEqual(response, b"task_id1") + self.assertIsNone(response) + # Add the task to the task table so that the result table lookup can + # succeed. + self.redis.execute_command("RAY.TASK_TABLE_ADD", "task_id1", 1, "local_scheduler_id1", "task_spec1") + response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id1") + self.assertEqual(response, [1, b"local_scheduler_id1", b"task_spec1"]) + # Doing it again should still work. + response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id1") + self.assertEqual(response, [1, b"local_scheduler_id1", b"task_spec1"]) + # Try another result table lookup. This should succeed. + self.redis.execute_command("RAY.TASK_TABLE_ADD", "task_id2", 2, "local_scheduler_id2", "task_spec2") self.redis.execute_command("RAY.RESULT_TABLE_ADD", "object_id2", "task_id2") response = self.redis.execute_command("RAY.RESULT_TABLE_LOOKUP", "object_id2") - self.assertEqual(response, b"task_id2") + self.assertEqual(response, [2, b"local_scheduler_id2", b"task_spec2"]) def testInvalidTaskTableAdd(self): # Check that Redis returns an error when RAY.TASK_TABLE_ADD is called with diff --git a/src/common/state/redis.c b/src/common/state/redis.c index d9043d544..07def309e 100644 --- a/src/common/state/redis.c +++ b/src/common/state/redis.c @@ -88,48 +88,6 @@ typedef struct { bool is_redis_reply; } redis_requests_info; -/** - * A header for callbacks similar to REDIS_CALLBACK_HEADER, but for operations - * that span multiple Redis commands. The differences are: - * - Instead of passing in the table operation's timer ID as the asynchronous - * command callback's `privdata` argument, the user must pass a pointer to a - * redis_requests_info instance. - * - The user must define an additional REQUEST_INFO variable name, which will - * hold a reference to the redis_requests_info passed into the Redis - * asynchronous command. - */ -#define REDIS_MULTI_CALLBACK_HEADER(DB, CB_DATA, REPLY, REQUEST_INFO) \ - db_handle *DB = c->data; \ - redis_requests_info *REQUEST_INFO = privdata; \ - DCHECK(REQUEST_INFO != NULL); \ - if ((REPLY) == NULL && REQUEST_INFO->is_redis_reply) { \ - free(REQUEST_INFO); \ - return; \ - } \ - table_callback_data *CB_DATA = \ - outstanding_callbacks_find(REQUEST_INFO->timer_id); \ - if (CB_DATA == NULL) { \ - /* the callback data structure has been \ - * already freed; just ignore this reply */ \ - free(privdata); \ - return; \ - } \ - do { \ - } while (0) - -/** - * A data structure to keep track of object IDs when doing object table - * lookups. - * TODO(swang): Remove this when we integrate a Redis module implementation. - */ -typedef struct { - /** The timer ID that uniquely identifies this table operation. All retry - * attempts of a table operation share the same timer ID. */ - int64_t timer_id; - /** The object ID that the request was for. */ - object_id object_id; -} object_table_get_entry_info; - db_handle *db_connect(const char *db_address, int db_port, const char *client_type, @@ -263,63 +221,6 @@ void db_attach(db_handle *db, event_loop *loop, bool reattach) { } } -/** - * An internal function to allocate a task object and parse a hashmap reply - * from Redis into the task object. If the Redis reply is malformed, an empty - * task with the given task ID is returned. - * - * @param id The ID of the task we're looking up. If the reply from Redis is - * well-formed, the reply's ID should match this ID. Else, the returned - * task will have its ID set to this ID. - * @param num_redis_replies The number of keys and values in the Redis hashmap. - * @param redis_replies A pointer to the Redis hashmap keys and values. - * @return A pointer to the parsed task. - */ -task *parse_redis_task_table_entry(task_id id, - int num_redis_replies, - redisReply **redis_replies) { - task *task_result; - if (num_redis_replies == 0) { - /* There was no information about this task. */ - return NULL; - } - /* Exit immediately if there weren't 6 fields, one for each key-value pair. - * The keys are "node", "state", and "task_spec". */ - DCHECK(num_redis_replies == 6); - /* Parse the task struct's fields. */ - scheduling_state state = 0; - node_id node = NIL_ID; - task_spec *spec = NULL; - for (int i = 0; i < num_redis_replies; i = i + 2) { - char *key = redis_replies[i]->str; - redisReply *value = redis_replies[i + 1]; - if (strcmp(key, "node") == 0) { - DCHECK(value->len == sizeof(node_id)); - memcpy(&node, value->str, value->len); - } else if (strcmp(key, "state") == 0) { - int scanned = sscanf(value->str, "%d", (int *) &state); - if (scanned != 1) { - LOG_FATAL("Scheduling state for task is malformed"); - state = 0; - } - } else if (strcmp(key, "task_spec") == 0) { - spec = malloc(value->len); - memcpy(spec, value->str, value->len); - } else { - LOG_FATAL("Found unexpected %s field in task log", key); - } - } - /* Exit immediately if we couldn't parse the task spec. */ - if (spec == NULL) { - LOG_FATAL("Could not parse task spec from task log"); - } - /* Build and return the task. */ - DCHECK(task_ids_equal(task_spec_id(spec), id)); - task_result = alloc_task(spec, state, node); - free_task_spec(spec); - return task_result; -} - /* * ==== object_table callbacks ==== */ @@ -417,11 +318,6 @@ void redis_object_table_lookup(table_callback_data *callback_data) { db_handle *db = callback_data->db_handle; object_id obj_id = callback_data->id; - // object_table_get_entry_info *context = - // malloc(sizeof(object_table_get_entry_info)); - // context->timer_id = callback_data->timer_id; - // context->object_id = id; - int status = redisAsyncCommand( db->context, redis_object_table_lookup_callback, (void *) callback_data->timer_id, "RAY.OBJECT_TABLE_LOOKUP %b", obj_id.id, @@ -436,8 +332,11 @@ void redis_result_table_add_callback(redisAsyncContext *c, void *privdata) { REDIS_CALLBACK_HEADER(db, callback_data, r); redisReply *reply = r; - CHECK(reply->type == REDIS_REPLY_STATUS || - reply->type == REDIS_REPLY_INTEGER); + /* Check that the command succeeded. */ + CHECK(reply->type != REDIS_REPLY_ERROR); + CHECKM(strncmp(reply->str, "OK", strlen("OK")) == 0, "reply->str is %s", + reply->str); + /* Call the done callback if there is one. */ if (callback_data->done_callback) { result_table_done_callback done_callback = callback_data->done_callback; done_callback(callback_data->id, callback_data->user_context); @@ -453,77 +352,78 @@ void redis_result_table_add(table_callback_data *callback_data) { /* Add the result entry to the result table. */ int status = redisAsyncCommand( db->context, redis_result_table_add_callback, - (void *) callback_data->timer_id, "SET result:%b %b", id.id, + (void *) callback_data->timer_id, "RAY.RESULT_TABLE_ADD %b %b", id.id, sizeof(id.id), result_task_id->id, sizeof(result_task_id->id)); if ((status == REDIS_ERR) || db->context->err) { LOG_REDIS_DEBUG(db->context, "Error in result table add"); } } -void redis_result_table_lookup_task_callback(redisAsyncContext *c, - void *r, - void *privdata) { - REDIS_CALLBACK_HEADER(db, callback_data, r); - redisReply *reply = r; - /* Check that we received a Redis hashmap. */ - if (reply->type != REDIS_REPLY_ARRAY) { - LOG_FATAL("Expected Redis array, received type %d %s", reply->type, - reply->str); +/* This allocates a task which must be freed by the caller, unless the returned + * task is NULL. This is used by both redis_result_table_lookup_callback and + * redis_task_table_get_task_callback. */ +task *parse_and_construct_task_from_redis_reply(redisReply *reply) { + task *task; + if (reply->type == REDIS_REPLY_NIL) { + /* There is no task in the reply, so return NULL. */ + task = NULL; + } else if (reply->type == REDIS_REPLY_ARRAY) { + /* Check that the reply is as expected. The 0th element is the scheduling + * state. The 1st element is the db_client_id of the associated local + * scheduler, and the 2nd element is the task_spec. */ + CHECK(reply->elements == 3); + CHECK(reply->element[0]->type == REDIS_REPLY_INTEGER); + CHECK(reply->element[1]->type == REDIS_REPLY_STRING); + CHECK(reply->element[2]->type == REDIS_REPLY_STRING); + /* Parse the scheduling state. */ + long long state = reply->element[0]->integer; + /* Parse the local scheduler db_client_id. */ + db_client_id local_scheduler_id; + CHECK(sizeof(local_scheduler_id) == reply->element[1]->len); + memcpy(local_scheduler_id.id, reply->element[1]->str, + reply->element[1]->len); + /* Parse the task spec. */ + task_spec *spec = malloc(reply->element[2]->len); + memcpy(spec, reply->element[2]->str, reply->element[2]->len); + CHECK(task_spec_size(spec) == reply->element[2]->len); + task = alloc_task(spec, state, local_scheduler_id); + /* Free the task spec. */ + free_task_spec(spec); + } else { + LOG_FATAL("Unexpected reply type %d", reply->type); } - /* If the user registered a success callback, construct the task object from - * the Redis reply and call the callback. */ - result_table_lookup_callback done_callback = callback_data->done_callback; - task_id *result_task_id = callback_data->data; - if (done_callback) { - task *task_reply = parse_redis_task_table_entry( - *result_task_id, reply->elements, reply->element); - done_callback(callback_data->id, task_reply, callback_data->user_context); - free_task(task_reply); - } - destroy_timer_callback(db->loop, callback_data); + /* Return the task. If it is not NULL, then it must be freed by the caller. */ + return task; } -void redis_result_table_lookup_object_callback(redisAsyncContext *c, - void *r, - void *privdata) { +void redis_result_table_lookup_callback(redisAsyncContext *c, + void *r, + void *privdata) { REDIS_CALLBACK_HEADER(db, callback_data, r); redisReply *reply = r; - - if (reply->type == REDIS_REPLY_STRING) { - /* If we found the object, get the spec of the task that created it. */ - DCHECK(reply->len == sizeof(task_id)); - task_id *result_task_id = malloc(sizeof(task_id)); - memcpy(result_task_id, reply->str, reply->len); - callback_data->data = (void *) result_task_id; - int status = - redisAsyncCommand(db->context, redis_result_table_lookup_task_callback, - (void *) callback_data->timer_id, "HGETALL task:%b", - result_task_id->id, sizeof(result_task_id->id)); - if ((status == REDIS_ERR) || db->context->err) { - LOG_REDIS_DEBUG(db->context, "Could not look up result table entry"); - } - } else if (reply->type == REDIS_REPLY_NIL) { - /* The object with the requested ID was not in the table. */ - LOG_INFO("Object's result not in table."); - result_table_lookup_callback done_callback = callback_data->done_callback; - if (done_callback) { - done_callback(callback_data->id, NULL, callback_data->user_context); - } - destroy_timer_callback(db->loop, callback_data); - return; - } else { - LOG_FATAL("expected string or nil, received type %d", reply->type); + /* Parse the task from the reply. */ + task *task = parse_and_construct_task_from_redis_reply(reply); + /* Call the done callback if there is one. */ + result_table_lookup_callback done_callback = callback_data->done_callback; + if (done_callback != NULL) { + done_callback(callback_data->id, task, callback_data->user_context); } + /* Free the task if it is not NULL. */ + if (task != NULL) { + free_task(task); + } + /* Clean up timer and callback. */ + destroy_timer_callback(db->loop, callback_data); } void redis_result_table_lookup(table_callback_data *callback_data) { CHECK(callback_data); db_handle *db = callback_data->db_handle; - /* First, lookup the ID of the task that created this object. */ object_id id = callback_data->id; - int status = redisAsyncCommand( - db->context, redis_result_table_lookup_object_callback, - (void *) callback_data->timer_id, "GET result:%b", id.id, sizeof(id.id)); + int status = + redisAsyncCommand(db->context, redis_result_table_lookup_callback, + (void *) callback_data->timer_id, + "RAY.RESULT_TABLE_LOOKUP %b", id.id, sizeof(id.id)); if ((status == REDIS_ERR) || db->context->err) { LOG_REDIS_DEBUG(db->context, "Error in result table lookup"); } @@ -822,143 +722,143 @@ void redis_task_table_get_task_callback(redisAsyncContext *c, void *privdata) { REDIS_CALLBACK_HEADER(db, callback_data, r); redisReply *reply = r; - /* Check that we received a Redis hashmap. */ - if (reply->type != REDIS_REPLY_ARRAY) { - LOG_FATAL("Expected Redis array, received type %d %s", reply->type, - reply->str); - } - /* If the user registered a success callback, construct the task object from - * the Redis reply and call the callback. */ - if (callback_data->done_callback) { - task_table_get_callback done_callback = callback_data->done_callback; - task *task_reply = parse_redis_task_table_entry( - callback_data->id, reply->elements, reply->element); - done_callback(task_reply, callback_data->user_context); - free_task(task_reply); + /* Parse the task from the reply. */ + task *task = parse_and_construct_task_from_redis_reply(reply); + /* Call the done callback if there is one. */ + task_table_get_callback done_callback = callback_data->done_callback; + if (done_callback != NULL) { + done_callback(task, callback_data->user_context); } + /* Free the task if it is not NULL. */ + free_task(task); + + /* Clean up the timer and callback. */ destroy_timer_callback(db->loop, callback_data); } void redis_task_table_get_task(table_callback_data *callback_data) { - CHECK(callback_data); db_handle *db = callback_data->db_handle; - task_id id = callback_data->id; - int status = - redisAsyncCommand(db->context, redis_task_table_get_task_callback, - (void *) callback_data->timer_id, "HGETALL task:%b", - id.id, sizeof(id.id)); + CHECK(callback_data->data == NULL); + task_id task_id = callback_data->id; + + int status = redisAsyncCommand( + db->context, redis_task_table_get_task_callback, + (void *) callback_data->timer_id, "RAY.TASK_TABLE_GET %b", task_id.id, + sizeof(task_id.id)); if ((status == REDIS_ERR) || db->context->err) { - LOG_REDIS_DEBUG(db->context, "Could not get task from task table"); + LOG_REDIS_DEBUG(db->context, "error in redis_task_table_get_task"); } } -void redis_task_table_publish(table_callback_data *callback_data, - bool task_added) { - db_handle *db = callback_data->db_handle; - task *task = callback_data->data; - task_id id = task_task_id(task); - node_id node = task_node(task); - scheduling_state state = task_state(task); - task_spec *spec = task_task_spec(task); +void redis_task_table_add_task_callback(redisAsyncContext *c, + void *r, + void *privdata) { + REDIS_CALLBACK_HEADER(db, callback_data, r); - LOG_DEBUG("Called log_publish callback"); - -/* Check whether the vector (requests_info) indicating the status of the - * requests has been allocated. - * If was not allocate it, allocate it and initialize it. - * This vector has an entry for each redis command, and it stores true if a - * reply for that command - * has been received, and false otherwise. - * The first entry in the callback corresponds to RPUSH, and the second entry to - * PUBLISH. - */ -#define NUM_PUBLISH_COMMANDS 2 -#define PUBLISH_PUSH_INDEX 0 -#define PUBLISH_PUBLISH_INDEX 1 - if (callback_data->requests_info == NULL) { - callback_data->requests_info = malloc(NUM_PUBLISH_COMMANDS * sizeof(bool)); - for (int i = 0; i < NUM_PUBLISH_COMMANDS; i++) { - ((bool *) callback_data->requests_info)[i] = false; - } - } - - if (((bool *) callback_data->requests_info)[PUBLISH_PUSH_INDEX] == false) { - /* If the task has already been added to the task table, only update the - * scheduling information fields. */ - int status = REDIS_OK; - if (task_added) { - status = redisAsyncCommand( - db->context, redis_task_table_publish_push_callback, - (void *) callback_data->timer_id, "HMSET task:%b state %d node %b", - (char *) id.id, sizeof(id.id), state, (char *) node.id, - sizeof(node.id)); - } else { - status = redisAsyncCommand( - db->context, redis_task_table_publish_push_callback, - (void *) callback_data->timer_id, - "HMSET task:%b state %d node %b task_spec %b", (char *) id.id, - sizeof(id.id), state, (char *) node.id, sizeof(node.id), - (char *) spec, task_spec_size(spec)); - } - if ((status == REDIS_ERR) || db->context->err) { - LOG_REDIS_DEBUG(db->context, "error setting task in task_table_add_task"); - } - } - - if (((bool *) callback_data->requests_info)[PUBLISH_PUBLISH_INDEX] == false) { - int status = redisAsyncCommand( - db->context, redis_task_table_publish_publish_callback, - (void *) callback_data->timer_id, "PUBLISH task:%b:%d %b", - (char *) node.id, sizeof(node.id), state, (char *) task, - task_size(task)); - - if ((status == REDIS_ERR) || db->context->err) { - LOG_REDIS_DEBUG(db->context, - "error publishing task in task_table_add_task"); - } + /* Do some minimal checking. */ + redisReply *reply = r; + CHECKM(strcmp(reply->str, "OK") == 0, "reply->str is %s", reply->str); + /* Call the done callback if there is one. */ + if (callback_data->done_callback != NULL) { + task_table_done_callback done_callback = callback_data->done_callback; + done_callback(callback_data->id, callback_data->user_context); } + /* Clean up the timer and callback. */ + destroy_timer_callback(db->loop, callback_data); } void redis_task_table_add_task(table_callback_data *callback_data) { - redis_task_table_publish(callback_data, false); + db_handle *db = callback_data->db_handle; + task *task = callback_data->data; + task_id task_id = task_task_id(task); + db_client_id local_scheduler_id = task_local_scheduler(task); + scheduling_state state = task_state(task); + task_spec *spec = task_task_spec(task); + + CHECKM(task != NULL, "NULL task passed to redis_task_table_add_task."); + int status = redisAsyncCommand( + db->context, redis_task_table_add_task_callback, + (void *) callback_data->timer_id, "RAY.TASK_TABLE_ADD %b %d %b %b", + task_id.id, sizeof(task_id.id), state, local_scheduler_id.id, + sizeof(local_scheduler_id.id), spec, task_spec_size(spec)); + if ((status == REDIS_ERR) || db->context->err) { + LOG_REDIS_DEBUG(db->context, "error in redis_task_table_add_task"); + } +} + +void redis_task_table_update_callback(redisAsyncContext *c, + void *r, + void *privdata) { + REDIS_CALLBACK_HEADER(db, callback_data, r); + + /* Do some minimal checking. */ + redisReply *reply = r; + CHECKM(strcmp(reply->str, "OK") == 0, "reply->str is %s", reply->str); + /* Call the done callback if there is one. */ + if (callback_data->done_callback != NULL) { + task_table_done_callback done_callback = callback_data->done_callback; + done_callback(callback_data->id, callback_data->user_context); + } + /* Clean up the timer and callback. */ + destroy_timer_callback(db->loop, callback_data); } void redis_task_table_update(table_callback_data *callback_data) { - redis_task_table_publish(callback_data, true); -} + db_handle *db = callback_data->db_handle; + task *task = callback_data->data; + task_id task_id = task_task_id(task); + db_client_id local_scheduler_id = task_local_scheduler(task); + scheduling_state state = task_state(task); -void redis_task_table_publish_push_callback(redisAsyncContext *c, - void *r, - void *privdata) { - LOG_DEBUG("Calling publish push callback"); - REDIS_CALLBACK_HEADER(db, callback_data, r); - CHECK(callback_data->requests_info != NULL); - ((bool *) callback_data->requests_info)[PUBLISH_PUSH_INDEX] = true; - - if (((bool *) callback_data->requests_info)[PUBLISH_PUBLISH_INDEX] == true) { - if (callback_data->done_callback) { - task_table_done_callback done_callback = callback_data->done_callback; - done_callback(callback_data->id, callback_data->user_context); - } - destroy_timer_callback(db->loop, callback_data); + CHECKM(task != NULL, "NULL task passed to redis_task_table_update."); + int status = redisAsyncCommand( + db->context, redis_task_table_update_callback, + (void *) callback_data->timer_id, "RAY.TASK_TABLE_UPDATE %b %d %b", + task_id.id, sizeof(task_id.id), state, local_scheduler_id.id, + sizeof(local_scheduler_id.id)); + if ((status == REDIS_ERR) || db->context->err) { + LOG_REDIS_DEBUG(db->context, "error in redis_task_table_update"); } } -void redis_task_table_publish_publish_callback(redisAsyncContext *c, - void *r, - void *privdata) { - LOG_DEBUG("Calling publish publish callback"); - REDIS_CALLBACK_HEADER(db, callback_data, r); - CHECK(callback_data->requests_info != NULL); - ((bool *) callback_data->requests_info)[PUBLISH_PUBLISH_INDEX] = true; - - if (((bool *) callback_data->requests_info)[PUBLISH_PUSH_INDEX] == true) { - if (callback_data->done_callback) { - task_table_done_callback done_callback = callback_data->done_callback; - done_callback(callback_data->id, callback_data->user_context); - } - destroy_timer_callback(db->loop, callback_data); - } +/* The format of the payload is described in ray_redis_module.c and is + * " ". TODO(rkn): + * Make this code nicer. */ +void parse_task_table_subscribe_callback(char *payload, + int length, + task_id *task_id, + int *state, + db_client_id *local_scheduler_id, + task_spec **spec) { + /* Note that the state is padded with spaces to consist of precisely two + * characters. */ + int task_spec_payload_size = + length - sizeof(*task_id) - 1 - 2 - 1 - sizeof(*local_scheduler_id) - 1; + int offset = 0; + /* Read in the task ID. */ + memcpy(task_id, &payload[offset], sizeof(*task_id)); + offset += sizeof(*task_id); + /* Read in a space. */ + char *space_str = " "; + CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0); + offset += strlen(space_str); + /* Read in the state, which is an integer left-padded with spaces to two + * characters. */ + CHECK(sscanf(&payload[offset], "%2d", state) == 1); + offset += 2; + /* Read in a space. */ + CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0); + offset += strlen(space_str); + /* Read in the local scheduler ID. */ + memcpy(local_scheduler_id, &payload[offset], sizeof(*local_scheduler_id)); + offset += sizeof(*local_scheduler_id); + /* Read in a space. */ + CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0); + offset += strlen(space_str); + /* Read in the task spec. */ + *spec = malloc(task_spec_payload_size); + memcpy(*spec, &payload[offset], task_spec_payload_size); + CHECK(task_spec_size(*spec) == task_spec_payload_size); } void redis_task_table_subscribe_callback(redisAsyncContext *c, @@ -968,53 +868,75 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c, redisReply *reply = r; CHECK(reply->type == REDIS_REPLY_ARRAY); - CHECK(reply->elements > 2); - /* First entry is message type, then possibly the regex we psubscribed to, - * then topic, then payload. */ + /* The number of elements is 3 for a reply to SUBSCRIBE, and 4 for a reply to + * PSUBSCRIBE. */ + CHECKM(reply->elements == 3 || reply->elements == 4, "reply->elements is %zu", + reply->elements); + /* The first element is the message type and the last entry is the payload. + * The middle one or middle two elements describe the channel that was + * published on. */ + redisReply *message_type = reply->element[0]; redisReply *payload = reply->element[reply->elements - 1]; - /* If this condition is true, we got the initial message that acknowledged the - * subscription. */ - if (payload->str == NULL) { - if (callback_data->done_callback) { + if (strcmp(message_type->str, "message") == 0 || + strcmp(message_type->str, "pmessage") == 0) { + /* Handle a task table event. Parse the payload and call the callback. */ + task_table_subscribe_data *data = callback_data->data; + /* Read out the information from the payload. */ + task_id task_id; + int state; + db_client_id local_scheduler_id; + task_spec *spec; + parse_task_table_subscribe_callback(payload->str, payload->len, &task_id, + &state, &local_scheduler_id, &spec); + task *task = alloc_task(spec, state, local_scheduler_id); + free(spec); + /* Call the subscribe callback if there is one. */ + if (data->subscribe_callback != NULL) { + data->subscribe_callback(task, data->subscribe_context); + } + free_task(task); + } else if (strcmp(message_type->str, "subscribe") == 0 || + strcmp(message_type->str, "psubscribe") == 0) { + /* If this condition is true, we got the initial message that acknowledged + * the subscription. */ + if (callback_data->done_callback != NULL) { task_table_done_callback done_callback = callback_data->done_callback; done_callback(callback_data->id, callback_data->user_context); } /* Note that we do not destroy the callback data yet because the * subscription callback needs this data. */ event_loop_remove_timer(db->loop, callback_data->timer_id); - return; + } else { + LOG_FATAL( + "Unexpected reply type from task table subscribe. Message type is %s.", + message_type->str); } - /* Otherwise, parse the task and call the callback. */ - task_table_subscribe_data *data = callback_data->data; - - task *task = malloc(payload->len); - memcpy(task, payload->str, payload->len); - if (data->subscribe_callback) { - data->subscribe_callback(task, data->subscribe_context); - } - free_task(task); } void redis_task_table_subscribe(table_callback_data *callback_data) { db_handle *db = callback_data->db_handle; task_table_subscribe_data *data = callback_data->data; - int status = REDIS_OK; - if (IS_NIL_ID(data->node)) { + /* TASK_CHANNEL_PREFIX is defined in ray_redis_module.c and must be kept in + * sync with that file. */ + const char *TASK_CHANNEL_PREFIX = "TT:"; + int status; + if (IS_NIL_ID(data->local_scheduler_id)) { /* TODO(swang): Implement the state_filter by translating the bitmask into * a Redis key-matching pattern. */ - status = - redisAsyncCommand(db->sub_context, redis_task_table_subscribe_callback, - (void *) callback_data->timer_id, - "PSUBSCRIBE task:*:%d", data->state_filter); - } else { - node_id node = data->node; status = redisAsyncCommand( db->sub_context, redis_task_table_subscribe_callback, - (void *) callback_data->timer_id, "SUBSCRIBE task:%b:%d", - (char *) node.id, sizeof(node.id), data->state_filter); + (void *) callback_data->timer_id, "PSUBSCRIBE %s*:%2d", + TASK_CHANNEL_PREFIX, data->state_filter); + } else { + db_client_id local_scheduler_id = data->local_scheduler_id; + status = redisAsyncCommand( + db->sub_context, redis_task_table_subscribe_callback, + (void *) callback_data->timer_id, "SUBSCRIBE %s%b:%2d", + TASK_CHANNEL_PREFIX, (char *) local_scheduler_id.id, + sizeof(local_scheduler_id.id), data->state_filter); } if ((status == REDIS_ERR) || db->sub_context->err) { - LOG_REDIS_DEBUG(db->sub_context, "error in task_table_register_callback"); + LOG_REDIS_DEBUG(db->sub_context, "error in redis_task_table_subscribe"); } } diff --git a/src/common/state/task_table.c b/src/common/state/task_table.c index e139662d9..4e67e4bf6 100644 --- a/src/common/state/task_table.c +++ b/src/common/state/task_table.c @@ -32,7 +32,7 @@ void task_table_update(db_handle *db_handle, /* TODO(swang): A corresponding task_table_unsubscribe. */ void task_table_subscribe(db_handle *db_handle, - node_id node, + db_client_id local_scheduler_id, scheduling_state state_filter, task_table_subscribe_callback subscribe_callback, void *subscribe_context, @@ -41,11 +41,11 @@ void task_table_subscribe(db_handle *db_handle, void *user_context) { task_table_subscribe_data *sub_data = malloc(sizeof(task_table_subscribe_data)); - sub_data->node = node; + sub_data->local_scheduler_id = local_scheduler_id; sub_data->state_filter = state_filter; sub_data->subscribe_callback = subscribe_callback; sub_data->subscribe_context = subscribe_context; - init_table_callback(db_handle, node, __func__, sub_data, retry, done_callback, - redis_task_table_subscribe, user_context); + init_table_callback(db_handle, local_scheduler_id, __func__, sub_data, retry, + done_callback, redis_task_table_subscribe, user_context); } diff --git a/src/common/state/task_table.h b/src/common/state/task_table.h index c7b788a57..2cf8b4eef 100644 --- a/src/common/state/task_table.h +++ b/src/common/state/task_table.h @@ -6,21 +6,26 @@ #include "task.h" /** - * The task table is a message bus that is used for all communication between - * local and global schedulers (and also persisted to the state database). - * Here are examples of events that are recorded by the task table: + * The task table is a message bus that is used for communication between local + * and global schedulers (and also persisted to the state database). Here are + * examples of events that are recorded by the task table: * - * 1) local scheduler writes when it submits a task to the global scheduler; - * 2) global scheduler reads it to get the task submitted by local schedulers; - * 3) global scheduler writes it when assigning the task to a local scheduler; - * 4) local scheduler reads it to get its tasks assigned by global scheduler; - * 5) local scheduler writes it when a task finishes execution; - * 6) global scheduler reads it to get the tasks that have finished; */ + * 1) Local schedulers write to it when submitting a task to the global + * scheduler. + * 2) The global scheduler subscribes to updates to the task table to get tasks + * submitted by local schedulers. + * 3) The global scheduler writes to it when assigning a task to a local + * scheduler. + * 4) Local schedulers subscribe to updates to the task table to get tasks + * assigned to them by the global scheduler. + * 5) Local schedulers write to it when a task finishes execution. + */ /* Callback called when a task table write operation completes. */ typedef void (*task_table_done_callback)(task_id task_id, void *user_context); -/* Callback called when a task table read operation completes. */ +/* Callback called when a task table read operation completes. If the task ID + * was not in the task table, then the task pointer will be NULL. */ typedef void (*task_table_get_callback)(task *task, void *user_context); /** @@ -41,9 +46,9 @@ void task_table_get_task(db_handle *db, void *user_context); /** - * Add a task entry, including task spec and scheduling information, to the - * task table. This will overwrite any task already in the task table with the - * same task ID. + * Add a task entry, including task spec and scheduling information, to the task + * table. This will overwrite any task already in the task table with the same + * task ID. * * @param db_handle Database handle. * @param task The task entry to add to the table. @@ -93,15 +98,16 @@ typedef void (*task_table_subscribe_callback)(task *task, void *user_context); * Register a callback for a task event. An event is any update of a task in * the task table, produced by task_table_add_task or task_table_add_task. * Events include changes to the task's scheduling state or changes to the - * task's node location. + * task's local scheduler ID. * * @param db_handle Database handle. * @param subscribe_callback Callback that will be called when the task table is * updated. * @param subscribe_context Context that will be passed into the * subscribe_callback. - * @param node Node whose events we want to listen to. If you want to register - * to updates from all nodes, set node = NIL_ID. + * @param local_scheduler_id The db_client_id of the local scheduler whose + * events we want to listen to. If you want to subscribe to updates from + * all local schedulers, pass in NIL_ID. * @param state_filter Flags for events we want to listen to. If you want * to listen to all events, use state_filter = TASK_WAITING | * TASK_SCHEDULED | TASK_RUNNING | TASK_DONE. @@ -112,7 +118,7 @@ typedef void (*task_table_subscribe_callback)(task *task, void *user_context); * @return Void. */ void task_table_subscribe(db_handle *db_handle, - node_id node, + db_client_id local_scheduler_id, scheduling_state state_filter, task_table_subscribe_callback subscribe_callback, void *subscribe_context, @@ -123,7 +129,7 @@ void task_table_subscribe(db_handle *db_handle, /* Data that is needed to register task table subscribe callbacks with the state * database. */ typedef struct { - node_id node; + db_client_id local_scheduler_id; scheduling_state state_filter; task_table_subscribe_callback subscribe_callback; void *subscribe_context; diff --git a/src/common/task.c b/src/common/task.c index 544ccb520..577816251 100644 --- a/src/common/task.c +++ b/src/common/task.c @@ -298,24 +298,18 @@ void print_task(task_spec *spec, UT_string *output) { struct task_impl { scheduling_state state; - node_id node; + db_client_id local_scheduler_id; task_spec spec; }; -bool node_ids_equal(node_id first_id, node_id second_id) { - return UNIQUE_ID_EQ(first_id, second_id); -} - -bool node_id_is_nil(node_id id) { - return node_ids_equal(id, NIL_NODE_ID); -} - -task *alloc_task(task_spec *spec, scheduling_state state, node_id node) { +task *alloc_task(task_spec *spec, + scheduling_state state, + db_client_id local_scheduler_id) { int64_t size = sizeof(task) - sizeof(task_spec) + task_spec_size(spec); task *result = malloc(size); memset(result, 0, size); result->state = state; - result->node = node; + result->local_scheduler_id = local_scheduler_id; memcpy(&result->spec, spec, task_spec_size(spec)); return result; } @@ -347,12 +341,12 @@ void task_set_state(task *task, scheduling_state state) { task->state = state; } -node_id task_node(task *task) { - return task->node; +db_client_id task_local_scheduler(task *task) { + return task->local_scheduler_id; } -void task_set_node(task *task, node_id node) { - task->node = node; +void task_set_local_scheduler(task *task, db_client_id local_scheduler_id) { + task->local_scheduler_id = local_scheduler_id; } task_spec *task_task_spec(task *task) { diff --git a/src/common/task.h b/src/common/task.h index d9a2ca41f..e8989a8b6 100644 --- a/src/common/task.h +++ b/src/common/task.h @@ -16,7 +16,6 @@ #define NIL_TASK_ID NIL_ID #define NIL_FUNCTION_ID NIL_ID -#define NIL_NODE_ID NIL_ID typedef unique_id function_id; @@ -28,9 +27,6 @@ typedef unique_id task_id; * particular execution of the task. */ typedef unique_id task_iid; -/** The node id is an identifier for the node the task is scheduled on. */ -typedef unique_id node_id; - /** * ==== Task specifications ==== * Contain all the information neccessary to execute the @@ -248,9 +244,9 @@ void print_task(task_spec *spec, UT_string *output); /** * ==== Task ==== - * Contains information about a scheduled task: The task iid, - * the task specification and the task status (WAITING, SCHEDULED, - * RUNNING, DONE) and which node the task is scheduled on. + * Contains information about a scheduled task: The task specification, the task + * schedulign state (WAITING, SCHEDULED, RUNNING, DONE), and which local + * scheduler the task is scheduled on. */ /** The scheduling_state can be used as a flag when we are listening @@ -262,27 +258,9 @@ typedef enum { TASK_STATUS_DONE = 8 } scheduling_state; -/** - * Compare two node IDs. - * - * @param first_id The first node ID to compare. - * @param second_id The first node ID to compare. - * @return True if the node IDs are the same and false - * otherwise. - */ -bool node_ids_equal(node_id first_id, node_id second_id); - -/** - * Compare a node ID to the nil ID. - * - * @param id The node ID to compare to nil. - * @return True if the node ID is equal to nil. - */ -bool node_id_is_nil(node_id id); - -/** A task is an execution of a task specification. It has a state of - * execution (see scheduling_state) and a node it is scheduled on or running - * on. */ +/** A task is an execution of a task specification. It has a state of execution + * (see scheduling_state) and the ID of the local scheduler it is scheduled on + * or running on. */ typedef struct task_impl task; /** @@ -290,9 +268,12 @@ typedef struct task_impl task; * * @param spec The task spec for the new task. * @param state The scheduling state for the new task. - * @param node The ID of the node that the task is scheduled on, if any. + * @param local_scheduler_id The ID of the local scheduler that the task is + * scheduled on, if any. */ -task *alloc_task(task_spec *spec, scheduling_state state, node_id node); +task *alloc_task(task_spec *spec, + scheduling_state state, + db_client_id local_scheduler_id); /** * Create a copy of the task. Must be freed with free_task after use. @@ -311,11 +292,11 @@ scheduling_state task_state(task *task); /** Update the schedule state of the task. */ void task_set_state(task *task, scheduling_state state); -/** Node this task has been assigned to or is running on. */ -node_id task_node(task *task); +/** Local scheduler this task has been assigned to or is running on. */ +db_client_id task_local_scheduler(task *task); -/** Set the node for this task. */ -void task_set_node(task *task, node_id node); +/** Set the local scheduler ID for this task. */ +void task_set_local_scheduler(task *task, db_client_id local_scheduler_id); /** Task specification of this task. */ task_spec *task_task_spec(task *task); @@ -333,7 +314,7 @@ void free_task(task *task); typedef struct { scheduling_state state; - node_id node; + db_client_id local_scheduler_id; } task_update; #endif diff --git a/src/common/test/db_tests.c b/src/common/test/db_tests.c index 52886ec7c..0ee27f128 100644 --- a/src/common/test/db_tests.c +++ b/src/common/test/db_tests.c @@ -127,7 +127,8 @@ int64_t task_table_delayed_add_task(event_loop *loop, .timeout = TIMEOUT, .fail_callback = task_table_test_fail_callback, }; - task_table_add_task(db, task_table_test_task, &retry, NULL, (void *) loop); + task_table_add_task(db, copy_task(task_table_test_task), &retry, NULL, + (void *) loop); return EVENT_LOOP_TIMER_DONE; } @@ -147,21 +148,23 @@ TEST task_table_test(void) { db_handle *db = db_connect("127.0.0.1", 6379, "local_scheduler", "127.0.0.1", 0, NULL); db_attach(db, loop, false); - node_id node = globally_unique_id(); + db_client_id local_scheduler_id = globally_unique_id(); task_spec *spec = example_task_spec(1, 1); - task_table_test_task = alloc_task(spec, TASK_STATUS_SCHEDULED, node); + task_table_test_task = + alloc_task(spec, TASK_STATUS_SCHEDULED, local_scheduler_id); free_task_spec(spec); retry_info retry = { .num_retries = NUM_RETRIES, .timeout = TIMEOUT, .fail_callback = task_table_test_fail_callback, }; - task_table_subscribe(db, node, TASK_STATUS_SCHEDULED, + task_table_subscribe(db, local_scheduler_id, TASK_STATUS_SCHEDULED, task_table_test_callback, (void *) loop, &retry, NULL, (void *) loop); event_loop_add_timer( loop, 200, (event_loop_timer_handler) task_table_delayed_add_task, db); event_loop_run(loop); + free_task(task_table_test_task); db_disconnect(db); destroy_outstanding_callbacks(loop); event_loop_destroy(loop); @@ -181,7 +184,7 @@ TEST task_table_all_test(void) { db_connect("127.0.0.1", 6379, "local_scheduler", "127.0.0.1", 0, NULL); db_attach(db, loop, false); task_spec *spec = example_task_spec(1, 1); - /* Schedule two tasks on different nodes. */ + /* Schedule two tasks on different local local schedulers. */ task *task1 = alloc_task(spec, TASK_STATUS_SCHEDULED, globally_unique_id()); task *task2 = alloc_task(spec, TASK_STATUS_SCHEDULED, globally_unique_id()); retry_info retry = { @@ -193,8 +196,8 @@ TEST task_table_all_test(void) { NULL); event_loop_run(loop); /* TODO(pcm): Get rid of this sleep once the robust pubsub is implemented. */ - task_table_update(db, task1, &retry, NULL, NULL); - task_table_update(db, task2, &retry, NULL, NULL); + task_table_add_task(db, task1, &retry, NULL, NULL); + task_table_add_task(db, task2, &retry, NULL, NULL); event_loop_add_timer(loop, 200, (event_loop_timer_handler) timeout_handler, NULL); event_loop_run(loop); diff --git a/src/common/test/object_table_tests.c b/src/common/test/object_table_tests.c index fff2c8622..39c337ff2 100644 --- a/src/common/test/object_table_tests.c +++ b/src/common/test/object_table_tests.c @@ -101,7 +101,7 @@ void new_object_no_task_lookup_callback(object_id object_id, } void new_object_no_task_callback(object_id object_id, void *user_context) { - CHECK(node_ids_equal(object_id, new_object_id)); + CHECK(object_ids_equal(object_id, new_object_id)); retry_info retry = { .num_retries = 5, .timeout = 100, diff --git a/src/common/test/task_table_tests.c b/src/common/test/task_table_tests.c index 0c2cbc6a0..06df8514b 100644 --- a/src/common/test/task_table_tests.c +++ b/src/common/test/task_table_tests.c @@ -184,8 +184,8 @@ TEST publish_timeout_test(void) { retry_info retry = { .num_retries = 5, .timeout = 100, .fail_callback = publish_fail_callback, }; - task_table_update(db, task, &retry, publish_done_callback, - (void *) publish_timeout_context); + task_table_add_task(db, task, &retry, publish_done_callback, + (void *) publish_timeout_context); /* Disconnect the database to see if the publish times out. */ close(db->context->c.fd); aeProcessEvents(g_loop, AE_TIME_EVENTS); @@ -295,8 +295,8 @@ TEST publish_retry_test(void) { .timeout = 100, .fail_callback = publish_retry_fail_callback, }; - task_table_update(db, task, &retry, publish_retry_done_callback, - (void *) publish_retry_context); + task_table_add_task(db, task, &retry, publish_retry_done_callback, + (void *) publish_retry_context); /* Disconnect the database to see if the publish times out. */ close(db->sub_context->c.fd); /* Install handler for reconnecting the database. */ @@ -389,8 +389,8 @@ TEST publish_late_test(void) { .timeout = 0, .fail_callback = publish_late_fail_callback, }; - task_table_update(db, task, &retry, publish_late_done_callback, - (void *) publish_late_context); + task_table_add_task(db, task, &retry, publish_late_done_callback, + (void *) publish_late_context); /* Install handler for terminating the event loop. */ event_loop_add_timer(g_loop, 750, (event_loop_timer_handler) terminate_event_loop_callback, diff --git a/src/global_scheduler/global_scheduler.c b/src/global_scheduler/global_scheduler.c index a8c9ad829..66ac936d4 100644 --- a/src/global_scheduler/global_scheduler.c +++ b/src/global_scheduler/global_scheduler.c @@ -20,12 +20,12 @@ UT_icd local_scheduler_icd = {sizeof(local_scheduler), NULL, NULL, NULL}; void assign_task_to_local_scheduler(global_scheduler_state *state, task *task, - node_id node_id) { + db_client_id local_scheduler_id) { char id_string[ID_STRING_SIZE]; - LOG_DEBUG("assigning task to node_id = %s", - object_id_to_string(node_id, id_string, ID_STRING_SIZE)); + LOG_DEBUG("assigning task to local_scheduler_id = %s", + object_id_to_string(local_scheduler_id, id_string, ID_STRING_SIZE)); task_set_state(task, TASK_STATUS_SCHEDULED); - task_set_node(task, node_id); + task_set_local_scheduler(task, local_scheduler_id); retry_info retry = { .num_retries = 0, .timeout = 100, .fail_callback = NULL, }; diff --git a/src/global_scheduler/global_scheduler.h b/src/global_scheduler/global_scheduler.h index dc27f03e3..6239896c2 100644 --- a/src/global_scheduler/global_scheduler.h +++ b/src/global_scheduler/global_scheduler.h @@ -54,6 +54,6 @@ typedef struct { void assign_task_to_local_scheduler(global_scheduler_state *state, task *task, - node_id node_id); + db_client_id local_scheduler_id); #endif /* GLOBAL_SCHEDULER_H */ diff --git a/src/global_scheduler/test/test.py b/src/global_scheduler/test/test.py index b4255a33d..508c683d9 100644 --- a/src/global_scheduler/test/test.py +++ b/src/global_scheduler/test/test.py @@ -28,9 +28,10 @@ TASK_STATUS_SCHEDULED = 2 TASK_STATUS_RUNNING = 4 TASK_STATUS_DONE = 8 -# DB_CLIENT_PREFIX is an implementation detail of ray_redis_module.c, so this +# These constants are an implementation detail of ray_redis_module.c, so this # must be kept in sync with that file. DB_CLIENT_PREFIX = "CL:" +TASK_PREFIX = "TT:" def random_task_id(): return photon.ObjectID(np.random.bytes(ID_SIZE)) @@ -155,7 +156,7 @@ class TestGlobalScheduler(unittest.TestCase): # local scheduler num_retries = 10 while num_retries > 0: - task_entries = self.redis_client.keys("task*") + task_entries = self.redis_client.keys("{}*".format(TASK_PREFIX)) self.assertLessEqual(len(task_entries), 1) if len(task_entries) == 1: task_contents = self.redis_client.hgetall(task_entries[0]) @@ -197,7 +198,7 @@ class TestGlobalScheduler(unittest.TestCase): num_retries = 10 num_tasks_done = 0 while num_retries > 0: - task_entries = self.redis_client.keys("task*") + task_entries = self.redis_client.keys("{}*".format(TASK_PREFIX)) self.assertLessEqual(len(task_entries), num_tasks) # First, check if all tasks made it to Redis. if len(task_entries) == num_tasks: diff --git a/src/photon/photon_scheduler.c b/src/photon/photon_scheduler.c index 66417c3b5..b47019e79 100644 --- a/src/photon/photon_scheduler.c +++ b/src/photon/photon_scheduler.c @@ -109,8 +109,13 @@ void assign_task_to_worker(local_scheduler_state *state, if (state->db != NULL) { task *task = alloc_task(spec, TASK_STATUS_RUNNING, get_db_client_id(state->db)); - task_table_update(state->db, task, (retry_info *) &photon_retry, NULL, - NULL); + if (from_global_scheduler) { + task_table_update(state->db, task, (retry_info *) &photon_retry, NULL, + NULL); + } else { + task_table_add_task(state->db, task, (retry_info *) &photon_retry, NULL, + NULL); + } /* Record which task this worker is executing. This will be freed in * process_message when the worker sends a GET_TASK message to the local * scheduler. */