diff --git a/src/global_scheduler/global_scheduler.cc b/src/global_scheduler/global_scheduler.cc index 0271cb95c..ffa9bc58d 100644 --- a/src/global_scheduler/global_scheduler.cc +++ b/src/global_scheduler/global_scheduler.cc @@ -13,13 +13,6 @@ #include "state/table.h" #include "state/task_table.h" -/* This is used to define the array of local schedulers used to define the - * GlobalSchedulerState type. */ -UT_icd local_scheduler_icd = {sizeof(LocalScheduler), NULL, NULL, NULL}; - -/* This is used to define the array of tasks that haven't been scheduled yet. */ -UT_icd pending_tasks_icd = {sizeof(Task *), NULL, NULL, NULL}; - /** * Assign the given task to the local scheduler, update Redis and scheduler data * structures. @@ -43,19 +36,38 @@ void assign_task_to_local_scheduler(GlobalSchedulerState *state, UNUSED(id_string); task_table_update(state->db, Task_copy(task), NULL, NULL, NULL); + /* Update the object table info to reflect the fact that the results of this + * task will be created on the machine that the task was assigned to. This can + * be used to improve locality-aware scheduling. */ + for (int64_t i = 0; i < TaskSpec_num_returns(spec); ++i) { + ObjectID return_id = TaskSpec_return(spec, i); + if (state->scheduler_object_info_table.find(return_id) == + state->scheduler_object_info_table.end()) { + SchedulerObjectInfo &obj_info_entry = + state->scheduler_object_info_table[return_id]; + /* The value -1 indicates that the size of the object is not known yet. */ + obj_info_entry.data_size = -1; + } + CHECK(state->local_scheduler_plasma_map.count(local_scheduler_id) == 1); + state->scheduler_object_info_table[return_id].object_locations.push_back( + state->local_scheduler_plasma_map[local_scheduler_id]); + } + /* TODO(rkn): We should probably pass around local_scheduler struct pointers * instead of db_client_id objects. */ /* Update the local scheduler info. */ - LocalScheduler *local_scheduler = - get_local_scheduler(state, local_scheduler_id); - local_scheduler->num_tasks_sent += 1; - local_scheduler->num_recent_tasks_sent += 1; + auto it = state->local_schedulers.find(local_scheduler_id); + CHECK(it != state->local_schedulers.end()); + + LocalScheduler &local_scheduler = it->second; + local_scheduler.num_tasks_sent += 1; + local_scheduler.num_recent_tasks_sent += 1; /* Resource accounting update for this local scheduler. */ for (int i = 0; i < ResourceIndex_MAX; i++) { /* Subtract task's resource from the cached dynamic resource capacity for * this local scheduler. This will be overwritten on the next heartbeat. */ - local_scheduler->info.dynamic_resources[i] = - MAX(0, local_scheduler->info.dynamic_resources[i] - + local_scheduler.info.dynamic_resources[i] = + MAX(0, local_scheduler.info.dynamic_resources[i] - TaskSpec_get_required_resource(spec, i)); } } @@ -64,66 +76,40 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop, const char *node_ip_address, const char *redis_primary_addr, int redis_primary_port) { - GlobalSchedulerState *state = - (GlobalSchedulerState *) malloc(sizeof(GlobalSchedulerState)); - /* Must initialize state to 0. Sets hashmap head(s) to NULL. */ - memset(state, 0, sizeof(GlobalSchedulerState)); + GlobalSchedulerState *state = new GlobalSchedulerState(); state->db = db_connect(std::string(redis_primary_addr), redis_primary_port, "global_scheduler", node_ip_address, 0, NULL); db_attach(state->db, loop, false); - utarray_new(state->local_schedulers, &local_scheduler_icd); state->policy_state = GlobalSchedulerPolicyState_init(); - /* Initialize the array of tasks that have not been scheduled yet. */ - utarray_new(state->pending_tasks, &pending_tasks_icd); return state; } void GlobalSchedulerState_free(GlobalSchedulerState *state) { - AuxAddressEntry *entry, *tmp; - db_disconnect(state->db); - utarray_free(state->local_schedulers); + state->local_schedulers.clear(); GlobalSchedulerPolicyState_free(state->policy_state); /* Delete the plasma to local scheduler association map. */ - HASH_ITER(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, entry, - tmp) { - HASH_DELETE(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, - entry); - /* The hash entry is shared with the local_scheduler_plasma hashmap and will - * be freed there. */ - free(entry->aux_address); - } + state->plasma_local_scheduler_map.clear(); /* Delete the local scheduler to plasma association map. */ - HASH_ITER(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, entry, - tmp) { - HASH_DELETE(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, - entry); - /* Now free the shared hash entry -- no longer needed. */ - free(entry); - } + state->local_scheduler_plasma_map.clear(); /* Free the scheduler object info table. */ - SchedulerObjectInfo *object_entry, *tmp_entry; - HASH_ITER(hh, state->scheduler_object_info_table, object_entry, tmp_entry) { - HASH_DELETE(hh, state->scheduler_object_info_table, object_entry); - utarray_free(object_entry->object_locations); - free(object_entry); - } + state->scheduler_object_info_table.clear(); /* Free the array of unschedulable tasks. */ - int64_t num_pending_tasks = utarray_len(state->pending_tasks); + int64_t num_pending_tasks = state->pending_tasks.size(); if (num_pending_tasks > 0) { LOG_WARN("There are %" PRId64 " remaining tasks in the pending tasks array.", num_pending_tasks); } for (int i = 0; i < num_pending_tasks; ++i) { - Task **pending_task = (Task **) utarray_eltptr(state->pending_tasks, i); - Task_free(*pending_task); + Task *pending_task = state->pending_tasks[i]; + Task_free(pending_task); } - utarray_free(state->pending_tasks); + state->pending_tasks.clear(); /* Free the global scheduler state. */ - free(state); + delete state; } /* We need this code so we can clean up when we get a SIGTERM signal. */ @@ -139,20 +125,6 @@ void signal_handler(int signal) { /* End of the cleanup code. */ -LocalScheduler *get_local_scheduler(GlobalSchedulerState *state, - DBClientID local_scheduler_id) { - LocalScheduler *local_scheduler_ptr; - for (int i = 0; i < utarray_len(state->local_schedulers); ++i) { - local_scheduler_ptr = - (LocalScheduler *) utarray_eltptr(state->local_schedulers, i); - if (DBClientID_equal(local_scheduler_ptr->id, local_scheduler_id)) { - LOG_DEBUG("local_scheduler_id matched cached local scheduler entry."); - return local_scheduler_ptr; - } - } - return NULL; -} - void process_task_waiting(Task *waiting_task, void *user_context) { GlobalSchedulerState *state = (GlobalSchedulerState *) user_context; LOG_DEBUG("Task waiting callback is called."); @@ -163,7 +135,7 @@ void process_task_waiting(Task *waiting_task, void *user_context) { * resubmit the tasks in this array. */ if (!successfully_assigned) { Task *task_copy = Task_copy(waiting_task); - utarray_push_back(state->pending_tasks, &task_copy); + state->pending_tasks.push_back(task_copy); } } @@ -172,35 +144,11 @@ void add_local_scheduler(GlobalSchedulerState *state, const char *aux_address) { /* Add plasma_manager ip:port -> local_scheduler_db_client_id association to * state. */ - AuxAddressEntry *plasma_local_scheduler_entry = - (AuxAddressEntry *) calloc(1, sizeof(AuxAddressEntry)); - plasma_local_scheduler_entry->aux_address = strdup(aux_address); - plasma_local_scheduler_entry->local_scheduler_db_client_id = db_client_id; - HASH_ADD_KEYPTR(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, - plasma_local_scheduler_entry->aux_address, - strlen(plasma_local_scheduler_entry->aux_address), - plasma_local_scheduler_entry); + state->plasma_local_scheduler_map[std::string(aux_address)] = db_client_id; /* Add local_scheduler_db_client_id -> plasma_manager ip:port association to * state. */ - HASH_ADD(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, - local_scheduler_db_client_id, - sizeof(plasma_local_scheduler_entry->local_scheduler_db_client_id), - plasma_local_scheduler_entry); - -#if (RAY_COMMON_LOG_LEVEL <= RAY_COMMON_DEBUG) - { - /* Print the local scheduler to plasma association map so far. */ - AuxAddressEntry *entry, *tmp; - LOG_DEBUG("Local scheduler to plasma hash map so far:"); - HASH_ITER(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, - entry, tmp) { - LOG_DEBUG("%s -> %s", entry->aux_address, - ObjectID_to_string(entry->local_scheduler_db_client_id, - id_string, ID_STRING_SIZE)); - } - } -#endif + state->local_scheduler_plasma_map[db_client_id] = std::string(aux_address); /* Add new local scheduler to the state. */ LocalScheduler local_scheduler; @@ -214,44 +162,40 @@ void add_local_scheduler(GlobalSchedulerState *state, sizeof(local_scheduler.info.dynamic_resources)); memset(local_scheduler.info.static_resources, 0, sizeof(local_scheduler.info.static_resources)); - utarray_push_back(state->local_schedulers, &local_scheduler); + state->local_schedulers[db_client_id] = local_scheduler; /* Allow the scheduling algorithm to process this event. */ handle_new_local_scheduler(state, state->policy_state, db_client_id); } -void remove_local_scheduler(GlobalSchedulerState *state, int index) { - LocalScheduler *active_worker = - (LocalScheduler *) utarray_eltptr(state->local_schedulers, index); - DBClientID db_client_id = active_worker->id; - utarray_erase(state->local_schedulers, index, 1); +std::unordered_map::iterator +remove_local_scheduler( + GlobalSchedulerState *state, + std::unordered_map::iterator + it) { + CHECK(it != state->local_schedulers.end()); + DBClientID local_scheduler_id = it->first; + it = state->local_schedulers.erase(it); - AuxAddressEntry *entry, *tmp; - HASH_ITER(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, entry, - tmp) { - if (DBClientID_equal(entry->local_scheduler_db_client_id, db_client_id)) { - HASH_DELETE(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, - entry); - /* The hash entry is shared with the local_scheduler_plasma hashmap and - * will be freed there. */ - free(entry->aux_address); - } - } + /* Remove the local scheduler from the mappings. This code only makes sense if + * there is a one-to-one mapping between local schedulers and plasma managers. + */ + std::string aux_address = + state->local_scheduler_plasma_map[local_scheduler_id]; + state->local_scheduler_plasma_map.erase(local_scheduler_id); + state->plasma_local_scheduler_map.erase(aux_address); - HASH_FIND(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, - &db_client_id, sizeof(db_client_id), entry); - CHECK(entry != NULL); - HASH_DELETE(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, - entry); - free(entry); - - handle_local_scheduler_removed(state, state->policy_state, db_client_id); + handle_local_scheduler_removed(state, state->policy_state, + local_scheduler_id); + return it; } /** * Process a notification about a new DB client connecting to Redis. - * @param aux_address: an ip:port pair for the plasma manager associated with - * this db client. + * + * @param aux_address An ip:port pair for the plasma manager associated with + * this db client. + * @return Void. */ void process_new_db_client(DBClient *db_client, void *user_context) { GlobalSchedulerState *state = (GlobalSchedulerState *) user_context; @@ -261,31 +205,20 @@ void process_new_db_client(DBClient *db_client, void *user_context) { UNUSED(id_string); if (strncmp(db_client->client_type, "local_scheduler", strlen("local_scheduler")) == 0) { + bool local_scheduler_present = + (state->local_schedulers.find(db_client->id) != + state->local_schedulers.end()); if (db_client->is_insertion) { /* This is a notification for an insert. We may receive duplicate * notifications since we read the entire table before processing * notifications. Filter out local schedulers that we already added. */ - for (LocalScheduler *scheduler = - (LocalScheduler *) utarray_front(state->local_schedulers); - scheduler != NULL; scheduler = (LocalScheduler *) utarray_next( - state->local_schedulers, scheduler)) { - if (UNIQUE_ID_EQ(scheduler->id, db_client->id)) { - return; - } + if (!local_scheduler_present) { + add_local_scheduler(state, db_client->id, db_client->aux_address); } - - add_local_scheduler(state, db_client->id, db_client->aux_address); } else { - int i = 0; - for (; i < utarray_len(state->local_schedulers); ++i) { - LocalScheduler *active_worker = - (LocalScheduler *) utarray_eltptr(state->local_schedulers, i); - if (DBClientID_equal(active_worker->id, db_client->id)) { - break; - } - } - if (i < utarray_len(state->local_schedulers)) { - remove_local_scheduler(state, i); + if (local_scheduler_present) { + remove_local_scheduler(state, + state->local_schedulers.find(db_client->id)); } } } @@ -316,22 +249,14 @@ void object_table_subscribe_callback(ObjectID object_id, for (int i = 0; i < manager_count; i++) { LOG_DEBUG("\t\t%s", manager_vector[i]); } - SchedulerObjectInfo *obj_info_entry = NULL; - HASH_FIND(hh, state->scheduler_object_info_table, &object_id, - sizeof(object_id), obj_info_entry); - - if (obj_info_entry == NULL) { + if (state->scheduler_object_info_table.find(object_id) == + state->scheduler_object_info_table.end()) { /* Construct a new object info hash table entry. */ - obj_info_entry = - (SchedulerObjectInfo *) malloc(sizeof(SchedulerObjectInfo)); - memset(obj_info_entry, 0, sizeof(*obj_info_entry)); + SchedulerObjectInfo &obj_info_entry = + state->scheduler_object_info_table[object_id]; + obj_info_entry.data_size = data_size; - obj_info_entry->object_id = object_id; - obj_info_entry->data_size = data_size; - - HASH_ADD(hh, state->scheduler_object_info_table, object_id, - sizeof(obj_info_entry->object_id), obj_info_entry); LOG_DEBUG("New object added to object_info_table with id = %s", ObjectID_to_string(object_id, id_string, ID_STRING_SIZE)); LOG_DEBUG("\tmanager locations:"); @@ -340,15 +265,13 @@ void object_table_subscribe_callback(ObjectID object_id, } } - /* In all cases, replace the object location vector on each callback. */ - if (obj_info_entry->object_locations != NULL) { - utarray_free(obj_info_entry->object_locations); - obj_info_entry->object_locations = NULL; - } + SchedulerObjectInfo &obj_info_entry = + state->scheduler_object_info_table[object_id]; - utarray_new(obj_info_entry->object_locations, &ut_str_icd); + /* In all cases, replace the object location vector on each callback. */ + obj_info_entry.object_locations.clear(); for (int i = 0; i < manager_count; i++) { - utarray_push_back(obj_info_entry->object_locations, &manager_vector[i]); + obj_info_entry.object_locations.push_back(std::string(manager_vector[i])); } } @@ -367,12 +290,13 @@ void local_scheduler_table_handler(DBClientID client_id, "total workers = %d, task queue length = %d, available workers = %d", info.total_num_workers, info.task_queue_length, info.available_workers); /* Update the local scheduler info struct. */ - LocalScheduler *local_scheduler_ptr = get_local_scheduler(state, client_id); - if (local_scheduler_ptr != NULL) { + auto it = state->local_schedulers.find(client_id); + if (it != state->local_schedulers.end()) { /* Reset the number of tasks sent since the last heartbeat. */ - local_scheduler_ptr->num_heartbeats_missed = 0; - local_scheduler_ptr->num_recent_tasks_sent = 0; - local_scheduler_ptr->info = info; + LocalScheduler &local_scheduler = it->second; + local_scheduler.num_heartbeats_missed = 0; + local_scheduler.num_recent_tasks_sent = 0; + local_scheduler.info = info; } else { LOG_WARN("client_id didn't match any cached local scheduler entries"); } @@ -380,18 +304,20 @@ void local_scheduler_table_handler(DBClientID client_id, int task_cleanup_handler(event_loop *loop, timer_id id, void *context) { GlobalSchedulerState *state = (GlobalSchedulerState *) context; - /* Loop over the pending tasks and resubmit them. */ - int64_t num_pending_tasks = utarray_len(state->pending_tasks); - for (int64_t i = num_pending_tasks - 1; i >= 0; --i) { - Task **pending_task = (Task **) utarray_eltptr(state->pending_tasks, i); + /* Loop over the pending tasks in reverse order and resubmit them. */ + auto it = state->pending_tasks.end(); + while (it != state->pending_tasks.begin()) { + it--; + Task *pending_task = *it; /* Pretend that the task has been resubmitted. */ bool successfully_assigned = - handle_task_waiting(state, state->policy_state, *pending_task); + handle_task_waiting(state, state->policy_state, pending_task); if (successfully_assigned) { /* The task was successfully assigned, so remove it from this list and - * free it. */ - utarray_erase(state->pending_tasks, i, 1); - free(*pending_task); + * free it. This uses the fact that pending_tasks is a vector and so erase + * returns an iterator to the next element in the vector. */ + it = state->pending_tasks.erase(it); + Task_free(pending_task); } } @@ -405,20 +331,21 @@ int heartbeat_timeout_handler(event_loop *loop, timer_id id, void *context) { * up. */ /* TODO(swang): If the local scheduler hasn't actually died, then it should * clean up its state and exit upon receiving this notification. */ - LocalScheduler *local_scheduler_ptr; - for (int i = utarray_len(state->local_schedulers) - 1; i >= 0; --i) { - local_scheduler_ptr = - (LocalScheduler *) utarray_eltptr(state->local_schedulers, i); - if (local_scheduler_ptr->num_heartbeats_missed >= NUM_HEARTBEATS_TIMEOUT) { + auto it = state->local_schedulers.begin(); + while (it != state->local_schedulers.end()) { + if (it->second.num_heartbeats_missed >= NUM_HEARTBEATS_TIMEOUT) { LOG_WARN( "Missed too many heartbeats from local scheduler, marking as dead."); /* Notify others by updating the global state. */ - db_client_table_remove(state->db, local_scheduler_ptr->id, NULL, NULL, - NULL); - /* Remove the scheduler from the local state. */ - remove_local_scheduler(state, i); + db_client_table_remove(state->db, it->second.id, NULL, NULL, NULL); + /* Remove the scheduler from the local state. The call to + * remove_local_scheduler modifies the container in place and returns the + * next iterator. */ + it = remove_local_scheduler(state, it); + } else { + it->second.num_heartbeats_missed += 1; + it++; } - ++local_scheduler_ptr->num_heartbeats_missed; } /* Reset the timer. */ diff --git a/src/global_scheduler/global_scheduler.h b/src/global_scheduler/global_scheduler.h index 80a70e94a..ff98094b6 100644 --- a/src/global_scheduler/global_scheduler.h +++ b/src/global_scheduler/global_scheduler.h @@ -3,10 +3,10 @@ #include "task.h" +#include + #include "state/db.h" #include "state/local_scheduler_table.h" -#include "utarray.h" -#include "uthash.h" /* The frequency with which the global scheduler checks if there are any tasks * that haven't been scheduled yet. */ @@ -36,30 +36,12 @@ typedef struct GlobalSchedulerPolicyState GlobalSchedulerPolicyState; * This defines a hash table used to cache information about different objects. */ typedef struct { - /** The object ID in question. */ - ObjectID object_id; /** The size in bytes of the object. */ int64_t data_size; - /** An array of object locations for this object. */ - UT_array *object_locations; - /** Handle for the uthash table. */ - UT_hash_handle hh; + /** A vector of object locations for this object. */ + std::vector object_locations; } SchedulerObjectInfo; -/** - * A struct used for caching local scheduler to Plasma association. - */ -typedef struct { - /** IP:port string for the plasma_manager. */ - char *aux_address; - /** Local scheduler db client id. */ - DBClientID local_scheduler_db_client_id; - /** Plasma_manager ip:port -> local_scheduler_db_client_id. */ - UT_hash_handle plasma_local_scheduler_hh; - /** local_scheduler_db_client_id -> plasma_manager ip:port. */ - UT_hash_handle local_scheduler_plasma_hh; -} AuxAddressEntry; - /** * Global scheduler state structure. */ @@ -68,20 +50,22 @@ typedef struct { event_loop *loop; /** The global state store database. */ DBHandle *db; - /** The local schedulers that are connected to Redis. TODO(rkn): This probably - * needs to be a hashtable since we often look up the local_scheduler struct - * based on its db_client_id. */ - UT_array *local_schedulers; + /** A hash table mapping local scheduler ID to the local schedulers that are + * connected to Redis. */ + std::unordered_map + local_schedulers; /** The state managed by the scheduling policy. */ GlobalSchedulerPolicyState *policy_state; /** The plasma_manager ip:port -> local_scheduler_db_client_id association. */ - AuxAddressEntry *plasma_local_scheduler_map; + std::unordered_map plasma_local_scheduler_map; /** The local_scheduler_db_client_id -> plasma_manager ip:port association. */ - AuxAddressEntry *local_scheduler_plasma_map; + std::unordered_map + local_scheduler_plasma_map; /** Objects cached by this global scheduler instance. */ - SchedulerObjectInfo *scheduler_object_info_table; + std::unordered_map + scheduler_object_info_table; /** An array of tasks that haven't been scheduled yet. */ - UT_array *pending_tasks; + std::vector pending_tasks; } GlobalSchedulerState; /** diff --git a/src/global_scheduler/global_scheduler_algorithm.cc b/src/global_scheduler/global_scheduler_algorithm.cc index 37018cdd8..30c692f28 100644 --- a/src/global_scheduler/global_scheduler_algorithm.cc +++ b/src/global_scheduler/global_scheduler_algorithm.cc @@ -6,8 +6,7 @@ #include "global_scheduler_algorithm.h" GlobalSchedulerPolicyState *GlobalSchedulerPolicyState_init(void) { - GlobalSchedulerPolicyState *policy_state = - (GlobalSchedulerPolicyState *) malloc(sizeof(GlobalSchedulerPolicyState)); + GlobalSchedulerPolicyState *policy_state = new GlobalSchedulerPolicyState(); policy_state->round_robin_index = 0; int num_weight_elem = @@ -23,7 +22,7 @@ GlobalSchedulerPolicyState *GlobalSchedulerPolicyState_init(void) { } void GlobalSchedulerPolicyState_free(GlobalSchedulerPolicyState *policy_state) { - free(policy_state); + delete policy_state; } /** @@ -45,158 +44,6 @@ bool constraints_satisfied_hard(const LocalScheduler *scheduler, return true; } -/** - * This is a helper method that assigns a task to the next local scheduler in a - * round robin fashion. - */ -void handle_task_round_robin(GlobalSchedulerState *state, - GlobalSchedulerPolicyState *policy_state, - Task *task) { - CHECKM(utarray_len(state->local_schedulers) > 0, - "No local schedulers. We currently don't handle this case."); - LocalScheduler *scheduler = NULL; - TaskSpec *task_spec = Task_task_spec(task); - int i; - int num_retries = 1; - bool task_satisfied = false; - - for (i = policy_state->round_robin_index; !task_satisfied && num_retries > 0; - i = (i + 1) % utarray_len(state->local_schedulers)) { - if (i == policy_state->round_robin_index) { - num_retries--; - } - scheduler = (LocalScheduler *) utarray_eltptr(state->local_schedulers, i); - task_satisfied = constraints_satisfied_hard(scheduler, task_spec); - } - - if (task_satisfied) { - /* Update next index to try and assign the task. Note that the counter i has - * been advanced. */ - policy_state->round_robin_index = i; - assign_task_to_local_scheduler(state, task, scheduler->id); - } else { - /* TODO(atumanov): propagate the error to the driver, which submitted - * this impossible task and/or cache the task to consider when new - * local schedulers register. */ - } -} - -ObjectSizeEntry *create_object_size_hashmap(GlobalSchedulerState *state, - TaskSpec *task_spec, - bool *has_args_by_ref, - int64_t *task_data_size) { - ObjectSizeEntry *s = NULL, *object_size_table = NULL; - *task_data_size = 0; - - for (int i = 0; i < TaskSpec_num_args(task_spec); i++) { - /* Object ids are only available for args by references. - * Args by value are serialized into the TaskSpec itself. - * We will only concern ourselves with args by ref for data size calculation - */ - if (!TaskSpec_arg_by_ref(task_spec, i)) { - continue; - } - *has_args_by_ref = true; - ObjectID obj_id = TaskSpec_arg_id(task_spec, i); - /* Look up this object ID in the global scheduler object cache. */ - SchedulerObjectInfo *obj_info_entry = NULL; - HASH_FIND(hh, state->scheduler_object_info_table, &obj_id, sizeof(obj_id), - obj_info_entry); - if (obj_info_entry == NULL) { - /* Global scheduler doesn't know anything about this object ID, so skip - * it. */ - LOG_DEBUG("Processing task with object ID not known to global scheduler"); - continue; - } - LOG_DEBUG("[GS] found object id, data_size = %" PRId64, - obj_info_entry->data_size); - /* Object is known to the scheduler. For each of its locations, add size. */ - int64_t object_size = obj_info_entry->data_size; - /* Add each object's size to task's size. */ - *task_data_size += object_size; - char **p = NULL; - char id_string[ID_STRING_SIZE]; - LOG_DEBUG("locations for an arg_by_ref obj_id = %s", - ObjectID_to_string(obj_id, id_string, ID_STRING_SIZE)); - UNUSED(id_string); - for (p = (char **) utarray_front(obj_info_entry->object_locations); - p != NULL; - p = (char **) utarray_next(obj_info_entry->object_locations, p)) { - const char *object_location = *p; - - LOG_DEBUG("\tobject location: %s", object_location); - - /* Look up this location in the local object size hash table. */ - HASH_FIND_STR(object_size_table, object_location, s); - if (NULL == s) { - /* This location not yet known, so add this object location. */ - s = (ObjectSizeEntry *) calloc(1, sizeof(ObjectSizeEntry)); - s->object_location = object_location; - HASH_ADD_KEYPTR(hh, object_size_table, s->object_location, - strlen(s->object_location), s); - } - /* At this point the object location exists in our hash table. */ - s->total_object_size += object_size; - } /* End for each object's location. */ - } /* End for each task's object. */ - - return object_size_table; -} - -void free_object_size_hashmap(ObjectSizeEntry *object_size_table) { - /* Destroy local state. */ - ObjectSizeEntry *tmp, *s = NULL; - HASH_ITER(hh, object_size_table, s, tmp) { - HASH_DEL(object_size_table, s); - /* NOTE: Do not free externally stored s->object_location. */ - free(s); - } -} - -DBClientID get_local_scheduler_id(GlobalSchedulerState *state, - const char *plasma_location) { - AuxAddressEntry *aux_entry = NULL; - DBClientID local_scheduler_id = NIL_ID; - if (plasma_location != NULL) { - LOG_DEBUG("max object size location found : %s", plasma_location); - /* Lookup association of plasma location to local scheduler. */ - HASH_FIND(plasma_local_scheduler_hh, state->plasma_local_scheduler_map, - plasma_location, uthash_strlen(plasma_location), aux_entry); - if (aux_entry) { - LOG_DEBUG( - "found local scheduler db client association for plasma ip:port = %s", - aux_entry->aux_address); - /* Plasma to local scheduler db client ID association found, get local - * scheduler ID. */ - local_scheduler_id = aux_entry->local_scheduler_db_client_id; - } else { - LOG_ERROR( - "local scheduler db client association not found for plasma " - "ip:port=%s", - plasma_location); - } - } - - char id_string[ID_STRING_SIZE]; - LOG_DEBUG("local scheduler ID found = %s", - ObjectID_to_string(local_scheduler_id, id_string, ID_STRING_SIZE)); - UNUSED(id_string); - - if (IS_NIL_ID(local_scheduler_id)) { - return local_scheduler_id; - } - - /* Check to make sure this local_scheduler_db_client_id matches one of the - * schedulers. */ - LocalScheduler *local_scheduler_ptr = - get_local_scheduler(state, local_scheduler_id); - if (local_scheduler_ptr == NULL) { - LOG_WARN( - "local_scheduler_id didn't match any cached local scheduler entries"); - } - return local_scheduler_id; -} - double inner_product(double a[], double b[], int size) { double result = 0; for (int i = 0; i < size; i++) { @@ -205,43 +52,6 @@ double inner_product(double a[], double b[], int size) { return result; } -double calculate_object_size_fraction(GlobalSchedulerState *state, - LocalScheduler *scheduler, - ObjectSizeEntry *object_size_table, - int64_t total_task_object_size) { - /* Look up its cached object size in the hashmap, normalize by total object - * size for this task. */ - /* Aggregate object size for this task. */ - double object_size_fraction = 0; - if (total_task_object_size > 0) { - /* Does this node contribute anything to this task object size? */ - /* Lookup scheduler->id in local_scheduler_plasma_map to get plasma aux - * address, which is used as the key for object_size_table. This uses the - * plasma aux address to locate the object_size this node contributes. */ - AuxAddressEntry *local_scheduler_plasma_pair = NULL; - HASH_FIND(local_scheduler_plasma_hh, state->local_scheduler_plasma_map, - &(scheduler->id), sizeof(scheduler->id), - local_scheduler_plasma_pair); - if (local_scheduler_plasma_pair != NULL) { - ObjectSizeEntry *s = NULL; - /* Found this node's local scheduler to plasma mapping. Use the - * corresponding plasma key to see if this node has any cached objects for - * this task. */ - HASH_FIND_STR(object_size_table, local_scheduler_plasma_pair->aux_address, - s); - if (s != NULL) { - /* This node has some of this task's objects. Calculate what fraction. - */ - CHECK(strcmp(s->object_location, - local_scheduler_plasma_pair->aux_address) == 0); - object_size_fraction = - MIN(1, (double) (s->total_object_size) / total_task_object_size); - } - } - } - return object_size_fraction; -} - double calculate_score_dynvec_normalized(GlobalSchedulerState *state, LocalScheduler *scheduler, const TaskSpec *task_spec, @@ -268,9 +78,73 @@ double calculate_score_dynvec_normalized(GlobalSchedulerState *state, return score; } +int64_t locally_available_data_size(const GlobalSchedulerState *state, + DBClientID local_scheduler_id, + TaskSpec *task_spec) { + /* This function will compute the total size of all the object dependencies + * for the given task that are already locally available to the specified + * local scheduler. */ + int64_t task_data_size = 0; + + CHECK(state->local_scheduler_plasma_map.count(local_scheduler_id) == 1); + + const std::string &plasma_manager = + state->local_scheduler_plasma_map.at(local_scheduler_id); + + /* TODO(rkn): Note that if the same object ID appears as multiple arguments, + * then it will be overcounted. */ + for (int64_t i = 0; i < TaskSpec_num_args(task_spec); ++i) { + if (!TaskSpec_arg_by_ref(task_spec, i)) { + /* Ignore arguments that are not object IDs since these are serialized as + * part of the task spec and so they don't require any data transfer. */ + continue; + } + + ObjectID object_id = TaskSpec_arg_id(task_spec, i); + + if (state->scheduler_object_info_table.count(object_id) == 0) { + /* If this global scheduler is not aware of this object ID, then ignore + * it.*/ + continue; + } + + const SchedulerObjectInfo &object_size_info = + state->scheduler_object_info_table.at(object_id); + + if (std::find(object_size_info.object_locations.begin(), + object_size_info.object_locations.end(), + plasma_manager) == object_size_info.object_locations.end()) { + /* This local scheduler does not have access to this object, so don't + * count this object. */ + continue; + } + + /* Look at the size of the object. */ + int64_t object_size = object_size_info.data_size; + if (object_size == -1) { + /* This means that this global scheduler does not know the object size + * yet, so assume that the object is one megabyte. TODO(rkn): Maybe we + * should instead use the average object size. */ + object_size = 1000000; + } + + /* If we get here, then this local scheduler has access to this object, so + * count the contribution of this object. */ + task_data_size += object_size; + } + + return task_data_size; +} + double calculate_cost_pending(const GlobalSchedulerState *state, - const LocalScheduler *scheduler) { - /* TODO: make sure that num_recent_tasks_sent is reset on each heartbeat. */ + const LocalScheduler *scheduler, + TaskSpec *task_spec) { + /* Calculate how much data is already present on this machine. TODO(rkn): Note + * that this information is not being used yet. Fix this. */ + int64_t data_size = + locally_available_data_size(state, scheduler->id, task_spec); + /* TODO(rkn): This logic does not load balance properly when the different + * machines have different sizes. Fix this. */ return scheduler->num_recent_tasks_sent + scheduler->info.task_queue_length; } @@ -281,41 +155,33 @@ bool handle_task_waiting(GlobalSchedulerState *state, CHECKM(task_spec != NULL, "task wait handler encounted a task with NULL spec"); - /* Local hash table to keep track of aggregate object sizes per local - * scheduler. */ - ObjectSizeEntry *object_size_table = NULL; - bool has_args_by_ref = false; + bool task_feasible = false; /* The total size of the task's data. */ int64_t task_object_size = 0; - object_size_table = create_object_size_hashmap( - state, task_spec, &has_args_by_ref, &task_object_size); - /* Go through all the nodes, calculate the score for each, pick max score. */ LocalScheduler *scheduler = NULL; double best_local_scheduler_score = INT32_MIN; CHECKM(best_local_scheduler_score < 0, "We might have a floating point underflow"); DBClientID best_local_scheduler_id = NIL_ID; /* best node to send this task */ - for (scheduler = (LocalScheduler *) utarray_front(state->local_schedulers); - scheduler != NULL; scheduler = (LocalScheduler *) utarray_next( - state->local_schedulers, scheduler)) { + for (auto it = state->local_schedulers.begin(); + it != state->local_schedulers.end(); it++) { /* For each local scheduler, calculate its score. Check hard constraints * first. */ + LocalScheduler *scheduler = &(it->second); if (!constraints_satisfied_hard(scheduler, task_spec)) { continue; } task_feasible = true; /* This node satisfies the hard capacity constraint. Calculate its score. */ - double score = -1 * calculate_cost_pending(state, scheduler); + double score = -1 * calculate_cost_pending(state, scheduler, task_spec); if (score > best_local_scheduler_score) { best_local_scheduler_score = score; best_local_scheduler_id = scheduler->id; } - } /* For each local scheduler. */ - - free_object_size_hashmap(object_size_table); + } if (!task_feasible) { char id_string[ID_STRING_SIZE]; diff --git a/src/global_scheduler/global_scheduler_algorithm.h b/src/global_scheduler/global_scheduler_algorithm.h index d57b53402..9d7e9335f 100644 --- a/src/global_scheduler/global_scheduler_algorithm.h +++ b/src/global_scheduler/global_scheduler_algorithm.h @@ -26,12 +26,6 @@ struct GlobalSchedulerPolicyState { double resource_attribute_weight[ResourceIndex_MAX + 1]; }; -typedef struct { - const char *object_location; - int64_t total_object_size; - UT_hash_handle hh; -} ObjectSizeEntry; - /** * Create the state of the global scheduler policy. This state must be freed by * the caller.