Clean up when a driver disconnects. (#462)

* Clean up state when drivers exit.

* Remove unnecessary field in ActorMapEntry struct.

* Have monitor release GPU resources in Redis when driver exits.

* Enable multiple drivers in multi-node tests and test driver cleanup.

* Make redis GPU allocation a redis transaction and small cleanups.

* Fix multi-node test.

* Small cleanups.

* Make global scheduler take node_ip_address so it appears in the right place in the client table.

* Cleanups.

* Fix linting and cleanups in local scheduler.

* Fix removed_driver_test.

* Fix bug related to vector -> list.

* Fix linting.

* Cleanup.

* Fix multi node tests.

* Fix jenkins tests.

* Add another multi node test with many drivers.

* Fix linting.

* Make the actor creation notification a flatbuffer message.

* Revert "Make the actor creation notification a flatbuffer message."

This reverts commit af99099c8084dbf9177fb4e34c0c9b1a12c78f39.

* Add comment explaining flatbuffer problems.

test/jenkins_tests/multi_node_docker_test.py

@@ -3,6 +3,7 @@ from __future__ import division
 from __future__ import print_function
 
 import argparse
+import numpy as np
 import os
 import re
 import subprocess
@@ -85,8 +86,8 @@ class DockerRunner(object):
     else:
       return m.group(1)
 
-  def _start_head_node(self, docker_image, mem_size, shm_size,
-                       development_mode):
+  def _start_head_node(self, docker_image, mem_size, shm_size, num_cpus,
+                       num_gpus, development_mode):
     """Start the Ray head node inside a docker container."""
     mem_arg = ["--memory=" + mem_size] if mem_size else []
     shm_arg = ["--shm-size=" + shm_size] if shm_size else []
@@ -94,10 +95,15 @@ class DockerRunner(object):
                    "{}:{}".format(os.path.dirname(os.path.realpath(__file__)),
                                   "/ray/test/jenkins_tests")]
                   if development_mode else [])
