Allow starting multiple local schedulers. (#86)

src/global_scheduler/test/test.py

@@ -50,11 +50,13 @@ class TestGlobalScheduler(unittest.TestCase):
     # Create a Redis client.
     self.redis_client = redis.StrictRedis(host=node_ip_address, port=redis_port)
     # Start the global scheduler.
-    self.p1 = global_scheduler.start_global_scheduler(redis_address, USE_VALGRIND)
+    self.p1 = global_scheduler.start_global_scheduler(redis_address, use_valgrind=USE_VALGRIND)
     # 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)
     # Start the local scheduler.
-    local_scheduler_name, self.p3 = photon.start_local_scheduler(plasma_store_name, redis_address=redis_address)
+    local_scheduler_name, self.p4 = photon.start_local_scheduler(plasma_store_name, plasma_manager_name=plasma_manager_name, redis_address=redis_address)
     # Connect to the scheduler.
     self.photon_client = photon.PhotonClient(local_scheduler_name)
 
@@ -68,16 +70,17 @@ class TestGlobalScheduler(unittest.TestCase):
       self.p1.kill()
     self.p2.kill()
     self.p3.kill()
+    self.p4.kill()
     # Kill Redis. In the event that we are using valgrind, this needs to happen
     # after we kill the global scheduler.
     self.redis_process.kill()
 
   def test_redis_contents(self):
-    # There should be two db clients, the global scheduler and the local
-    # scheduler.
-    self.assertEqual(len(self.redis_client.keys("db_clients*")), 2)
+    # There should be two db clients, the global scheduler, the local scheduler,
+    # and the plasma manager.
+    self.assertEqual(len(self.redis_client.keys("db_clients*")), 3)
     # There should not be anything else in Redis yet.
-    self.assertEqual(len(self.redis_client.keys("*")), 2)
+    self.assertEqual(len(self.redis_client.keys("*")), 3)
 
     # Submit a task to Redis.
     task = photon.Task(random_function_id(), [], 0, random_task_id(), 0)

src/photon/photon/photon_services.py

@@ -8,11 +8,14 @@ import time
 def random_name():
   return str(random.randint(0, 99999999))
 
-def start_local_scheduler(plasma_store_name, redis_address=None, use_valgrind=False, use_profiler=False):
+def start_local_scheduler(plasma_store_name, plasma_manager_name=None, redis_address=None, use_valgrind=False, use_profiler=False):
   """Start a local scheduler process.
 
   Args:
     plasma_store_name (str): The name of the plasma store socket to connect to.
+    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.
     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
@@ -24,11 +27,15 @@ def start_local_scheduler(plasma_store_name, redis_address=None, use_valgrind=Fa
     A tuple of the name of the local scheduler socket and the process ID of the
       local scheduler process.
   """
+  if (plasma_manager_name == None) != (redis_address == None):
+    raise Exception("If one of the plasma_manager_name and the redis_address is provided, then both must be provided.")
   if use_valgrind and use_profiler:
     raise Exception("Cannot use valgrind and profiler at the same time.")
   local_scheduler_executable = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../build/photon_scheduler")
   local_scheduler_name = "/tmp/scheduler{}".format(random_name())
   command = [local_scheduler_executable, "-s", local_scheduler_name, "-p", plasma_store_name]
+  if plasma_manager_name is not None:
+    command += ["-m", plasma_manager_name]
   if redis_address is not None:
     command += ["-r", redis_address]
   if use_valgrind:

src/photon/photon_scheduler.c

@@ -23,21 +23,14 @@ UT_icd worker_icd = {sizeof(worker), NULL, NULL, NULL};
 
 UT_icd byte_icd = {sizeof(uint8_t), NULL, NULL, NULL};
 
-local_scheduler_state *init_local_scheduler(event_loop *loop,
-                                            const char *redis_addr,
-                                            int redis_port,
-                                            const char *plasma_socket_name) {
+local_scheduler_state *init_local_scheduler(
+    event_loop *loop,
+    const char *redis_addr,
+    int redis_port,
+    const char *plasma_store_socket_name,
+    const char *plasma_manager_socket_name) {
   local_scheduler_state *state = malloc(sizeof(local_scheduler_state));
   state->loop = loop;
-  /* Connect to Plasma. This method will retry if Plasma hasn't started yet.
-   * Pass in a NULL manager address and port. */
-  state->plasma_conn =
-      plasma_connect(plasma_socket_name, NULL, PLASMA_DEFAULT_RELEASE_DELAY);
-  /* Subscribe to notifications about sealed objects. */
-  int plasma_fd = plasma_subscribe(state->plasma_conn);
-  /* Add the callback that processes the notification to the event loop. */
-  event_loop_add_file(loop, plasma_fd, EVENT_LOOP_READ,
-                      process_plasma_notification, state);
   state->worker_index = NULL;
   /* Add scheduler info. */
   utarray_new(state->workers, &worker_icd);
@@ -48,6 +41,16 @@ local_scheduler_state *init_local_scheduler(event_loop *loop,
   } else {
     state->db = NULL;
   }
+  /* Connect to Plasma. This method will retry if Plasma hasn't started yet.
+   * Pass in a NULL manager address and port. */
+  state->plasma_conn =
+      plasma_connect(plasma_store_socket_name, plasma_manager_socket_name,
+                     PLASMA_DEFAULT_RELEASE_DELAY);
+  /* Subscribe to notifications about sealed objects. */
+  int plasma_fd = plasma_subscribe(state->plasma_conn);
+  /* Add the callback that processes the notification to the event loop. */
+  event_loop_add_file(loop, plasma_fd, EVENT_LOOP_READ,
+                      process_plasma_notification, state);
   /* Add scheduler state. */
   state->algorithm_state = make_scheduling_algorithm_state();
   utarray_new(state->input_buffer, &byte_icd);
@@ -187,11 +190,13 @@ void handle_task_scheduled_callback(task *original_task, void *user_context) {
 void start_server(const char *socket_name,
                   const char *redis_addr,
                   int redis_port,
-                  const char *plasma_socket_name) {
+                  const char *plasma_store_socket_name,
+                  const char *plasma_manager_socket_name) {
   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_socket_name);
+  g_state = init_local_scheduler(loop, redis_addr, redis_port,
+                                 plasma_store_socket_name,
+                                 plasma_manager_socket_name);
 
   /* Register a callback for registering new clients. */
   event_loop_add_file(loop, fd, EVENT_LOOP_READ, new_client_connection,
@@ -221,9 +226,11 @@ int main(int argc, char *argv[]) {
   /* IP address and port of redis. */
   char *redis_addr_port = NULL;
   /* Socket name for the local Plasma store. */
-  char *plasma_socket_name = NULL;
+  char *plasma_store_socket_name = NULL;
+  /* Socket name for the local Plasma manager. */
+  char *plasma_manager_socket_name = NULL;
   int c;
-  while ((c = getopt(argc, argv, "s:r:p:")) != -1) {
+  while ((c = getopt(argc, argv, "s:r:p:m:")) != -1) {
     switch (c) {
     case 's':
       scheduler_socket_name = optarg;
@@ -232,7 +239,10 @@ int main(int argc, char *argv[]) {
       redis_addr_port = optarg;
       break;
     case 'p':
-      plasma_socket_name = optarg;
+      plasma_store_socket_name = optarg;
+      break;
+    case 'm':
+      plasma_manager_socket_name = optarg;
       break;
     default:
       LOG_FATAL("unknown option %c", c);
@@ -241,13 +251,20 @@ int main(int argc, char *argv[]) {
   if (!scheduler_socket_name) {
     LOG_FATAL("please specify socket for incoming connections with -s switch");
   }
-  if (!plasma_socket_name) {
-    LOG_FATAL("please specify socket for connecting to Plasma with -p switch");
+  if (!plasma_store_socket_name) {
+    LOG_FATAL(
+        "please specify socket for connecting to Plasma store with -p switch");
   }
   if (!redis_addr_port) {
     /* Start the local scheduler without connecting to Redis. In this case, all
      * submitted tasks will be queued and scheduled locally. */
-    start_server(scheduler_socket_name, NULL, -1, plasma_socket_name);
+    if (plasma_manager_socket_name) {
+      LOG_FATAL(
+          "if a plasma manager socket name is provided with the -m switch, "
+          "then a redis address must be provided with the -r switch");
+    }
+    start_server(scheduler_socket_name, NULL, -1, plasma_store_socket_name,
+                 NULL);
   } else {
     /* Parse the Redis address into an IP address and a port. */
     char redis_addr[16] = {0};
@@ -259,7 +276,12 @@ int main(int argc, char *argv[]) {
           "if a redis address is provided with the -r switch, it should be "
           "formatted like 127.0.0.1:6379");
     }
+    if (!plasma_manager_socket_name) {
+      LOG_FATAL(
+          "please specify socket for connecting to Plasma manager with -m "
+          "switch");
+    }
     start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port),
-                 plasma_socket_name);
+                 plasma_store_socket_name, plasma_manager_socket_name);
   }
 }