Pipe num_cpus and num_gpus through from start_ray.py. (#275)

* Pipe num_cpus and num_gpus through from start_ray.py.

* Improve load balancing tests.

* Fix bug.

* Factor out some testing code.

test/multi_node_test.py

@@ -127,11 +127,23 @@ class StartRayScriptTest(unittest.TestCase):
                              "--object-manager-port", "12345"])
     subprocess.Popen([stop_ray_script]).wait()
 
+    # Test starting Ray with the number of CPUs specified.
+    subprocess.check_output([start_ray_script, "--head",
+                             "--num-cpus", "100"])
+    subprocess.Popen([stop_ray_script]).wait()
+
+    # Test starting Ray with the number of GPUs specified.
+    subprocess.check_output([start_ray_script, "--head",
+                             "--num-gpus", "100"])
+    subprocess.Popen([stop_ray_script]).wait()
+
     # Test starting Ray with all arguments specified.
     subprocess.check_output([start_ray_script, "--head",
                              "--num-workers", "20",
                              "--redis-port", "6379",
-                             "--object-manager-port", "12345"])
+                             "--object-manager-port", "12345",
+                             "--num-cpus", "100",
+                             "--num-gpus", "0"])
     subprocess.Popen([stop_ray_script]).wait()
 
     # Test starting Ray with invalid arguments.

test/runtest.py

@@ -1061,7 +1061,23 @@ class ResourcesTest(unittest.TestCase):
 
 class SchedulingAlgorithm(unittest.TestCase):
 
+  def attempt_to_load_balance(self, remote_function, args, total_tasks,
+                              num_local_schedulers, minimum_count,
+                              num_attempts=20):
+    attempts = 0
+    while attempts < num_attempts:
+      locations = ray.get([remote_function.remote(*args) for _ in range(total_tasks)])
+      names = set(locations)
+      counts = [locations.count(name) for name in names]
+      print("Counts are {}.".format(counts))
+      if len(names) == num_local_schedulers and all([count >= minimum_count for count in counts]):
+        break
+      attempts += 1
+    self.assertLess(attempts, num_attempts)
+
   def testLoadBalancing(self):
+    # This test ensures that tasks are being assigned to all local schedulers in
+    # a roughly equal manner.
     num_workers = 21
     num_local_schedulers = 3
     ray.worker._init(start_ray_local=True, num_workers=num_workers, num_local_schedulers=num_local_schedulers)
@@ -1071,23 +1087,14 @@ class SchedulingAlgorithm(unittest.TestCase):
       time.sleep(0.001)
       return ray.worker.global_worker.plasma_client.store_socket_name
 
-    locations = ray.get([f.remote() for _ in range(100)])
-    names = set(locations)
-    self.assertEqual(len(names), num_local_schedulers)
-    counts = [locations.count(name) for name in names]
-    for count in counts:
-      self.assertGreater(count, 30)
-
-    locations = ray.get([f.remote() for _ in range(1000)])
-    names = set(locations)
-    self.assertEqual(len(names), num_local_schedulers)
-    counts = [locations.count(name) for name in names]
-    for count in counts:
-      self.assertGreater(count, 200)
+    self.attempt_to_load_balance(f, [], 100, num_local_schedulers, 25)
+    self.attempt_to_load_balance(f, [], 1000, num_local_schedulers, 250)
 
     ray.worker.cleanup()
 
   def testLoadBalancingWithDependencies(self):
+    # This test ensures that tasks are being assigned to all local schedulers in
+    # a roughly equal manner even when the tasks have dependencies.
     num_workers = 3
     num_local_schedulers = 3
     ray.worker._init(start_ray_local=True, num_workers=num_workers, num_local_schedulers=num_local_schedulers)
@@ -1100,12 +1107,7 @@ class SchedulingAlgorithm(unittest.TestCase):
     # doesn't prevent tasks from being scheduled on other local schedulers.
     x = ray.put(np.zeros(1000000))
 
-    locations = ray.get([f.remote(x) for _ in range(100)])
-    names = set(locations)
-    self.assertEqual(len(names), num_local_schedulers)
-    counts = [locations.count(name) for name in names]
-    for count in counts:
-      self.assertGreater(count, 30)
+    self.attempt_to_load_balance(f, [x], 100, num_local_schedulers, 25)
 
     ray.worker.cleanup()