wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-08 11:25:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

Convert task_spec to flatbuffers (#255)

Browse Source 
* convert Ray to C++

* convert task_spec to flatbuffers

* fix

* it compiles

* latest

* tests are passing

* task2 -> task

* fix

* fix

* fix

* fix

* fix

* linting

* fix valgrind

* upgrade flatbuffers

* use debug mode for valgrind

* fix naming and comments

* downgrade flatbuffers

* fix linting

* reintroduce TaskSpec_free

* rename TaskSpec -> TaskInfo

* refactoring

* linting
This commit is contained in:
Philipp Moritz
2017-03-05 02:05:02 -08:00
committed by Robert Nishihara
parent 65a8659f3d
commit 0b8d279ef2
36 changed files with 1054 additions and 931 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/local_scheduler/CMakeLists.txt
+1 -1
View File
@@ -35,7 +35,7 @@ endif(APPLE)
add_library(local_scheduler_client STATIC local_scheduler_client.cc)
target_link_libraries(local_scheduler_library local_scheduler_client ${COMMON_LIB} ${PYTHON_LIBRARIES})
target_link_libraries(local_scheduler_library local_scheduler_client common ${PYTHON_LIBRARIES})
add_executable(local_scheduler local_scheduler.cc local_scheduler_algorithm.cc)
target_link_libraries(local_scheduler local_scheduler_client common ${HIREDIS_LIB} plasma_lib)
src/local_scheduler/local_scheduler.cc
+41 -37
View File
@@ -37,20 +37,20 @@ UT_icd byte_icd = {sizeof(uint8_t), NULL, NULL, NULL};
* @return Void.
*/
void print_resource_info(const LocalSchedulerState *state,
const task_spec *spec) {
const TaskSpec *spec) {
#if RAY_COMMON_LOG_LEVEL <= RAY_COMMON_DEBUG
/* Print information about available and requested resources. */
char buftotal[256], bufavail[256], bufresreq[256];
snprintf(bufavail, sizeof(bufavail), "%8.4f %8.4f",
state->dynamic_resources[CPU_RESOURCE_INDEX],
state->dynamic_resources[GPU_RESOURCE_INDEX]);
state->dynamic_resources[ResourceIndex_CPU],
state->dynamic_resources[ResourceIndex_GPU]);
snprintf(buftotal, sizeof(buftotal), "%8.4f %8.4f",
state->static_resources[CPU_RESOURCE_INDEX],
state->static_resources[GPU_RESOURCE_INDEX]);
state->static_resources[ResourceIndex_CPU],
state->static_resources[ResourceIndex_GPU]);
if (spec) {
snprintf(bufresreq, sizeof(bufresreq), "%8.4f %8.4f",
task_spec_get_required_resource(spec, CPU_RESOURCE_INDEX),
task_spec_get_required_resource(spec, GPU_RESOURCE_INDEX));
task_spec_get_required_resource(spec, ResourceIndex_CPU),
task_spec_get_required_resource(spec, ResourceIndex_GPU));
}
LOG_DEBUG("Resources: [total=%s][available=%s][requested=%s]", buftotal,
bufavail, spec ? bufresreq : "n/a");
@@ -124,7 +124,7 @@ void kill_worker(LocalSchedulerClient *worker, bool cleanup) {
/* Clean up the task in progress. */
if (worker->task_in_progress) {
/* Return the resources that the worker was using. */
task_spec *spec = Task_task_spec(worker->task_in_progress);
TaskSpec *spec = Task_task_spec(worker->task_in_progress);
update_dynamic_resources(state, spec, true);
/* Update the task table to reflect that the task failed to complete. */
if (state->db != NULL) {
@@ -388,7 +388,7 @@ LocalSchedulerState *LocalSchedulerState_init(
utarray_new(state->input_buffer, &byte_icd);
/* Initialize resource vectors. */
for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
for (int i = 0; i < ResourceIndex_MAX; i++) {
state->static_resources[i] = state->dynamic_resources[i] =
static_resource_conf[i];
}
@@ -405,10 +405,10 @@ LocalSchedulerState *LocalSchedulerState_init(
}
void update_dynamic_resources(LocalSchedulerState *state,
task_spec *spec,
TaskSpec *spec,
bool return_resources) {
for (int i = 0; i < MAX_RESOURCE_INDEX; ++i) {
double resource = task_spec_get_required_resource(spec, i);
for (int i = 0; i < ResourceIndex_MAX; ++i) {
double resource = TaskSpec_get_required_resource(spec, i);
if (!return_resources) {
/* If we are not returning resources, we are leasing them, so we want to
* subtract the resource quantities from our accounting. */
@@ -428,9 +428,10 @@ void update_dynamic_resources(LocalSchedulerState *state,
}
void assign_task_to_worker(LocalSchedulerState *state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
LocalSchedulerClient *worker) {
if (write_message(worker->sock, EXECUTE_TASK, task_spec_size(spec),
if (write_message(worker->sock, EXECUTE_TASK, task_spec_size,
(uint8_t *) spec) < 0) {
if (errno == EPIPE || errno == EBADF) {
/* TODO(rkn): If this happens, the task should be added back to the task
@@ -447,7 +448,7 @@ void assign_task_to_worker(LocalSchedulerState *state,
/* Resource accounting:
* Update dynamic resource vector in the local scheduler state. */
update_dynamic_resources(state, spec, false);
Task *task = Task_alloc(spec, TASK_STATUS_RUNNING,
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_RUNNING,
state->db ? get_db_client_id(state->db) : NIL_ID);
/* 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
@@ -497,16 +498,17 @@ void reconstruct_task_update_callback(Task *task, void *user_context) {
/* Otherwise, the test-and-set succeeded, so resubmit the task for execution
* to ensure that reconstruction will happen. */
LocalSchedulerState *state = (LocalSchedulerState *) user_context;
task_spec *spec = Task_task_spec(task);
TaskSpec *spec = Task_task_spec(task);
/* If the task is an actor task, then we currently do not reconstruct it.
* TODO(rkn): Handle this better. */
CHECK(ActorID_equal(task_spec_actor_id(spec), NIL_ACTOR_ID));
CHECK(ActorID_equal(TaskSpec_actor_id(spec), NIL_ACTOR_ID));
/* Resubmit the task. */
handle_task_submitted(state, state->algorithm_state, spec);
handle_task_submitted(state, state->algorithm_state, spec,
Task_task_spec_size(task));
/* Recursively reconstruct the task's inputs, if necessary. */
for (int64_t i = 0; i < task_num_args(spec); ++i) {
if (task_arg_type(spec, i) == ARG_BY_REF) {
ObjectID arg_id = task_arg_id(spec, i);
for (int64_t i = 0; i < TaskSpec_num_args(spec); ++i) {
if (TaskSpec_arg_by_ref(spec, i)) {
ObjectID arg_id = TaskSpec_arg_id(spec, i);
reconstruct_object(state, arg_id);
}
}
@@ -601,22 +603,22 @@ void process_message(event_loop *loop,
switch (type) {
case SUBMIT_TASK: {
task_spec *spec = (task_spec *) utarray_front(state->input_buffer);
TaskSpec *spec = (TaskSpec *) utarray_front(state->input_buffer);
/* Update the result table, which holds mappings of object ID -> ID of the
* task that created it. */
if (state->db != NULL) {
TaskID task_id = task_spec_id(spec);
for (int64_t i = 0; i < task_num_returns(spec); ++i) {
ObjectID return_id = task_return(spec, i);
TaskID task_id = TaskSpec_task_id(spec);
for (int64_t i = 0; i < TaskSpec_num_returns(spec); ++i) {
ObjectID return_id = TaskSpec_return(spec, i);
result_table_add(state->db, return_id, task_id, NULL, NULL, NULL);
}
}
/* Handle the task submission. */
if (ActorID_equal(task_spec_actor_id(spec), NIL_ACTOR_ID)) {
handle_task_submitted(state, state->algorithm_state, spec);
if (ActorID_equal(TaskSpec_actor_id(spec), NIL_ACTOR_ID)) {
handle_task_submitted(state, state->algorithm_state, spec, length);
} else {
handle_actor_task_submitted(state, state->algorithm_state, spec);
handle_actor_task_submitted(state, state->algorithm_state, spec, length);
}
} break;
@@ -693,7 +695,7 @@ void process_message(event_loop *loop,
case GET_TASK: {
/* If this worker reports a completed task: account for resources. */
if (worker->task_in_progress != NULL) {
task_spec *spec = Task_task_spec(worker->task_in_progress);
TaskSpec *spec = Task_task_spec(worker->task_in_progress);
/* Return dynamic resources back for the task in progress. */
update_dynamic_resources(state, spec, true);
/* If we're connected to Redis, update tables. */
@@ -798,14 +800,16 @@ void signal_handler(int signal) {
/* End of the cleanup code. */
void handle_task_scheduled_callback(Task *original_task, void *user_context) {
task_spec *spec = Task_task_spec(original_task);
if (ActorID_equal(task_spec_actor_id(spec), NIL_ACTOR_ID)) {
TaskSpec *spec = Task_task_spec(original_task);
if (ActorID_equal(TaskSpec_actor_id(spec), NIL_ACTOR_ID)) {
/* This task does not involve an actor. Handle it normally. */
handle_task_scheduled(g_state, g_state->algorithm_state, spec);
handle_task_scheduled(g_state, g_state->algorithm_state, spec,
Task_task_spec_size(original_task));
} else {
/* This task involves an actor. Call the scheduling algorithm's actor
* handler. */
handle_actor_task_scheduled(g_state, g_state->algorithm_state, spec);
handle_actor_task_scheduled(g_state, g_state->algorithm_state, spec,
Task_task_spec_size(original_task));
}
}
@@ -932,7 +936,7 @@ int main(int argc, char *argv[]) {
char *node_ip_address = NULL;
/* Comma-separated list of configured resource capabilities for this node. */
char *static_resource_list = NULL;
double static_resource_conf[MAX_RESOURCE_INDEX];
double static_resource_conf[ResourceIndex_MAX];
/* The command to run when starting new workers. */
char *start_worker_command = NULL;
/* The number of workers to start. */
@@ -978,15 +982,15 @@ int main(int argc, char *argv[]) {
if (!static_resource_list) {
/* Use defaults for this node's static resource configuration. */
memset(&static_resource_conf[0], 0, sizeof(static_resource_conf));
static_resource_conf[CPU_RESOURCE_INDEX] = DEFAULT_NUM_CPUS;
static_resource_conf[GPU_RESOURCE_INDEX] = DEFAULT_NUM_GPUS;
static_resource_conf[ResourceIndex_CPU] = DEFAULT_NUM_CPUS;
static_resource_conf[ResourceIndex_GPU] = DEFAULT_NUM_GPUS;
} else {
/* Tokenize the string. */
const char delim[2] = ",";
char *token;
int idx = 0; /* Index into the resource vector. */
token = strtok(static_resource_list, delim);
while (token != NULL && idx < MAX_RESOURCE_INDEX) {
while (token != NULL && idx < ResourceIndex_MAX) {
static_resource_conf[idx++] = atoi(token);
/* Attempt to get the next token. */
token = strtok(NULL, delim);
src/local_scheduler/local_scheduler.h
+4 -3
View File
@@ -39,7 +39,8 @@ void new_client_connection(event_loop *loop,
* @return Void.
*/
void assign_task_to_worker(LocalSchedulerState *state,
task_spec *task,
TaskSpec *task,
int64_t task_spec_size,
LocalSchedulerClient *worker);
/**
@@ -69,7 +70,7 @@ void process_plasma_notification(event_loop *loop,
*/
void reconstruct_object(LocalSchedulerState *state, ObjectID object_id);
void print_resource_info(const LocalSchedulerState *s, const task_spec *spec);
void print_resource_info(const LocalSchedulerState *s, const TaskSpec *spec);
/**
* Kill a worker.
@@ -106,7 +107,7 @@ void start_worker(LocalSchedulerState *state, ActorID actor_id);
* @return Void.
*/
void update_dynamic_resources(LocalSchedulerState *state,
task_spec *spec,
TaskSpec *spec,
bool return_resources);
/** The following methods are for testing purposes only. */
src/local_scheduler/local_scheduler_algorithm.cc
+113 -78
View File
@@ -16,7 +16,8 @@ void remove_actor(SchedulingAlgorithmState *algorithm_state, ActorID actor_id);
typedef struct task_queue_entry {
/** The task that is queued. */
task_spec *spec;
TaskSpec *spec;
int64_t task_spec_size;
struct task_queue_entry *prev;
struct task_queue_entry *next;
} task_queue_entry;
@@ -39,7 +40,9 @@ UT_icd task_queue_entry_icd = {sizeof(task_queue_entry *), NULL, NULL, NULL};
/** This is used to define the queue of actor task specs for which the
* corresponding local scheduler is unknown. */
UT_icd task_spec_icd = {sizeof(task_spec *), NULL, NULL, NULL};
UT_icd task_spec_icd = {sizeof(TaskSpec *), NULL, NULL, NULL};
/** This is used to keep track of task spec sizes in the above queue. */
UT_icd task_spec_size_icd = {sizeof(int64_t), NULL, NULL, NULL};
/** This is used to define the queue of available workers. */
UT_icd worker_icd = {sizeof(LocalSchedulerClient *), NULL, NULL, NULL};
@@ -83,6 +86,8 @@ struct SchedulingAlgorithmState {
* about a new local scheduler arrives, we will resubmit all of these tasks
* locally. */
UT_array *cached_submitted_actor_tasks;
/** An array of task sizes of cached_submitted_actor_tasks. */
UT_array *cached_submitted_actor_task_sizes;
/** An array of pointers to workers in the worker pool. These are workers
* that have registered a PID with us and that are now waiting to be
* assigned a task to execute. */
@@ -119,6 +124,9 @@ SchedulingAlgorithmState *SchedulingAlgorithmState_init(void) {
algorithm_state->dispatch_task_queue = NULL;
utarray_new(algorithm_state->cached_submitted_actor_tasks, &task_spec_icd);
utarray_new(algorithm_state->cached_submitted_actor_task_sizes,
&task_spec_size_icd);
algorithm_state->local_actor_infos = NULL;
utarray_new(algorithm_state->available_workers, &worker_icd);
@@ -132,13 +140,13 @@ void SchedulingAlgorithmState_free(SchedulingAlgorithmState *algorithm_state) {
task_queue_entry *elt, *tmp1;
DL_FOREACH_SAFE(algorithm_state->waiting_task_queue, elt, tmp1) {
DL_DELETE(algorithm_state->waiting_task_queue, elt);
free_task_spec(elt->spec);
free(elt->spec);
free(elt);
}
/* Free all the tasks in the dispatch queue. */
DL_FOREACH_SAFE(algorithm_state->dispatch_task_queue, elt, tmp1) {
DL_DELETE(algorithm_state->dispatch_task_queue, elt);
free_task_spec(elt->spec);
free(elt->spec);
free(elt);
}
/* Remove all of the remaining actors. */
@@ -152,11 +160,12 @@ void SchedulingAlgorithmState_free(SchedulingAlgorithmState *algorithm_state) {
/* Free the list of cached actor task specs and the task specs themselves. */
for (int i = 0;
i < utarray_len(algorithm_state->cached_submitted_actor_tasks); ++i) {
task_spec **spec = (task_spec **) utarray_eltptr(
TaskSpec **spec = (TaskSpec **) utarray_eltptr(
algorithm_state->cached_submitted_actor_tasks, i);
free(*spec);
}
utarray_free(algorithm_state->cached_submitted_actor_tasks);
utarray_free(algorithm_state->cached_submitted_actor_task_sizes);
/* Free the list of available workers. */
utarray_free(algorithm_state->available_workers);
utarray_free(algorithm_state->executing_workers);
@@ -195,7 +204,7 @@ void provide_scheduler_info(LocalSchedulerState *state,
waiting_task_queue_length + dispatch_task_queue_length;
info->available_workers = utarray_len(algorithm_state->available_workers);
/* Copy static and dynamic resource information. */
for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
for (int i = 0; i < ResourceIndex_MAX; i++) {
info->dynamic_resources[i] = state->dynamic_resources[i];
info->static_resources[i] = state->static_resources[i];
}
@@ -259,7 +268,7 @@ void remove_actor(SchedulingAlgorithmState *algorithm_state, ActorID actor_id) {
task_queue_entry *task_queue_elt, *tmp;
DL_FOREACH_SAFE(entry->task_queue, task_queue_elt, tmp) {
DL_DELETE(entry->task_queue, task_queue_elt);
free_task_spec(task_queue_elt->spec);
free(task_queue_elt->spec);
free(task_queue_elt);
}
/* Remove the entry from the hash table and free it. */
@@ -310,9 +319,10 @@ void handle_actor_worker_disconnect(LocalSchedulerState *state,
*/
void add_task_to_actor_queue(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
bool from_global_scheduler) {
ActorID actor_id = task_spec_actor_id(spec);
ActorID actor_id = TaskSpec_actor_id(spec);
char tmp[ID_STRING_SIZE];
ObjectID_to_string(actor_id, tmp, ID_STRING_SIZE);
DCHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
@@ -332,7 +342,7 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
CHECK(entry != NULL);
}
int64_t task_counter = task_spec_actor_counter(spec);
int64_t task_counter = TaskSpec_actor_counter(spec);
/* As a sanity check, the counter of the new task should be greater than the
* number of tasks that have executed on this actor so far (since we are
* guaranteeing in-order execution of the tasks on the actor). TODO(rkn): This
@@ -342,8 +352,9 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
/* Create a new task queue entry. */
task_queue_entry *elt = (task_queue_entry *) malloc(sizeof(task_queue_entry));
elt->spec = (task_spec *) malloc(task_spec_size(spec));
memcpy(elt->spec, spec, task_spec_size(spec));
elt->spec = (TaskSpec *) malloc(task_spec_size);
memcpy(elt->spec, spec, task_spec_size);
elt->task_spec_size = task_spec_size;
/* Add the task spec to the actor's task queue in a manner that preserves the
* order of the actor task counters. Iterate from the beginning of the queue
* to find the right place to insert the task queue entry. TODO(pcm): This
@@ -351,15 +362,15 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
* be optimized. */
task_queue_entry *current_entry = entry->task_queue;
while (current_entry != NULL && current_entry->next != NULL &&
task_counter > task_spec_actor_counter(current_entry->spec)) {
task_counter > TaskSpec_actor_counter(current_entry->spec)) {
current_entry = current_entry->next;
}
DL_APPEND_ELEM(entry->task_queue, current_entry, elt);
/* Update the task table. */
if (state->db != NULL) {
Task *task =
Task_alloc(spec, TASK_STATUS_QUEUED, get_db_client_id(state->db));
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
get_db_client_id(state->db));
if (from_global_scheduler) {
/* If the task is from the global scheduler, it's already been added to
* the task table, so just update the entry. */
@@ -403,7 +414,7 @@ bool dispatch_actor_task(LocalSchedulerState *state,
* the actor. */
return false;
}
int64_t next_task_counter = task_spec_actor_counter(entry->task_queue->spec);
int64_t next_task_counter = TaskSpec_actor_counter(entry->task_queue->spec);
if (next_task_counter != entry->task_counter) {
/* We cannot execute the next task on this actor without violating the
* in-order execution guarantee for actor tasks. */
@@ -418,12 +429,13 @@ bool dispatch_actor_task(LocalSchedulerState *state,
* as unavailable. */
task_queue_entry *first_task = entry->task_queue;
entry->task_counter += 1;
assign_task_to_worker(state, first_task->spec, entry->worker);
assign_task_to_worker(state, first_task->spec, first_task->task_spec_size,
entry->worker);
entry->worker_available = false;
/* Remove the task from the actor's task queue. */
DL_DELETE(entry->task_queue, first_task);
/* Free the task spec and the task queue entry. */
free_task_spec(first_task->spec);
free(first_task->spec);
free(first_task);
return true;
}
@@ -479,12 +491,12 @@ void fetch_missing_dependency(LocalSchedulerState *state,
void fetch_missing_dependencies(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_queue_entry *task_entry) {
task_spec *task = task_entry->spec;
int64_t num_args = task_num_args(task);
TaskSpec *task = task_entry->spec;
int64_t num_args = TaskSpec_num_args(task);
int num_missing_dependencies = 0;
for (int i = 0; i < num_args; ++i) {
if (task_arg_type(task, i) == ARG_BY_REF) {
ObjectID obj_id = task_arg_id(task, i);
if (TaskSpec_arg_by_ref(task, i)) {
ObjectID obj_id = TaskSpec_arg_id(task, i);
object_entry *entry;
HASH_FIND(hh, algorithm_state->local_objects, &obj_id, sizeof(obj_id),
entry);
@@ -508,11 +520,11 @@ void fetch_missing_dependencies(LocalSchedulerState *state,
* task are present in the local object store, otherwise it returns
* false.
*/
bool can_run(SchedulingAlgorithmState *algorithm_state, task_spec *task) {
int64_t num_args = task_num_args(task);
bool can_run(SchedulingAlgorithmState *algorithm_state, TaskSpec *task) {
int64_t num_args = TaskSpec_num_args(task);
for (int i = 0; i < num_args; ++i) {
if (task_arg_type(task, i) == ARG_BY_REF) {
ObjectID obj_id = task_arg_id(task, i);
if (TaskSpec_arg_by_ref(task, i)) {
ObjectID obj_id = TaskSpec_arg_id(task, i);
object_entry *entry;
HASH_FIND(hh, algorithm_state->local_objects, &obj_id, sizeof(obj_id),
entry);
@@ -580,7 +592,7 @@ void dispatch_tasks(LocalSchedulerState *state,
}
/* Terminate early if there are no more resources available. */
bool resources_available = false;
for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
for (int i = 0; i < ResourceIndex_MAX; i++) {
if (state->dynamic_resources[i] > 0) {
/* There are still resources left, continue checking tasks. */
resources_available = true;
@@ -593,8 +605,8 @@ void dispatch_tasks(LocalSchedulerState *state,
}
/* Skip to the next task if this task cannot currently be satisfied. */
bool task_satisfied = true;
for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
if (task_spec_get_required_resource(elt->spec, i) >
for (int i = 0; i < ResourceIndex_MAX; i++) {
if (TaskSpec_get_required_resource(elt->spec, i) >
state->dynamic_resources[i]) {
/* Insufficient capacity for this task, proceed to the next task. */
task_satisfied = false;
@@ -612,7 +624,7 @@ void dispatch_tasks(LocalSchedulerState *state,
LocalSchedulerClient **worker = (LocalSchedulerClient **) utarray_back(
algorithm_state->available_workers);
/* Tell the available worker to execute the task. */
assign_task_to_worker(state, elt->spec, *worker);
assign_task_to_worker(state, elt->spec, elt->task_spec_size, *worker);
/* Remove the worker from the available queue, and add it to the executing
* workers. */
utarray_pop_back(algorithm_state->available_workers);
@@ -620,7 +632,7 @@ void dispatch_tasks(LocalSchedulerState *state,
/* Dequeue the task and free the struct. */
print_resource_info(state, elt->spec);
DL_DELETE(algorithm_state->dispatch_task_queue, elt);
free_task_spec(elt->spec);
free(elt->spec);
free(elt);
} /* End for each task in the dispatch queue. */
}
@@ -641,20 +653,22 @@ void dispatch_tasks(LocalSchedulerState *state,
*/
task_queue_entry *queue_task(LocalSchedulerState *state,
task_queue_entry **task_queue,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
bool from_global_scheduler) {
/* Copy the spec and add it to the task queue. The allocated spec will be
* freed when it is assigned to a worker. */
task_queue_entry *elt = (task_queue_entry *) malloc(sizeof(task_queue_entry));
elt->spec = (task_spec *) malloc(task_spec_size(spec));
memcpy(elt->spec, spec, task_spec_size(spec));
elt->spec = (TaskSpec *) malloc(task_spec_size);
memcpy(elt->spec, spec, task_spec_size);
elt->task_spec_size = task_spec_size;
DL_APPEND((*task_queue), elt);
/* The task has been added to a local scheduler queue. Write the entry in the
* task table to notify others that we have queued it. */
if (state->db != NULL) {
Task *task =
Task_alloc(spec, TASK_STATUS_QUEUED, get_db_client_id(state->db));
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
get_db_client_id(state->db));
if (from_global_scheduler) {
/* If the task is from the global scheduler, it's already been added to
* the task table, so just update the entry. */
@@ -684,11 +698,13 @@ task_queue_entry *queue_task(LocalSchedulerState *state,
*/
void queue_waiting_task(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
bool from_global_scheduler) {
LOG_DEBUG("Queueing task in waiting queue");
task_queue_entry *task_entry = queue_task(
state, &algorithm_state->waiting_task_queue, spec, from_global_scheduler);
task_queue_entry *task_entry =
queue_task(state, &algorithm_state->waiting_task_queue, spec,
task_spec_size, from_global_scheduler);
/* If we're queueing this task in the waiting queue, there must be at least
* one missing dependency, so record it. */
fetch_missing_dependencies(state, algorithm_state, task_entry);
@@ -707,10 +723,11 @@ void queue_waiting_task(LocalSchedulerState *state,
*/
void queue_dispatch_task(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
bool from_global_scheduler) {
LOG_DEBUG("Queueing task in dispatch queue");
queue_task(state, &algorithm_state->dispatch_task_queue, spec,
queue_task(state, &algorithm_state->dispatch_task_queue, spec, task_spec_size,
from_global_scheduler);
}
@@ -728,14 +745,17 @@ void queue_dispatch_task(LocalSchedulerState *state,
*/
void queue_task_locally(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
bool from_global_scheduler) {
if (can_run(algorithm_state, spec)) {
/* Dependencies are ready, so push the task to the dispatch queue. */
queue_dispatch_task(state, algorithm_state, spec, from_global_scheduler);
queue_dispatch_task(state, algorithm_state, spec, task_spec_size,
from_global_scheduler);
} else {
/* Dependencies are not ready, so push the task to the waiting queue. */
queue_waiting_task(state, algorithm_state, spec, from_global_scheduler);
queue_waiting_task(state, algorithm_state, spec, task_spec_size,
from_global_scheduler);
}
}
@@ -751,7 +771,8 @@ void queue_task_locally(LocalSchedulerState *state,
*/
void give_task_to_local_scheduler(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec,
TaskSpec *spec,
int64_t task_spec_size,
DBClientID local_scheduler_id) {
if (DBClientID_equal(local_scheduler_id, get_db_client_id(state->db))) {
LOG_WARN("Local scheduler is trying to assign a task to itself.");
@@ -759,7 +780,8 @@ void give_task_to_local_scheduler(LocalSchedulerState *state,
CHECK(state->db != NULL);
/* Assign the task to the relevant local scheduler. */
DCHECK(state->config.global_scheduler_exists);
Task *task = Task_alloc(spec, TASK_STATUS_SCHEDULED, local_scheduler_id);
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_SCHEDULED,
local_scheduler_id);
task_table_add_task(state->db, task, NULL, NULL, NULL);
}
@@ -773,27 +795,27 @@ void give_task_to_local_scheduler(LocalSchedulerState *state,
*/
void give_task_to_global_scheduler(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec) {
TaskSpec *spec,
int64_t task_spec_size) {
if (state->db == NULL || !state->config.global_scheduler_exists) {
/* A global scheduler is not available, so queue the task locally. */
queue_task_locally(state, algorithm_state, spec, false);
queue_task_locally(state, algorithm_state, spec, task_spec_size, false);
return;
}
/* Pass on the task to the global scheduler. */
DCHECK(state->config.global_scheduler_exists);
Task *task = Task_alloc(spec, TASK_STATUS_WAITING, NIL_ID);
Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_WAITING, NIL_ID);
DCHECK(state->db != NULL);
task_table_add_task(state->db, task, NULL, NULL, NULL);
}
bool resource_constraints_satisfied(LocalSchedulerState *state,
task_spec *spec) {
TaskSpec *spec) {
/* At the local scheduler, if required resource vector exceeds either static
* or dynamic resource vector, the resource constraint is not satisfied. */
for (int i = 0; i < MAX_RESOURCE_INDEX; i++) {
if (task_spec_get_required_resource(spec, i) > state->static_resources[i] ||
task_spec_get_required_resource(spec, i) >
state->dynamic_resources[i]) {
for (int i = 0; i < ResourceIndex_MAX; i++) {
if (TaskSpec_get_required_resource(spec, i) > state->static_resources[i] ||
TaskSpec_get_required_resource(spec, i) > state->dynamic_resources[i]) {
return false;
}
}
@@ -802,7 +824,8 @@ bool resource_constraints_satisfied(LocalSchedulerState *state,
void handle_task_submitted(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec) {
TaskSpec *spec,
int64_t task_spec_size) {
/* TODO(atumanov): if static is satisfied and local objects ready, but dynamic
* resource is currently unavailable, then consider queueing task locally and
* recheck dynamic next time. */
@@ -814,10 +837,10 @@ void handle_task_submitted(LocalSchedulerState *state,
if (resource_constraints_satisfied(state, spec) &&
(utarray_len(algorithm_state->available_workers) > 0) &&
can_run(algorithm_state, spec)) {
queue_dispatch_task(state, algorithm_state, spec, false);
queue_dispatch_task(state, algorithm_state, spec, task_spec_size, false);
} else {
/* Give the task to the global scheduler to schedule, if it exists. */
give_task_to_global_scheduler(state, algorithm_state, spec);
give_task_to_global_scheduler(state, algorithm_state, spec, task_spec_size);
}
/* Try to dispatch tasks, since we may have added one to the queue. */
@@ -826,8 +849,9 @@ void handle_task_submitted(LocalSchedulerState *state,
void handle_actor_task_submitted(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec) {
ActorID actor_id = task_spec_actor_id(spec);
TaskSpec *spec,
int64_t task_spec_size) {
ActorID actor_id = TaskSpec_actor_id(spec);
CHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
/* Find the local scheduler responsible for this actor. */
@@ -840,6 +864,8 @@ void handle_actor_task_submitted(LocalSchedulerState *state,
* will be resubmitted (internally by the local scheduler) whenever a new
* actor notification arrives. */
utarray_push_back(algorithm_state->cached_submitted_actor_tasks, &spec);
utarray_push_back(algorithm_state->cached_submitted_actor_task_sizes,
&task_spec_size);
return;
}
@@ -847,13 +873,14 @@ void handle_actor_task_submitted(LocalSchedulerState *state,
get_db_client_id(state->db))) {
/* This local scheduler is responsible for the actor, so handle the task
* locally. */
add_task_to_actor_queue(state, algorithm_state, spec, false);
add_task_to_actor_queue(state, algorithm_state, spec, task_spec_size,
false);
/* Attempt to dispatch tasks to this actor. */
dispatch_actor_task(state, algorithm_state, actor_id);
} else {
/* This local scheduler is not responsible for the task, so assign the task
* directly to the actor that is responsible. */
give_task_to_local_scheduler(state, algorithm_state, spec,
give_task_to_local_scheduler(state, algorithm_state, spec, task_spec_size,
entry->local_scheduler_id);
}
}
@@ -864,35 +891,43 @@ void handle_actor_creation_notification(
ActorID actor_id) {
int num_cached_actor_tasks =
utarray_len(algorithm_state->cached_submitted_actor_tasks);
CHECK(num_cached_actor_tasks ==
utarray_len(algorithm_state->cached_submitted_actor_task_sizes));
for (int i = 0; i < num_cached_actor_tasks; ++i) {
task_spec **spec = (task_spec **) utarray_eltptr(
TaskSpec **spec = (TaskSpec **) utarray_eltptr(
algorithm_state->cached_submitted_actor_tasks, i);
int64_t *task_spec_size = (int64_t *) utarray_eltptr(
algorithm_state->cached_submitted_actor_task_sizes, i);
/* Note that handle_actor_task_submitted may append the spec to the end of
* the cached_submitted_actor_tasks array. */
handle_actor_task_submitted(state, algorithm_state, *spec);
handle_actor_task_submitted(state, algorithm_state, *spec, *task_spec_size);
}
/* Remove all the tasks that were resubmitted. This does not erase the tasks
* that were newly appended to the cached_submitted_actor_tasks array. */
utarray_erase(algorithm_state->cached_submitted_actor_tasks, 0,
num_cached_actor_tasks);
utarray_erase(algorithm_state->cached_submitted_actor_task_sizes, 0,
num_cached_actor_tasks);
}
void handle_task_scheduled(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec) {
TaskSpec *spec,
int64_t task_spec_size) {
/* This callback handles tasks that were assigned to this local scheduler by
* the global scheduler, so we can safely assert that there is a connection to
* the database. */
DCHECK(state->db != NULL);
DCHECK(state->config.global_scheduler_exists);
/* Push the task to the appropriate queue. */
queue_task_locally(state, algorithm_state, spec, true);
queue_task_locally(state, algorithm_state, spec, task_spec_size, true);
dispatch_tasks(state, algorithm_state);
}
void handle_actor_task_scheduled(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec) {
TaskSpec *spec,
int64_t task_spec_size) {
/* This callback handles tasks that were assigned to this local scheduler by
* the global scheduler or by other workers, so we can safely assert that
* there is a connection to the database. */
@@ -900,7 +935,7 @@ void handle_actor_task_scheduled(LocalSchedulerState *state,
DCHECK(state->config.global_scheduler_exists);
/* Check that the task is meant to run on an actor that this local scheduler
* is responsible for. */
ActorID actor_id = task_spec_actor_id(spec);
ActorID actor_id = TaskSpec_actor_id(spec);
DCHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
actor_map_entry *entry;
HASH_FIND(hh, state->actor_mapping, &actor_id, sizeof(actor_id), entry);
@@ -918,7 +953,7 @@ void handle_actor_task_scheduled(LocalSchedulerState *state,
"corresponding actor_map_entry is not present. This should be rare.");
}
/* Push the task to the appropriate queue. */
add_task_to_actor_queue(state, algorithm_state, spec, true);
add_task_to_actor_queue(state, algorithm_state, spec, task_spec_size, true);
dispatch_actor_task(state, algorithm_state, actor_id);
}
@@ -1047,7 +1082,7 @@ void handle_worker_blocked(LocalSchedulerState *state,
/* Return the resources that the blocked worker was using. */
CHECK(worker->task_in_progress != NULL);
task_spec *spec = Task_task_spec(worker->task_in_progress);
TaskSpec *spec = Task_task_spec(worker->task_in_progress);
update_dynamic_resources(state, spec, true);
/* Add the worker to the list of blocked workers. */
worker->is_blocked = true;
@@ -1089,7 +1124,7 @@ void handle_worker_unblocked(LocalSchedulerState *state,
* fixed by having blocked workers explicitly yield and wait to be given
* back resources before continuing execution. */
CHECK(worker->task_in_progress != NULL);
task_spec *spec = Task_task_spec(worker->task_in_progress);
TaskSpec *spec = Task_task_spec(worker->task_in_progress);
update_dynamic_resources(state, spec, false);
/* Add the worker to the list of executing workers. */
worker->is_blocked = false;
@@ -1171,11 +1206,11 @@ void handle_object_removed(LocalSchedulerState *state,
task_queue_entry *elt, *tmp;
/* Track the dependency for tasks that were in the waiting queue. */
DL_FOREACH(algorithm_state->waiting_task_queue, elt) {
task_spec *task = elt->spec;
int64_t num_args = task_num_args(task);
TaskSpec *task = elt->spec;
int64_t num_args = TaskSpec_num_args(task);
for (int i = 0; i < num_args; ++i) {
if (task_arg_type(task, i) == ARG_BY_REF) {
ObjectID arg_id = task_arg_id(task, i);
if (TaskSpec_arg_by_ref(task, i)) {
ObjectID arg_id = TaskSpec_arg_id(task, i);
if (ObjectID_equal(arg_id, removed_object_id)) {
fetch_missing_dependency(state, algorithm_state, elt,
removed_object_id);
@@ -1186,11 +1221,11 @@ void handle_object_removed(LocalSchedulerState *state,
/* Track the dependency for tasks that were in the dispatch queue. Remove
* these tasks from the dispatch queue and push them to the waiting queue. */
DL_FOREACH_SAFE(algorithm_state->dispatch_task_queue, elt, tmp) {
task_spec *task = elt->spec;
int64_t num_args = task_num_args(task);
TaskSpec *task = elt->spec;
int64_t num_args = TaskSpec_num_args(task);
for (int i = 0; i < num_args; ++i) {
if (task_arg_type(task, i) == ARG_BY_REF) {
ObjectID arg_id = task_arg_id(task, i);
if (TaskSpec_arg_by_ref(task, i)) {
ObjectID arg_id = TaskSpec_arg_id(task, i);
if (ObjectID_equal(arg_id, removed_object_id)) {
LOG_DEBUG("Moved task from dispatch queue back to waiting queue");
DL_DELETE(algorithm_state->dispatch_task_queue, elt);
src/local_scheduler/local_scheduler_algorithm.h
+8 -4
View File
@@ -58,7 +58,8 @@ void provide_scheduler_info(LocalSchedulerState *state,
*/
void handle_task_submitted(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec);
TaskSpec *spec,
int64_t task_spec_size);
/**
* This version of handle_task_submitted is used when the task being submitted
@@ -71,7 +72,8 @@ void handle_task_submitted(LocalSchedulerState *state,
*/
void handle_actor_task_submitted(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec);
TaskSpec *spec,
int64_t task_spec_size);
/**
* This function will be called when the local scheduler receives a notification
@@ -99,7 +101,8 @@ void handle_actor_creation_notification(
*/
void handle_task_scheduled(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec);
TaskSpec *spec,
int64_t task_spec_size);
/**
* This function will be called when an actor task is assigned by the global
@@ -113,7 +116,8 @@ void handle_task_scheduled(LocalSchedulerState *state,
*/
void handle_actor_task_scheduled(LocalSchedulerState *state,
SchedulingAlgorithmState *algorithm_state,
task_spec *spec);
TaskSpec *spec,
int64_t task_spec_size);
/**
* This function is called if a new object becomes available in the local
src/local_scheduler/local_scheduler_client.cc
+8 -8
View File
@@ -48,22 +48,22 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn,
free(message);
}
void local_scheduler_submit(LocalSchedulerConnection *conn, task_spec *task) {
write_message(conn->conn, SUBMIT_TASK, task_spec_size(task),
(uint8_t *) task);
void local_scheduler_submit(LocalSchedulerConnection *conn,
TaskSpec *task,
int64_t task_size) {
write_message(conn->conn, SUBMIT_TASK, task_size, (uint8_t *) task);
}
task_spec *local_scheduler_get_task(LocalSchedulerConnection *conn) {
TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
int64_t *task_size) {
write_message(conn->conn, GET_TASK, 0, NULL);
int64_t type;
int64_t length;
uint8_t *message;
/* Receive a task from the local scheduler. This will block until the local
* scheduler gives this client a task. */
read_message(conn->conn, &type, &length, &message);
read_message(conn->conn, &type, task_size, &message);
CHECK(type == EXECUTE_TASK);
task_spec *task = (task_spec *) message;
CHECK(length == task_spec_size(task));
TaskSpec *task = (TaskSpec *) message;
return task;
}
src/local_scheduler/local_scheduler_client.h
+5 -2
View File
@@ -39,7 +39,9 @@ void LocalSchedulerConnection_free(LocalSchedulerConnection *conn);
* @param task The address of the task to submit.
* @return Void.
*/
void local_scheduler_submit(LocalSchedulerConnection *conn, task_spec *task);
void local_scheduler_submit(LocalSchedulerConnection *conn,
TaskSpec *task,
int64_t task_size);
/**
* Log an event to the event log. This will call RPUSH key value. We use RPUSH
@@ -70,7 +72,8 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn,
* @param conn The connection information.
* @return The address of the assigned task.
*/
task_spec *local_scheduler_get_task(LocalSchedulerConnection *conn);
TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
int64_t *task_size);
/**
* Tell the local scheduler that the client has finished executing a task.
src/local_scheduler/local_scheduler_extension.cc
+8 -5
View File
@@ -41,20 +41,21 @@ static PyObject *PyLocalSchedulerClient_submit(PyObject *self, PyObject *args) {
}
local_scheduler_submit(
((PyLocalSchedulerClient *) self)->local_scheduler_connection,
((PyTask *) py_task)->spec);
((PyTask *) py_task)->spec, ((PyTask *) py_task)->size);
Py_RETURN_NONE;
}
// clang-format off
static PyObject *PyLocalSchedulerClient_get_task(PyObject *self) {
task_spec *task_spec;
TaskSpec *task_spec;
/* Drop the global interpreter lock while we get a task because
* local_scheduler_get_task may block for a long time. */
int64_t task_size;
Py_BEGIN_ALLOW_THREADS
task_spec = local_scheduler_get_task(
((PyLocalSchedulerClient *) self)->local_scheduler_connection);
((PyLocalSchedulerClient *) self)->local_scheduler_connection, &task_size);
Py_END_ALLOW_THREADS
return PyTask_make(task_spec);
return PyTask_make(task_spec, task_size);
}
// clang-format on
@@ -154,7 +155,7 @@ static PyMethodDef local_scheduler_methods[] = {
"Return the object ID for a put call within a task."},
{"task_from_string", PyTask_from_string, METH_VARARGS,
"Creates a Python PyTask object from a string representation of "
"task_spec."},
"TaskSpec."},
{"task_to_string", PyTask_to_string, METH_VARARGS,
"Translates a PyTask python object to a byte string."},
{NULL} /* Sentinel */
@@ -223,6 +224,8 @@ MOD_INIT(liblocal_scheduler_library) {
PyModule_AddObject(m, "LocalSchedulerClient",
(PyObject *) &PyLocalSchedulerClientType);
g_task_builder = make_task_builder();
char local_scheduler_error[] = "local_scheduler.error";
LocalSchedulerError = PyErr_NewException(local_scheduler_error, NULL, NULL);
Py_INCREF(LocalSchedulerError);
src/local_scheduler/local_scheduler_shared.h
+2 -2
View File
@@ -94,10 +94,10 @@ typedef struct {
UT_array *input_buffer;
/** Vector of static attributes associated with the node owned by this local
* scheduler. */
double static_resources[MAX_RESOURCE_INDEX];
double static_resources[ResourceIndex_MAX];
/** Vector of dynamic attributes associated with the node owned by this local
* scheduler. */
double dynamic_resources[MAX_RESOURCE_INDEX];
double dynamic_resources[ResourceIndex_MAX];
} LocalSchedulerState;
/** Contains all information associated with a local scheduler client. */
src/local_scheduler/test/local_scheduler_tests.cc
+77 -55
View File
@@ -9,6 +9,7 @@
#include "common.h"
#include "test/test_common.h"
#include "test/example_task.h"
#include "event_loop.h"
#include "io.h"
#include "utstring.h"
@@ -23,6 +24,8 @@
SUITE(local_scheduler_tests);
TaskBuilder *g_task_builder = NULL;
const char *plasma_store_socket_name = "/tmp/plasma_store_socket_1";
const char *plasma_manager_socket_name_format = "/tmp/plasma_manager_socket_%d";
const char *local_scheduler_socket_name_format =
@@ -56,8 +59,8 @@ LocalSchedulerMock *LocalSchedulerMock_init(int num_workers,
const char *node_ip_address = "127.0.0.1";
const char *redis_addr = node_ip_address;
int redis_port = 6379;
const double static_resource_conf[MAX_RESOURCE_INDEX] = {DEFAULT_NUM_CPUS,
DEFAULT_NUM_GPUS};
const double static_resource_conf[ResourceIndex_MAX] = {DEFAULT_NUM_CPUS,
DEFAULT_NUM_GPUS};
LocalSchedulerMock *mock =
(LocalSchedulerMock *) malloc(sizeof(LocalSchedulerMock));
memset(mock, 0, sizeof(LocalSchedulerMock));
@@ -155,8 +158,9 @@ TEST object_reconstruction_test(void) {
LocalSchedulerConnection *worker = local_scheduler->conns[0];
/* Create a task with zero dependencies and one return value. */
task_spec *spec = example_task_spec(0, 1);
ObjectID return_id = task_return(spec, 0);
int64_t task_size;
TaskSpec *spec = example_task_spec(0, 1, &task_size);
ObjectID return_id = TaskSpec_return(spec, 0);
/* Add an empty object table entry for the object we want to reconstruct, to
* simulate it having been created and evicted. */
@@ -176,15 +180,21 @@ TEST object_reconstruction_test(void) {
if (pid == 0) {
/* Make sure we receive the task twice. First from the initial submission,
* and second from the reconstruct request. */
local_scheduler_submit(worker, spec);
task_spec *task_assigned = local_scheduler_get_task(worker);
ASSERT_EQ(memcmp(task_assigned, spec, task_spec_size(spec)), 0);
task_spec *reconstruct_task = local_scheduler_get_task(worker);
ASSERT_EQ(memcmp(reconstruct_task, spec, task_spec_size(spec)), 0);
int64_t task_assigned_size;
local_scheduler_submit(worker, spec, task_size);
TaskSpec *task_assigned =
local_scheduler_get_task(worker, &task_assigned_size);
ASSERT_EQ(memcmp(task_assigned, spec, task_size), 0);
ASSERT_EQ(task_assigned_size, task_size);
int64_t reconstruct_task_size;
TaskSpec *reconstruct_task =
local_scheduler_get_task(worker, &reconstruct_task_size);
ASSERT_EQ(memcmp(reconstruct_task, spec, task_size), 0);
ASSERT_EQ(reconstruct_task_size, task_size);
/* Clean up. */
free_task_spec(reconstruct_task);
free_task_spec(task_assigned);
free_task_spec(spec);
free(reconstruct_task);
free(task_assigned);
TaskSpec_free(spec);
LocalSchedulerMock_free(local_scheduler);
exit(0);
} else {
@@ -196,12 +206,12 @@ TEST object_reconstruction_test(void) {
/* Set the task's status to TASK_STATUS_DONE to prevent the race condition
* that would suppress object reconstruction. */
Task *task = Task_alloc(
spec, TASK_STATUS_DONE,
spec, task_size, TASK_STATUS_DONE,
get_db_client_id(local_scheduler->local_scheduler_state->db));
task_table_add_task(local_scheduler->local_scheduler_state->db, task, NULL,
NULL, NULL);
/* Trigger reconstruction, and run the event loop again. */
ObjectID return_id = task_return(spec, 0);
ObjectID return_id = TaskSpec_return(spec, 0);
local_scheduler_reconstruct_object(worker, return_id);
event_loop_add_timer(local_scheduler->loop, 500,
(event_loop_timer_handler) timeout_handler, NULL);
@@ -209,7 +219,7 @@ TEST object_reconstruction_test(void) {
/* Wait for the child process to exit and check that there are no tasks
* left in the local scheduler's task queue. Then, clean up. */
wait(NULL);
free_task_spec(spec);
TaskSpec_free(spec);
ASSERT_EQ(num_waiting_tasks(
local_scheduler->local_scheduler_state->algorithm_state),
0);
@@ -232,14 +242,15 @@ TEST object_reconstruction_recursive_test(void) {
/* Create a chain of tasks, each one dependent on the one before it. Mark
* each object as available so that tasks will run immediately. */
const int NUM_TASKS = 10;
task_spec *specs[NUM_TASKS];
specs[0] = example_task_spec(0, 1);
TaskSpec *specs[NUM_TASKS];
int64_t task_sizes[NUM_TASKS];
specs[0] = example_task_spec(0, 1, &task_sizes[0]);
for (int i = 1; i < NUM_TASKS; ++i) {
ObjectID arg_id = task_return(specs[i - 1], 0);
ObjectID arg_id = TaskSpec_return(specs[i - 1], 0);
handle_object_available(
local_scheduler->local_scheduler_state,
local_scheduler->local_scheduler_state->algorithm_state, arg_id);
specs[i] = example_task_spec_with_args(1, 1, &arg_id);
specs[i] = example_task_spec_with_args(1, 1, &arg_id, &task_sizes[i]);
}
/* Add an empty object table entry for each object we want to reconstruct, to
@@ -247,7 +258,7 @@ TEST object_reconstruction_recursive_test(void) {
const char *client_id = "clientid";
redisContext *context = redisConnect("127.0.0.1", 6379);
for (int i = 0; i < NUM_TASKS; ++i) {
ObjectID return_id = task_return(specs[i], 0);
ObjectID return_id = TaskSpec_return(specs[i], 0);
redisReply *reply = (redisReply *) redisCommand(
context, "RAY.OBJECT_TABLE_ADD %b %ld %b %s", return_id.id,
sizeof(return_id.id), 1, NIL_DIGEST, (size_t) DIGEST_SIZE, client_id);
@@ -263,32 +274,34 @@ TEST object_reconstruction_recursive_test(void) {
if (pid == 0) {
/* Submit the tasks, and make sure each one gets assigned to a worker. */
for (int i = 0; i < NUM_TASKS; ++i) {
local_scheduler_submit(worker, specs[i]);
local_scheduler_submit(worker, specs[i], task_sizes[i]);
}
/* Make sure we receive each task from the initial submission. */
for (int i = 0; i < NUM_TASKS; ++i) {
task_spec *task_assigned = local_scheduler_get_task(worker);
ASSERT_EQ(memcmp(task_assigned, specs[i], task_spec_size(task_assigned)),
0);
free_task_spec(task_assigned);
int64_t task_size;
TaskSpec *task_assigned = local_scheduler_get_task(worker, &task_size);
ASSERT_EQ(memcmp(task_assigned, specs[i], task_sizes[i]), 0);
ASSERT_EQ(task_size, task_sizes[i]);
free(task_assigned);
}
/* Check that the workers receive all tasks in the final return object's
* lineage during reconstruction. */
for (int i = 0; i < NUM_TASKS; ++i) {
task_spec *task_assigned = local_scheduler_get_task(worker);
int64_t task_assigned_size;
TaskSpec *task_assigned =
local_scheduler_get_task(worker, &task_assigned_size);
bool found = false;
for (int j = 0; j < NUM_TASKS; ++j) {
if (specs[j] == NULL) {
continue;
}
if (memcmp(task_assigned, specs[j], task_spec_size(task_assigned)) ==
0) {
if (memcmp(task_assigned, specs[j], task_assigned_size) == 0) {
found = true;
free_task_spec(specs[j]);
TaskSpec_free(specs[j]);
specs[j] = NULL;
}
}
free_task_spec(task_assigned);
free(task_assigned);
ASSERT(found);
}
LocalSchedulerMock_free(local_scheduler);
@@ -302,13 +315,13 @@ TEST object_reconstruction_recursive_test(void) {
/* Set the final task's status to TASK_STATUS_DONE to prevent the race
* condition that would suppress object reconstruction. */
Task *last_task = Task_alloc(
specs[NUM_TASKS - 1], TASK_STATUS_DONE,
specs[NUM_TASKS - 1], task_sizes[NUM_TASKS - 1], TASK_STATUS_DONE,
get_db_client_id(local_scheduler->local_scheduler_state->db));
task_table_add_task(local_scheduler->local_scheduler_state->db, last_task,
NULL, NULL, NULL);
/* Trigger reconstruction for the last object, and run the event loop
* again. */
ObjectID return_id = task_return(specs[NUM_TASKS - 1], 0);
ObjectID return_id = TaskSpec_return(specs[NUM_TASKS - 1], 0);
local_scheduler_reconstruct_object(worker, return_id);
event_loop_add_timer(local_scheduler->loop, 500,
(event_loop_timer_handler) timeout_handler, NULL);
@@ -323,7 +336,7 @@ TEST object_reconstruction_recursive_test(void) {
local_scheduler->local_scheduler_state->algorithm_state),
0);
for (int i = 0; i < NUM_TASKS; ++i) {
free_task_spec(specs[i]);
TaskSpec_free(specs[i]);
}
LocalSchedulerMock_free(local_scheduler);
PASS();
@@ -334,35 +347,41 @@ TEST object_reconstruction_recursive_test(void) {
* Test that object reconstruction gets suppressed when there is a location
* listed for the object in the object table.
*/
task_spec *object_reconstruction_suppression_spec;
TaskSpec *object_reconstruction_suppression_spec;
int64_t object_reconstruction_suppression_size;
void object_reconstruction_suppression_callback(ObjectID object_id,
void *user_context) {
/* Submit the task after adding the object to the object table. */
LocalSchedulerConnection *worker = (LocalSchedulerConnection *) user_context;
local_scheduler_submit(worker, object_reconstruction_suppression_spec);
local_scheduler_submit(worker, object_reconstruction_suppression_spec,
object_reconstruction_suppression_size);
}
TEST object_reconstruction_suppression_test(void) {
LocalSchedulerMock *local_scheduler = LocalSchedulerMock_init(0, 1);
LocalSchedulerConnection *worker = local_scheduler->conns[0];
object_reconstruction_suppression_spec = example_task_spec(0, 1);
ObjectID return_id = task_return(object_reconstruction_suppression_spec, 0);
object_reconstruction_suppression_spec =
example_task_spec(0, 1, &object_reconstruction_suppression_size);
ObjectID return_id =
TaskSpec_return(object_reconstruction_suppression_spec, 0);
pid_t pid = fork();
if (pid == 0) {
/* Make sure we receive the task once. This will block until the
* object_table_add callback completes. */
task_spec *task_assigned = local_scheduler_get_task(worker);
int64_t task_assigned_size;
TaskSpec *task_assigned =
local_scheduler_get_task(worker, &task_assigned_size);
ASSERT_EQ(memcmp(task_assigned, object_reconstruction_suppression_spec,
task_spec_size(object_reconstruction_suppression_spec)),
object_reconstruction_suppression_size),
0);
/* Trigger a reconstruction. We will check that no tasks get queued as a
* result of this line in the event loop process. */
local_scheduler_reconstruct_object(worker, return_id);
/* Clean up. */
free_task_spec(task_assigned);
free_task_spec(object_reconstruction_suppression_spec);
free(task_assigned);
TaskSpec_free(object_reconstruction_suppression_spec);
LocalSchedulerMock_free(local_scheduler);
exit(0);
} else {
@@ -389,7 +408,7 @@ TEST object_reconstruction_suppression_test(void) {
ASSERT_EQ(num_dispatch_tasks(
local_scheduler->local_scheduler_state->algorithm_state),
0);
free_task_spec(object_reconstruction_suppression_spec);
TaskSpec_free(object_reconstruction_suppression_spec);
db_disconnect(db);
LocalSchedulerMock_free(local_scheduler);
PASS();
@@ -403,12 +422,13 @@ TEST task_dependency_test(void) {
/* Get the first worker. */
LocalSchedulerClient *worker =
*((LocalSchedulerClient **) utarray_eltptr(state->workers, 0));
task_spec *spec = example_task_spec(1, 1);
ObjectID oid = task_arg_id(spec, 0);
int64_t task_size;
TaskSpec *spec = example_task_spec(1, 1, &task_size);
ObjectID oid = TaskSpec_arg_id(spec, 0);
/* Check that the task gets queued in the waiting queue if the task is
* submitted, but the input and workers are not available. */
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 1);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 0);
/* Once the input is available, the task gets moved to the dispatch queue. */
@@ -424,7 +444,7 @@ TEST task_dependency_test(void) {
/* Check that the task gets queued in the waiting queue if the task is
* submitted and a worker is available, but the input is not. */
handle_object_removed(state, oid);
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
handle_worker_available(state, algorithm_state, worker);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 1);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 0);
@@ -436,7 +456,7 @@ TEST task_dependency_test(void) {
/* Check that the task gets queued in the dispatch queue if the task is
* submitted and the input is available, but no worker is available yet. */
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 0);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 1);
/* Once a worker is available, the task gets assigned. */
@@ -448,7 +468,7 @@ TEST task_dependency_test(void) {
/* If an object gets removed, check the first scenario again, where the task
* gets queued in the waiting task if the task is submitted and a worker is
* available, but the input is not. */
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 0);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 1);
/* If the input is removed while a task is in the dispatch queue, the task
@@ -466,7 +486,7 @@ TEST task_dependency_test(void) {
ASSERT_EQ(num_waiting_tasks(algorithm_state), 0);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 0);
free_task_spec(spec);
TaskSpec_free(spec);
LocalSchedulerMock_free(local_scheduler);
PASS();
}
@@ -478,13 +498,14 @@ TEST task_multi_dependency_test(void) {
/* Get the first worker. */
LocalSchedulerClient *worker =
*((LocalSchedulerClient **) utarray_eltptr(state->workers, 0));
task_spec *spec = example_task_spec(2, 1);
ObjectID oid1 = task_arg_id(spec, 0);
ObjectID oid2 = task_arg_id(spec, 1);
int64_t task_size;
TaskSpec *spec = example_task_spec(2, 1, &task_size);
ObjectID oid1 = TaskSpec_arg_id(spec, 0);
ObjectID oid2 = TaskSpec_arg_id(spec, 1);
/* Check that the task gets queued in the waiting queue if the task is
* submitted, but the inputs and workers are not available. */
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 1);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 0);
/* Check that the task stays in the waiting queue if only one input becomes
@@ -504,7 +525,7 @@ TEST task_multi_dependency_test(void) {
/* Check that the task gets queued in the dispatch queue if the task is
* submitted and the inputs are available, but no worker is available yet. */
handle_task_submitted(state, algorithm_state, spec);
handle_task_submitted(state, algorithm_state, spec, task_size);
ASSERT_EQ(num_waiting_tasks(algorithm_state), 0);
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 1);
/* If any input is removed while a task is in the dispatch queue, the task
@@ -540,7 +561,7 @@ TEST task_multi_dependency_test(void) {
ASSERT_EQ(num_dispatch_tasks(algorithm_state), 0);
reset_worker(local_scheduler, worker);
free_task_spec(spec);
TaskSpec_free(spec);
LocalSchedulerMock_free(local_scheduler);
PASS();
}
@@ -633,6 +654,7 @@ SUITE(local_scheduler_tests) {
GREATEST_MAIN_DEFS();
int main(int argc, char **argv) {
g_task_builder = make_task_builder();
GREATEST_MAIN_BEGIN();
RUN_SUITE(local_scheduler_tests);
GREATEST_MAIN_END();