-    proc = subprocess.Popen(["docker", "run", "-d"] + mem_arg + shm_arg +
-                            volume_arg +
-                            [docker_image, "/ray/scripts/start_ray.sh",
-                             "--head", "--redis-port=6379"],
+
+    command = (["docker", "run", "-d"] + mem_arg + shm_arg + volume_arg +
+               [docker_image, "/ray/scripts/start_ray.sh", "--head",
+                "--redis-port=6379",
+                "--num-cpus={}".format(num_cpus),
+                "--num-gpus={}".format(num_gpus)])
+    print("Starting head node with command:{}".format(command))
+
+    proc = subprocess.Popen(command,
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     stdout_data, _ = wait_for_output(proc)
     container_id = self._get_container_id(stdout_data)
@@ -105,33 +111,34 @@ class DockerRunner(object):
       raise RuntimeError("Failed to find container ID.")
     self.head_container_id = container_id
     self.head_container_ip = self._get_container_ip(container_id)
-    print("start_node", {"container_id": container_id,
-                         "is_head": True,
-                         "shm_size": shm_size,
-                         "ip_address": self.head_container_ip})
-    return container_id
 
-  def _start_worker_node(self, docker_image, mem_size, shm_size):
+  def _start_worker_node(self, docker_image, mem_size, shm_size, num_cpus,
+                         num_gpus, development_mode):
     """Start a Ray worker node inside a docker container."""
     mem_arg = ["--memory=" + mem_size] if mem_size else []
     shm_arg = ["--shm-size=" + shm_size] if shm_size else []
-    proc = subprocess.Popen(["docker", "run", "-d"] + mem_arg + shm_arg +
-                            ["--shm-size=" + shm_size, docker_image,
-                             "/ray/scripts/start_ray.sh",
-                             "--redis-address={:s}:6379".format(
-                                 self.head_container_ip)],
-                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    volume_arg = (["-v",
+                   "{}:{}".format(os.path.dirname(os.path.realpath(__file__)),
+                                  "/ray/test/jenkins_tests")]
+                  if development_mode else [])
+    command = (["docker", "run", "-d"] + mem_arg + shm_arg + volume_arg +
+               ["--shm-size=" + shm_size, docker_image,
+                "/ray/scripts/start_ray.sh",
+                "--redis-address={:s}:6379".format(self.head_container_ip),
+                "--num-cpus={}".format(num_cpus),
+                "--num-gpus={}".format(num_gpus)])
+    print("Starting worker node with command:{}".format(command))
+    proc = subprocess.Popen(command, stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE)
     stdout_data, _ = wait_for_output(proc)
     container_id = self._get_container_id(stdout_data)
     if container_id is None:
       raise RuntimeError("Failed to find container id")
     self.worker_container_ids.append(container_id)
-    print("start_node", {"container_id": container_id,
-                         "is_head": False,
-                         "shm_size": shm_size})
 
-  def start_ray(self, docker_image, mem_size, shm_size, num_nodes,
-                development_mode):
+  def start_ray(self, docker_image=None, mem_size=None, shm_size=None,
+                num_nodes=None, num_cpus=None, num_gpus=None,
+                development_mode=None):
     """Start a Ray cluster within docker.
 
     This starts one docker container running the head node and num_nodes - 1
@@ -146,15 +153,23 @@ class DockerRunner(object):
         with. This will be passed into `docker run` as the `--shm-size` flag.
       num_nodes: The number of nodes to use in the cluster (this counts the
         head node as well).
+      num_cpus: A list of the number of CPUs to start each node with.
+      num_gpus: A list of the number of GPUs to start each node with.
       development_mode: True if you want to mount the local copy of
         test/jenkins_test on the head node so we can avoid rebuilding docker
         images during development.
     """
+    assert len(num_cpus) == num_nodes
+    assert len(num_gpus) == num_nodes
+
     # Launch the head node.
-    self._start_head_node(docker_image, mem_size, shm_size, development_mode)
+    self._start_head_node(docker_image, mem_size, shm_size, num_cpus[0],
+                          num_gpus[0], development_mode)
     # Start the worker nodes.
-    for _ in range(num_nodes - 1):
-      self._start_worker_node(docker_image, mem_size, shm_size)
+    for i in range(num_nodes - 1):
+      self._start_worker_node(docker_image, mem_size, shm_size,
+                              num_cpus[1 + i], num_gpus[1 + i],
+                              development_mode)
 
   def _stop_node(self, container_id):
     """Stop a node in the Ray cluster."""
@@ -181,32 +196,51 @@ class DockerRunner(object):
     for container_id in self.worker_container_ids:
       self._stop_node(container_id)
 
-  def run_test(self, test_script, run_in_docker=False):
+  def run_test(self, test_script, num_drivers, driver_locations=None):
     """Run a test script.
 
     Run a test using the Ray cluster.
 
     Args:
       test_script: The test script to run.
-      run_in_docker: If true then the test script will be run in a docker
-        container. If false, it will be run regularly.
+      num_drivers: The number of copies of the test script to run.
+      driver_locations: A list of the indices of the containers that the
+        different copies of the test script should be run on. If this is None,
+        then the containers will be chosen randomly.
 
     Returns:
       A dictionary with information about the test script run.
     """
-    print("Starting to run test script {}.".format(test_script))
-    proc = subprocess.Popen(["docker", "exec", self.head_container_id,
-                             "/bin/bash", "-c",
-                             "RAY_REDIS_ADDRESS={}:6379 "
-                             "python {}".format(self.head_container_ip,
-                                                test_script)],
-                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    stdout_data, stderr_data = wait_for_output(proc)
-    print("STDOUT:")
-    print(stdout_data)
-    print("STDERR:")
-    print(stderr_data)
-    return {"success": proc.returncode == 0, "return_code": proc.returncode}
+    all_container_ids = [self.head_container_id] + self.worker_container_ids
+    if driver_locations is None:
+      driver_locations = [np.random.randint(0, len(all_container_ids))
+                          for _ in range(num_drivers)]
+
+    # Start the different drivers.
+    driver_processes = []
+    for i in range(len(driver_locations)):
+      # Get the container ID to run the ith driver in.
+      container_id = all_container_ids[driver_locations[i]]
+      command = ["docker", "exec", container_id, "/bin/bash", "-c",
+                 ("RAY_REDIS_ADDRESS={}:6379 RAY_DRIVER_INDEX={} python {}"
+                  .format(self.head_container_ip, i, test_script))]
+      print("Starting driver with command {}.".format(test_script))
+      # Start the driver.
+      p = subprocess.Popen(command, stdout=subprocess.PIPE,
+                           stderr=subprocess.PIPE)
+      driver_processes.append(p)
+
+    # Wait for the drivers to finish.
+    results = []
+    for p in driver_processes:
+      stdout_data, stderr_data = wait_for_output(p)
+      print("STDOUT:")
+      print(stdout_data)
+      print("STDERR:")
+      print(stderr_data)
+      results.append({"success": p.returncode == 0,
+                      "return_code": p.returncode})
+    return results
 
 
 if __name__ == "__main__":
@@ -218,23 +252,49 @@ if __name__ == "__main__":
   parser.add_argument("--shm-size", default="1G", help="shared memory size")
   parser.add_argument("--num-nodes", default=1, type=int,
                       help="number of nodes to use in the cluster")
+  parser.add_argument("--num-cpus", type=str,
+                      help=("a comma separated list of values representing "
+                            "the number of CPUs to start each node with"))
+  parser.add_argument("--num-gpus", type=str,
+                      help=("a comma separated list of values representing "
+                            "the number of GPUs to start each node with"))
+  parser.add_argument("--num-drivers", default=1, type=int,
+                      help="number of drivers to run")
+  parser.add_argument("--driver-locations", type=str,
+                      help=("a comma separated list of indices of the "
+                            "containers to run the drivers in"))
   parser.add_argument("--test-script", required=True, help="test script")
   parser.add_argument("--development-mode", action="store_true",
                       help="use local copies of the test scripts")
   args = parser.parse_args()
 
+  # Parse the number of CPUs and GPUs to use for each worker.
+  num_nodes = args.num_nodes
+  num_cpus = ([int(i) for i in args.num_cpus.split(",")]
+              if args.num_cpus is not None else num_nodes * [10])
+  num_gpus = ([int(i) for i in args.num_gpus.split(",")]
+              if args.num_gpus is not None else num_nodes * [0])
+
   d = DockerRunner()
-  d.start_ray(mem_size=args.mem_size, shm_size=args.shm_size,
-              num_nodes=args.num_nodes, docker_image=args.docker_image,
+  d.start_ray(docker_image=args.docker_image, mem_size=args.mem_size,
+              shm_size=args.shm_size, num_nodes=num_nodes,
+              num_cpus=num_cpus, num_gpus=num_gpus,
               development_mode=args.development_mode)
   try:
-    run_result = d.run_test(args.test_script)
+    run_results = d.run_test(args.test_script, args.num_drivers,
+                             driver_locations=args.driver_locations)
   finally:
     d.stop_ray()
 
-  if "success" in run_result and run_result["success"]:
-    print("RESULT: Test {} succeeded.".format(args.test_script))
-    sys.exit(0)
-  else:
-    print("RESULT: Test {} failed.".format(args.test_script))
+  any_failed = False
+  for run_result in run_results:
+    if "success" in run_result and run_result["success"]:
+      print("RESULT: Test {} succeeded.".format(args.test_script))
+    else:
+      print("RESULT: Test {} failed.".format(args.test_script))
+      any_failed = True
+
+  if any_failed:
     sys.exit(1)
+  else:
+    sys.exit(0)

test/jenkins_tests/multi_node_tests/many_drivers_test.py

@@ -0,0 +1,80 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import time
+
+import ray
+from ray.test.multi_node_tests import (_wait_for_nodes_to_join,
+                                       _broadcast_event,
+                                       _wait_for_event)
+
+# This test should be run with 5 nodes, which have 0, 0, 5, 6, and 50 GPUs for
+# a total of 61 GPUs. It should be run with a large number of drivers (e.g.,
+# 100). At most 10 drivers will run at a time, and each driver will use at most
+# 5 GPUs (this is ceil(61 / 15), which guarantees that we will always be able
+# to make progress).
+total_num_nodes = 5
+max_concurrent_drivers = 15
+num_gpus_per_driver = 5
+
+
+@ray.actor(num_gpus=1)
+class Actor1(object):
+  def __init__(self):
+    assert len(ray.get_gpu_ids()) == 1
+
+  def check_ids(self):
+    assert len(ray.get_gpu_ids()) == 1
+
+
+def driver(redis_address, driver_index):
+  """The script for driver 0.
+
+  This driver should create five actors that each use one GPU and some actors
+  that use no GPUs. After a while, it should exit.
+  """
+  ray.init(redis_address=redis_address)
+
+  # Wait for all the nodes to join the cluster.
+  _wait_for_nodes_to_join(total_num_nodes)
+
+  # Limit the number of drivers running concurrently.
+  for i in range(driver_index - max_concurrent_drivers + 1):
+    _wait_for_event("DRIVER_{}_DONE".format(i), redis_address)
+
+  def try_to_create_actor(actor_class, timeout=100):
+    # Try to create an actor, but allow failures while we wait for the monitor
+    # to release the resources for the removed drivers.
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+      try:
+        actor = actor_class()
+      except Exception as e:
+        time.sleep(0.1)
+      else:
+        return actor
+    # If we are here, then we timed out while looping.
+    raise Exception("Timed out while trying to create actor.")
+
+  # Create some actors that require one GPU.
+  actors_one_gpu = []
+  for _ in range(num_gpus_per_driver):
+    actors_one_gpu.append(try_to_create_actor(Actor1))
+
+  for _ in range(100):
+    ray.get([actor.check_ids() for actor in actors_one_gpu])
+
+  _broadcast_event("DRIVER_{}_DONE".format(driver_index), redis_address)
+
+
+if __name__ == "__main__":
+  driver_index = int(os.environ["RAY_DRIVER_INDEX"])
+  redis_address = os.environ["RAY_REDIS_ADDRESS"]
+  print("Driver {} started at {}.".format(driver_index, time.time()))
+
+  # In this test, all drivers will run the same script.
+  driver(redis_address, driver_index)
+
+  print("Driver {} finished at {}.".format(driver_index, time.time()))

test/jenkins_tests/multi_node_tests/remove_driver_test.py

@@ -0,0 +1,153 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import time
+
+import ray
+from ray.test.multi_node_tests import (_wait_for_nodes_to_join,
+                                       _broadcast_event,
+                                       _wait_for_event)
+
+# This test should be run with 5 nodes, which have 0, 1, 2, 3, and 4 GPUs for a
+# total of 10 GPUs. It shoudl be run with 3 drivers.
+total_num_nodes = 5
+
+
+@ray.actor
+class Actor0(object):
+  def __init__(self):
+    assert len(ray.get_gpu_ids()) == 0
+
+  def check_ids(self):
+    assert len(ray.get_gpu_ids()) == 0
+
+
+@ray.actor(num_gpus=1)
+class Actor1(object):
+  def __init__(self):
+    assert len(ray.get_gpu_ids()) == 1
+
+  def check_ids(self):
+    assert len(ray.get_gpu_ids()) == 1
+
+
+@ray.actor(num_gpus=2)
+class Actor2(object):
+  def __init__(self):
+    assert len(ray.get_gpu_ids()) == 2
+
+  def check_ids(self):
+    assert len(ray.get_gpu_ids()) == 2
+
+
+def driver_0(redis_address):
+  """The script for driver 0.
+
+  This driver should create five actors that each use one GPU and some actors
+  that use no GPUs. After a while, it should exit.
+  """
+  ray.init(redis_address=redis_address)
+
+  # Wait for all the nodes to join the cluster.
+  _wait_for_nodes_to_join(total_num_nodes)
+
+  # Create some actors that require one GPU.
+  actors_one_gpu = [Actor1() for _ in range(5)]
+  # Create some actors that don't require any GPUs.
+  actors_no_gpus = [Actor0() for _ in range(5)]
+
+  for _ in range(1000):
+    ray.get([actor.check_ids() for actor in actors_one_gpu])
+    ray.get([actor.check_ids() for actor in actors_no_gpus])
+
+  _broadcast_event("DRIVER_0_DONE", redis_address)
+
+
+def driver_1(redis_address):
+  """The script for driver 1.
+
+  This driver should create one actor that uses two GPUs, three actors that
+  each use one GPU (the one requiring two must be created first), and some
+  actors that don't use any GPUs. After a while, it should exit.
+  """
+  ray.init(redis_address=redis_address)
+
+  # Wait for all the nodes to join the cluster.
+  _wait_for_nodes_to_join(total_num_nodes)
+
+  # Create an actor that requires two GPUs.
+  actors_two_gpus = [Actor2() for _ in range(1)]
+  # Create some actors that require one GPU.
+  actors_one_gpu = [Actor1() for _ in range(3)]
+  # Create some actors that don't require any GPUs.
+  actors_no_gpus = [Actor0() for _ in range(5)]
+
+  for _ in range(1000):
+    ray.get([actor.check_ids() for actor in actors_two_gpus])
+    ray.get([actor.check_ids() for actor in actors_one_gpu])
+    ray.get([actor.check_ids() for actor in actors_no_gpus])
+
+  _broadcast_event("DRIVER_1_DONE", redis_address)
+
+
+def driver_2(redis_address):
+  """The script for driver 2.
+
+  This driver should wait for the first two drivers to finish. Then it should
+  create some actors that use a total of ten GPUs.
+  """
+  ray.init(redis_address=redis_address)
+
+  _wait_for_event("DRIVER_0_DONE", redis_address)
+  _wait_for_event("DRIVER_1_DONE", redis_address)
+
+  def try_to_create_actor(actor_class, timeout=20):
+    # Try to create an actor, but allow failures while we wait for the monitor
+    # to release the resources for the removed drivers.
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+      try:
+        actor = actor_class()
+      except Exception as e:
+        time.sleep(0.1)
+      else:
+        return actor
+    # If we are here, then we timed out while looping.
+    raise Exception("Timed out while trying to create actor.")
+
+  # Create some actors that require two GPUs.
+  actors_two_gpus = []
+  for _ in range(3):
+    actors_two_gpus.append(try_to_create_actor(Actor2))
+  # Create some actors that require one GPU.
+  actors_one_gpu = []
+  for _ in range(4):
+    actors_one_gpu.append(try_to_create_actor(Actor1))
+  # Create some actors that don't require any GPUs.
+  actors_no_gpus = [Actor0() for _ in range(5)]
+
+  for _ in range(1000):
+    ray.get([actor.check_ids() for actor in actors_two_gpus])
+    ray.get([actor.check_ids() for actor in actors_one_gpu])
+    ray.get([actor.check_ids() for actor in actors_no_gpus])
+
+  _broadcast_event("DRIVER_2_DONE", redis_address)
+
+
+if __name__ == "__main__":
+  driver_index = int(os.environ["RAY_DRIVER_INDEX"])
+  redis_address = os.environ["RAY_REDIS_ADDRESS"]
+  print("Driver {} started at {}.".format(driver_index, time.time()))
+
+  if driver_index == 0:
+    driver_0(redis_address)
+  elif driver_index == 1:
+    driver_1(redis_address)
+  elif driver_index == 2:
+    driver_2(redis_address)
+  else:
+    raise Exception("This code should be unreachable.")
+
+  print("Driver {} finished at {}.".format(driver_index, time.time()))

test/jenkins_tests/multi_node_tests/test_0.py

@@ -15,7 +15,11 @@ def f():
 
 
 if __name__ == "__main__":
-  ray.init(redis_address=os.environ["RAY_REDIS_ADDRESS"])
+  driver_index = int(os.environ["RAY_DRIVER_INDEX"])
+  redis_address = os.environ["RAY_REDIS_ADDRESS"]
+  print("Driver {} started at {}.".format(driver_index, time.time()))
+
+  ray.init(redis_address=redis_address)
   # Check that tasks are scheduled on all nodes.
   num_attempts = 30
   for i in range(num_attempts):
@@ -26,3 +30,5 @@ if __name__ == "__main__":
     if len(counts) == 5:
       break
   assert len(counts) == 5
+
+  print("Driver {} finished at {}.".format(driver_index, time.time()))

test/jenkins_tests/run_multi_node_tests.sh

@@ -13,6 +13,20 @@ python $ROOT_DIR/multi_node_docker_test.py \
     --num-nodes=5 \
     --test-script=/ray/test/jenkins_tests/multi_node_tests/test_0.py
 
+python $ROOT_DIR/multi_node_docker_test.py \
+    --docker-image=$DOCKER_SHA \
+    --num-nodes=5 \
+    --num-gpus=0,1,2,3,4 \
+    --num-drivers=3 \
+    --test-script=/ray/test/jenkins_tests/multi_node_tests/remove_driver_test.py
+
+python $ROOT_DIR/multi_node_docker_test.py \
+    --docker-image=$DOCKER_SHA \
+    --num-nodes=5 \
+    --num-gpus=0,0,5,6,50 \
+    --num-drivers=100 \
+    --test-script=/ray/test/jenkins_tests/multi_node_tests/many_drivers_test.py
+
 python $ROOT_DIR/multi_node_docker_test.py \
     --docker-image=$DOCKER_SHA \
     --num-nodes=1 \

test/runtest.py

@@ -1433,8 +1433,7 @@ class GlobalStateAPI(unittest.TestCase):
 
     client_table = ray.global_state.client_table()
     node_ip_address = ray.worker.global_worker.node_ip_address
-    self.assertEqual(len(client_table[node_ip_address]), 2)
-    self.assertEqual(len(client_table[":"]), 1)
+    self.assertEqual(len(client_table[node_ip_address]), 3)
     manager_client = [c for c in client_table[node_ip_address]
                       if c["ClientType"] == "plasma_manager"][0]