Remove num_local_schedulers argument from ray.worker._init. (#3704)

* Remove num_local_schedulers argument from ray.worker._init.

* Fix

* Fix tests.

test/component_failures_test.py

@@ -12,22 +12,10 @@ import numpy as np
 import pytest
 
 import ray
-from ray.parameter import RayParams
 from ray.test.cluster_utils import Cluster
 from ray.test.test_utils import run_string_as_driver_nonblocking
 
 
-@pytest.fixture
-def ray_start_workers_separate():
-    # Start the Ray processes.
-    ray_params = RayParams(
-        num_cpus=1, start_ray_local=True, redirect_output=True)
-    ray.worker._init(ray_params)
-    yield None
-    # The code after the yield will run as teardown code.
-    ray.shutdown()
-
-
 @pytest.fixture
 def shutdown_only():
     yield None
@@ -38,21 +26,22 @@ def shutdown_only():
 @pytest.fixture
 def ray_start_cluster():
     node_args = {
-        "resources": dict(CPU=8),
+        "num_cpus": 8,
         "_internal_config": json.dumps({
             "initial_reconstruction_timeout_milliseconds": 1000,
             "num_heartbeats_timeout": 10
         })
     }
     # Start with 4 worker nodes and 8 cores each.
-    g = Cluster(initialize_head=True, connect=True, head_node_args=node_args)
+    cluster = Cluster(
+        initialize_head=True, connect=True, head_node_args=node_args)
     workers = []
     for _ in range(4):
-        workers.append(g.add_node(**node_args))
-    g.wait_for_nodes()
-    yield g
+        workers.append(cluster.add_node(**node_args))
+    cluster.wait_for_nodes()
+    yield cluster
     ray.shutdown()
-    g.shutdown()
+    cluster.shutdown()
 
 
 # This test checks that when a worker dies in the middle of a get, the plasma
@@ -235,23 +224,22 @@ ray.wait([ray.ObjectID(ray.utils.hex_to_binary("{}"))])
 
 @pytest.fixture(params=[(1, 4), (4, 4)])
 def ray_start_workers_separate_multinode(request):
-    num_local_schedulers = request.param[0]
+    num_nodes = request.param[0]
     num_initial_workers = request.param[1]
     # Start the Ray processes.
-    ray_params = RayParams(
-        num_local_schedulers=num_local_schedulers,
-        start_ray_local=True,
-        num_cpus=[num_initial_workers] * num_local_schedulers,
-        redirect_output=True)
-    ray.worker._init(ray_params)
-    yield num_local_schedulers, num_initial_workers
+    cluster = Cluster()
+    for _ in range(num_nodes):
+        cluster.add_node(num_cpus=num_initial_workers)
+    ray.init(redis_address=cluster.redis_address)
+
+    yield num_nodes, num_initial_workers
     # The code after the yield will run as teardown code.
     ray.shutdown()
+    cluster.shutdown()
 
 
 def test_worker_failed(ray_start_workers_separate_multinode):
-    num_local_schedulers, num_initial_workers = (
-        ray_start_workers_separate_multinode)
+    num_nodes, num_initial_workers = (ray_start_workers_separate_multinode)
 
     @ray.remote
     def f(x):
@@ -260,9 +248,7 @@ def test_worker_failed(ray_start_workers_separate_multinode):
 
     # Submit more tasks than there are workers so that all workers and
     # cores are utilized.
-    object_ids = [
-        f.remote(i) for i in range(num_initial_workers * num_local_schedulers)
-    ]
+    object_ids = [f.remote(i) for i in range(num_initial_workers * num_nodes)]
     object_ids += [f.remote(object_id) for object_id in object_ids]
     # Allow the tasks some time to begin executing.
     time.sleep(0.1)
@@ -276,22 +262,30 @@ def test_worker_failed(ray_start_workers_separate_multinode):
     ray.get(object_ids)
 
 
+@pytest.fixture
+def ray_initialize_cluster():
+    # Start with 4 workers and 4 cores.
+    num_nodes = 4
+    num_workers_per_scheduler = 8
+
+    cluster = Cluster()
+    for _ in range(num_nodes):
+        cluster.add_node(
+            num_cpus=num_workers_per_scheduler,
+            _internal_config=json.dumps({
+                "initial_reconstruction_timeout_milliseconds": 1000,
+                "num_heartbeats_timeout": 10,
+            }))
+    ray.init(redis_address=cluster.redis_address)
+
+    yield None
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
 def _test_component_failed(component_type):
     """Kill a component on all worker nodes and check workload succeeds."""
-    # Start with 4 workers and 4 cores.
-    num_local_schedulers = 4
-    num_workers_per_scheduler = 8
-    ray_params = RayParams(
-        num_local_schedulers=num_local_schedulers,
-        start_ray_local=True,
-        num_cpus=[num_workers_per_scheduler] * num_local_schedulers,
-        redirect_output=True,
-        _internal_config=json.dumps({
-            "initial_reconstruction_timeout_milliseconds": 1000,
-            "num_heartbeats_timeout": 10,
-        }))
-    ray.worker._init(ray_params)
-
     # Submit many tasks with many dependencies.
     @ray.remote
     def f(x):
@@ -346,20 +340,18 @@ def check_components_alive(component_type, check_component_alive):
             assert not component.poll() is None
 
 
-def test_raylet_failed():
+def test_raylet_failed(ray_initialize_cluster):
     # Kill all local schedulers on worker nodes.
     _test_component_failed(ray.services.PROCESS_TYPE_RAYLET)
 
     # The plasma stores should still be alive on the worker nodes.
     check_components_alive(ray.services.PROCESS_TYPE_PLASMA_STORE, True)
 
-    ray.shutdown()
-
 
 @pytest.mark.skipif(
     os.environ.get("RAY_USE_NEW_GCS") == "on",
     reason="Hanging with new GCS API.")
-def test_plasma_store_failed():
+def test_plasma_store_failed(ray_initialize_cluster):
     # Kill all plasma stores on worker nodes.
     _test_component_failed(ray.services.PROCESS_TYPE_PLASMA_STORE)
 
@@ -367,8 +359,6 @@ def test_plasma_store_failed():
     check_components_alive(ray.services.PROCESS_TYPE_PLASMA_STORE, False)
     check_components_alive(ray.services.PROCESS_TYPE_RAYLET, False)
 
-    ray.shutdown()
-
 
 def test_actor_creation_node_failure(ray_start_cluster):
     # TODO(swang): Refactor test_raylet_failed, etc to reuse the below code.