diff --git a/lib/python/ray/services.py b/lib/python/ray/services.py index bfc789b78..57367fc72 100644 --- a/lib/python/ray/services.py +++ b/lib/python/ray/services.py @@ -115,7 +115,7 @@ def start_global_scheduler(redis_address, cleanup=True): if cleanup: all_processes.append(p) -def start_local_scheduler(redis_address, plasma_store_name, plasma_manager_name, cleanup=True): +def start_local_scheduler(redis_address, plasma_store_name, plasma_manager_name, plasma_address=None, cleanup=True): """Start a local scheduler process. Args: @@ -130,7 +130,7 @@ def start_local_scheduler(redis_address, plasma_store_name, plasma_manager_name, Return: The name of the local scheduler socket. """ - local_scheduler_name, p = photon.start_local_scheduler(plasma_store_name, plasma_manager_name, redis_address=redis_address, use_profiler=RUN_PHOTON_PROFILER) + local_scheduler_name, p = photon.start_local_scheduler(plasma_store_name, plasma_manager_name, redis_address=redis_address, plasma_address=plasma_address, use_profiler=RUN_PHOTON_PROFILER) if cleanup: all_processes.append(p) return local_scheduler_name @@ -243,7 +243,8 @@ def start_ray_local(node_ip_address="127.0.0.1", num_workers=0, num_local_schedu object_store_manager_names.append(object_store_manager_name) time.sleep(0.1) # Start the local scheduler. - local_scheduler_name = start_local_scheduler(redis_address, object_store_name, object_store_manager_name, cleanup=True) + plasma_address = "{}:{}".format(node_ip_address, object_store_manager_port) + local_scheduler_name = start_local_scheduler(redis_address, object_store_name, object_store_manager_name, plasma_address=plasma_address, cleanup=True) local_scheduler_names.append(local_scheduler_name) time.sleep(0.1) # Aggregate the address information together. diff --git a/src/common/state/db.h b/src/common/state/db.h index 2f0032604..50634e68f 100644 --- a/src/common/state/db.h +++ b/src/common/state/db.h @@ -25,7 +25,12 @@ db_handle *db_connect(const char *db_address, const char *client_type, const char *client_addr, int client_port); - +db_handle *db_connect_extended(const char *db_address, + int db_port, + const char *client_type, + const char *client_addr, + int client_port, + const char *aux_address); /** * Attach global system store connection to an event loop. Callbacks from * queries to the global system store will trigger events in the event loop. diff --git a/src/common/state/db_client_table.h b/src/common/state/db_client_table.h index 3ea3949a3..2f5ffb4ac 100644 --- a/src/common/state/db_client_table.h +++ b/src/common/state/db_client_table.h @@ -14,6 +14,7 @@ typedef void (*db_client_table_done_callback)(db_client_id db_client_id, /* Callback for subscribing to the db client table. */ typedef void (*db_client_table_subscribe_callback)(db_client_id db_client_id, const char *client_type, + const char *aux_address, void *user_context); /** diff --git a/src/common/state/redis.c b/src/common/state/redis.c index ab3bcd60e..941530e60 100644 --- a/src/common/state/redis.c +++ b/src/common/state/redis.c @@ -134,11 +134,26 @@ db_handle *db_connect(const char *address, const char *client_type, const char *client_addr, int client_port) { + return db_connect_extended(address, port, client_type, client_addr, + client_port, ":"); +} + +db_handle *db_connect_extended(const char *address, + int port, + const char *client_type, + const char *client_addr, + int client_port, + const char *aux_address) { db_handle *db = malloc(sizeof(db_handle)); /* Sync connection for initial handshake */ redisReply *reply; int connection_attempts = 0; redisContext *context = redisConnect(address, port); + /* Sanity check aux_address. */ + if (aux_address == NULL || strlen(aux_address) == 0) { + LOG_WARN("db_connect: received empty aux_address, replacing with ':'"); + aux_address = ":"; + } while (context == NULL || context->err) { if (connection_attempts >= REDIS_DB_CONNECT_RETRIES) { break; @@ -155,7 +170,7 @@ db_handle *db_connect(const char *address, address, port); /* Enable keyspace events. */ reply = redisCommand(context, "CONFIG SET notify-keyspace-events AKE"); - CHECK(reply != NULL); + CHECKM(reply != NULL, "db_connect failed on CONFIG SET"); freeReplyObject(reply); /* Add new client using optimistic locking. */ db_client_id client = globally_unique_id(); @@ -166,16 +181,29 @@ db_handle *db_connect(const char *address, freeReplyObject(reply); reply = redisCommand(context, "MULTI"); freeReplyObject(reply); - reply = redisCommand( - context, - "HMSET db_clients:%b client_type %s address %s:%d db_client_id %b", - (char *) client.id, sizeof(client.id), client_type, client_addr, - client_port, (char *) client.id, sizeof(client.id)); - freeReplyObject(reply); - reply = redisCommand(context, "PUBLISH db_clients %b:%s", - (char *) client.id, sizeof(client.id), client_type); + reply = redisCommand(context, + "HMSET db_clients:%b client_type %s address %s:%d " + "db_client_id %b aux_address %s", + (char *) client.id, sizeof(client.id), client_type, + client_addr, client_port, (char *) client.id, + sizeof(client.id), aux_address); + CHECKM(reply != NULL, "db_connect failed on HMSET"); freeReplyObject(reply); + + { + UT_string *tmpbuf; + utstring_new(tmpbuf); + utstring_printf(tmpbuf, "%s %s", client_type, aux_address); + reply = + redisCommand(context, "PUBLISH db_clients %b:%s", (char *) client.id, + sizeof(client.id), utstring_body(tmpbuf)); + CHECKM(reply != NULL, "db_connect failed on PUBLISH"); + freeReplyObject(reply); + utstring_free(tmpbuf); + } + reply = redisCommand(context, "EXEC"); + CHECKM(reply != NULL, "db_connect failed on EXEC"); CHECK(reply); if (reply->type != REDIS_REPLY_NIL) { freeReplyObject(reply); @@ -550,7 +578,8 @@ void redis_get_cached_db_client(db_handle *db, redisReply *reply = redisCommand(db->sync_context, "HGET db_clients:%b address", (char *) db_client_id.id, sizeof(db_client_id.id)); - CHECK(reply->type == REDIS_REPLY_STRING); + CHECKM(reply->type == REDIS_REPLY_STRING, "REDIS reply type=%d", + reply->type); entry = malloc(sizeof(db_client_cache_entry)); entry->db_client_id = db_client_id; entry->addr = strdup(reply->str); @@ -963,11 +992,21 @@ void redis_db_client_table_subscribe_callback(redisAsyncContext *c, * client_type string, and we add 1 to null-terminate the string. */ int client_type_length = payload->len - 1 - sizeof(client.id) + 1; char *client_type = malloc(client_type_length); - memcpy(client_type, &payload->str[1 + sizeof(client.id)], client_type_length); + char *aux_address = malloc(client_type_length); + memset(aux_address, 0, client_type_length); + /* Published message format: */ + int rv = sscanf(&payload->str[1 + sizeof(client.id)], "%s %s", client_type, + aux_address); + CHECKM(rv == 2, + "redis_db_client_table_subscribe_callback: expected 2 parsed args, " + "Got %d instead.", + rv); if (data->subscribe_callback) { - data->subscribe_callback(client, client_type, data->subscribe_context); + data->subscribe_callback(client, client_type, aux_address, + data->subscribe_context); } free(client_type); + free(aux_address); } void redis_db_client_table_subscribe(table_callback_data *callback_data) { diff --git a/src/global_scheduler/global_scheduler.c b/src/global_scheduler/global_scheduler.c index ccaca61de..4c8901bba 100644 --- a/src/global_scheduler/global_scheduler.c +++ b/src/global_scheduler/global_scheduler.c @@ -30,6 +30,8 @@ global_scheduler_state *init_global_scheduler(event_loop *loop, const char *redis_addr, int redis_port) { global_scheduler_state *state = malloc(sizeof(global_scheduler_state)); + /* Must initialize state to 0. Sets hashmap head(s) to NULL. */ + memset(state, 0, sizeof(global_scheduler_state)); state->db = db_connect(redis_addr, redis_port, "global_scheduler", "", -1); db_attach(state->db, loop, false); utarray_new(state->local_schedulers, &local_scheduler_icd); @@ -38,9 +40,18 @@ global_scheduler_state *init_global_scheduler(event_loop *loop, } void free_global_scheduler(global_scheduler_state *state) { + aux_address_entry *entry, *tmp; + db_disconnect(state->db); utarray_free(state->local_schedulers); destroy_global_scheduler_policy(state->policy_state); + /* delete the plasma 2 photon association map */ + HASH_ITER(hh, state->plasma_photon_map, entry, tmp) { + HASH_DELETE(hh, state->plasma_photon_map, entry); + /* Now deallocate hash table entry. */ + free(entry->aux_address); + free(entry); + } free(state); } @@ -62,11 +73,23 @@ void process_task_waiting(task *task, void *user_context) { handle_task_waiting(state, state->policy_state, task); } +/** + * 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. + */ void process_new_db_client(db_client_id db_client_id, const char *client_type, + const char *aux_address, void *user_context) { global_scheduler_state *state = (global_scheduler_state *) user_context; - if (strcmp(client_type, "photon") == 0) { + if (strncmp(client_type, "photon", strlen("photon")) == 0) { + /* Add plasma_manager ip:port -> photon_db_client_id association to state. + */ + aux_address_entry *plasma_photon_entry = malloc(sizeof(aux_address_entry)); + plasma_photon_entry->aux_address = strdup(aux_address); + plasma_photon_entry->photon_db_client_id = db_client_id; + HASH_ADD_STR(state->plasma_photon_map, aux_address, plasma_photon_entry); handle_new_local_scheduler(state, state->policy_state, db_client_id); } } diff --git a/src/global_scheduler/global_scheduler.h b/src/global_scheduler/global_scheduler.h index 435f36d1b..00db7fac4 100644 --- a/src/global_scheduler/global_scheduler.h +++ b/src/global_scheduler/global_scheduler.h @@ -5,6 +5,7 @@ #include "state/db.h" #include "utarray.h" +#include "uthash.h" /** Contains all information that is associated with a local scheduler. */ typedef struct { @@ -14,6 +15,12 @@ typedef struct { typedef struct global_scheduler_policy_state global_scheduler_policy_state; +typedef struct { + char *aux_address; /* Key */ + db_client_id photon_db_client_id; + UT_hash_handle hh; +} aux_address_entry; + typedef struct { /** The global scheduler event loop. */ event_loop *loop; @@ -23,6 +30,7 @@ typedef struct { UT_array *local_schedulers; /** The state managed by the scheduling policy. */ global_scheduler_policy_state *policy_state; + aux_address_entry *plasma_photon_map; } global_scheduler_state; void assign_task_to_local_scheduler(global_scheduler_state *state, diff --git a/src/global_scheduler/test/test.py b/src/global_scheduler/test/test.py index 405ae4d14..4cd16f677 100644 --- a/src/global_scheduler/test/test.py +++ b/src/global_scheduler/test/test.py @@ -56,9 +56,10 @@ class TestGlobalScheduler(unittest.TestCase): # Start the Plasma store. plasma_store_name, self.p2 = plasma.start_plasma_store() # Start the Plasma manager. - plasma_manager_name, self.p3, _ = plasma.start_plasma_manager(plasma_store_name, redis_address) + plasma_manager_name, self.p3, plasma_manager_port = plasma.start_plasma_manager(plasma_store_name, redis_address) + plasma_address = "{}:{}".format(node_ip_address, plasma_manager_port) # Start the local scheduler. - local_scheduler_name, self.p4 = photon.start_local_scheduler(plasma_store_name, plasma_manager_name=plasma_manager_name, redis_address=redis_address) + local_scheduler_name, self.p4 = photon.start_local_scheduler(plasma_store_name, plasma_manager_name=plasma_manager_name, plasma_address=plasma_address, redis_address=redis_address) # Connect to the scheduler. self.photon_client = photon.PhotonClient(local_scheduler_name) diff --git a/src/photon/photon/photon_services.py b/src/photon/photon/photon_services.py index 59db729a4..ec66f52da 100644 --- a/src/photon/photon/photon_services.py +++ b/src/photon/photon/photon_services.py @@ -10,7 +10,7 @@ import time def random_name(): return str(random.randint(0, 99999999)) -def start_local_scheduler(plasma_store_name, plasma_manager_name=None, redis_address=None, use_valgrind=False, use_profiler=False): +def start_local_scheduler(plasma_store_name, plasma_manager_name=None, plasma_address=None, redis_address=None, use_valgrind=False, use_profiler=False): """Start a local scheduler process. Args: @@ -18,6 +18,9 @@ def start_local_scheduler(plasma_store_name, plasma_manager_name=None, redis_add plasma_manager_name (str): The name of the plasma manager to connect to. This does not need to be provided, but if it is, then the Redis address must be provided as well. + plasma_address (str): The address of the plasma manager to connect to. This + is only used by the global scheduler to figure out which plasma managers + are connected to which local schedulers. redis_address (str): The address of the Redis instance to connect to. If this is not provided, then the local scheduler will not connect to Redis. use_valgrind (bool): True if the local scheduler should be started inside of @@ -40,6 +43,8 @@ def start_local_scheduler(plasma_store_name, plasma_manager_name=None, redis_add command += ["-m", plasma_manager_name] if redis_address is not None: command += ["-r", redis_address] + if plasma_address is not None: + command += ["-a", plasma_address] if use_valgrind: pid = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full", "--show-leak-kinds=all", "--error-exitcode=1"] + command) time.sleep(1.0) diff --git a/src/photon/photon_scheduler.c b/src/photon/photon_scheduler.c index d44265cd2..a66eabe83 100644 --- a/src/photon/photon_scheduler.c +++ b/src/photon/photon_scheduler.c @@ -30,6 +30,7 @@ local_scheduler_state *init_local_scheduler( int redis_port, const char *plasma_store_socket_name, const char *plasma_manager_socket_name, + const char *plasma_manager_address, bool global_scheduler_exists) { local_scheduler_state *state = malloc(sizeof(local_scheduler_state)); state->loop = loop; @@ -38,7 +39,8 @@ local_scheduler_state *init_local_scheduler( utarray_new(state->workers, &worker_icd); /* Connect to Redis if a Redis address is provided. */ if (redis_addr != NULL) { - state->db = db_connect(redis_addr, redis_port, "photon", "", -1); + state->db = db_connect_extended(redis_addr, redis_port, "photon", "", -1, + plasma_manager_address); db_attach(state->db, loop, false); } else { state->db = NULL; @@ -151,11 +153,10 @@ void reconstruct_object_task_lookup_callback(object_id reconstruct_object_id, handle_task_submitted(state, state->algorithm_state, spec); } -void reconstruct_object_object_lookup_callback( - object_id reconstruct_object_id, - int manager_count, - const char *manager_vector[], - void *user_context) { +void reconstruct_object_object_lookup_callback(object_id reconstruct_object_id, + int manager_count, + const char *manager_vector[], + void *user_context) { /* Only continue reconstruction if we find that the object doesn't exist on * any nodes. NOTE: This codepath is not responsible for checking if the * object table entry is up-to-date. */ @@ -263,12 +264,14 @@ void start_server(const char *socket_name, int redis_port, const char *plasma_store_socket_name, const char *plasma_manager_socket_name, + const char *plasma_manager_address, bool global_scheduler_exists) { int fd = bind_ipc_sock(socket_name, true); event_loop *loop = event_loop_create(); - g_state = init_local_scheduler( - loop, redis_addr, redis_port, plasma_store_socket_name, - plasma_manager_socket_name, global_scheduler_exists); + g_state = + init_local_scheduler(loop, redis_addr, redis_port, + plasma_store_socket_name, plasma_manager_socket_name, + plasma_manager_address, global_scheduler_exists); /* Register a callback for registering new clients. */ event_loop_add_file(loop, fd, EVENT_LOOP_READ, new_client_connection, @@ -304,9 +307,11 @@ int main(int argc, char *argv[]) { char *plasma_store_socket_name = NULL; /* Socket name for the local Plasma manager. */ char *plasma_manager_socket_name = NULL; + /* Address for the plasma manager associated with this Photon instance. */ + char *plasma_manager_address = NULL; int c; bool global_scheduler_exists = true; - while ((c = getopt(argc, argv, "s:r:p:m:g:")) != -1) { + while ((c = getopt(argc, argv, "s:r:p:m:ga:")) != -1) { switch (c) { case 's': scheduler_socket_name = optarg; @@ -323,6 +328,9 @@ int main(int argc, char *argv[]) { case 'g': global_scheduler_exists = false; break; + case 'a': + plasma_manager_address = optarg; + break; default: LOG_FATAL("unknown option %c", c); } @@ -343,7 +351,7 @@ int main(int argc, char *argv[]) { "then a redis address must be provided with the -r switch"); } start_server(scheduler_socket_name, NULL, -1, plasma_store_socket_name, - NULL, global_scheduler_exists); + NULL, plasma_manager_address, global_scheduler_exists); } else { /* Parse the Redis address into an IP address and a port. */ char redis_addr[16] = {0}; @@ -362,7 +370,7 @@ int main(int argc, char *argv[]) { } start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port), plasma_store_socket_name, plasma_manager_socket_name, - global_scheduler_exists); + plasma_manager_address, global_scheduler_exists); } } #endif diff --git a/src/photon/photon_scheduler.h b/src/photon/photon_scheduler.h index 3be85948b..1fbcad4e2 100644 --- a/src/photon/photon_scheduler.h +++ b/src/photon/photon_scheduler.h @@ -70,6 +70,7 @@ local_scheduler_state *init_local_scheduler( int redis_port, const char *plasma_manager_socket_name, const char *plasma_store_socket_name, + const char *plasma_manager_address, bool global_scheduler_exists); void free_local_scheduler(local_scheduler_state *state); diff --git a/src/photon/test/photon_tests.c b/src/photon/test/photon_tests.c index 0dd5ce948..bbf2735f2 100644 --- a/src/photon/test/photon_tests.c +++ b/src/photon/test/photon_tests.c @@ -57,10 +57,10 @@ photon_mock *init_photon_mock() { UT_string *photon_socket_name = bind_ipc_sock_retry(photon_socket_name_format, &mock->photon_fd); CHECK(mock->plasma_fd >= 0 && mock->photon_fd >= 0); - mock->photon_state = - init_local_scheduler(mock->loop, redis_addr, redis_port, - utstring_body(plasma_manager_socket_name), - utstring_body(plasma_store_socket_name), false); + mock->photon_state = init_local_scheduler( + mock->loop, redis_addr, redis_port, + utstring_body(plasma_manager_socket_name), + utstring_body(plasma_store_socket_name), NULL, false); /* Connect a Photon client. */ mock->conn = photon_connect(utstring_body(photon_socket_name)); new_client_connection(mock->loop, mock->photon_fd,