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

Use redismodules for task table and result table. (#156)

Browse Source 
* 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.
This commit is contained in:
Robert Nishihara
2016-12-25 23:57:05 -08:00
committed by Philipp Moritz
parent d6695c867a
commit 985c424172
14 changed files with 456 additions and 500 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/redis_module/ray_redis_module.c
+127 -98
View File
@@ -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 <object id>
*
* @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 <task ID> <state> <node ID> <task spec>
* RAY.TASK_TABLE_ADD <task ID> <state> <local scheduler ID> <task spec>
*
* @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 <task ID> <state> <node ID>
* RAY.TASK_TABLE_UPDATE <task ID> <state> <local scheduler ID>
*
* @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 <task ID>
* RAY.TASK_TABLE_GET <task ID>
*
* @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;
}
src/common/redis_module/runtest.py
+19 -4
View File
@@ -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
src/common/state/redis.c
+224 -302
View File
@@ -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
* "<task ID> <state> <local scheduler ID> <task specification>". 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");
}
}
src/common/state/task_table.c
+4 -4
View File
@@ -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);
}
src/common/state/task_table.h
+24 -18
View File
@@ -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;
src/common/task.c
+9 -15
View File
@@ -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) {
src/common/task.h
+16 -35
View File
@@ -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
src/common/test/db_tests.c
+10 -7
View File
@@ -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);
src/common/test/object_table_tests.c
+1 -1
View File
@@ -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,
src/common/test/task_table_tests.c
+6 -6
View File
@@ -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,
src/global_scheduler/global_scheduler.c
+4 -4
View File
@@ -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,
};
src/global_scheduler/global_scheduler.h
+1 -1
View File
@@ -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 */
src/global_scheduler/test/test.py
+4 -3
View File
@@ -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:
src/photon/photon_scheduler.c
+7 -2
View File
@@ -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. */