Move ray.experimental.multiprocessing to ray.util.multiprocessing (#7149)

python/ray/tests/test_multiprocessing.py

@@ -1,21 +1,18 @@
 import os
+import sys
 import pytest
 import tempfile
 import time
 import random
-import subprocess
 from collections import defaultdict
 import queue
 
 import ray
-from ray.experimental.multiprocessing import Pool, TimeoutError
+from ray.util.multiprocessing import Pool, TimeoutError
 
 
-@pytest.fixture
-def cleanup_only():
-    yield None
-    ray.shutdown()
-    subprocess.check_output(["ray", "stop"])
+def teardown_function(function):
+    # Delete environment variable if set.
     if "RAY_ADDRESS" in os.environ:
         del os.environ["RAY_ADDRESS"]
 
@@ -36,7 +33,7 @@ def pool_4_processes():
     ray.shutdown()
 
 
-def test_initialize_ray(cleanup_only):
+def test_ray_init(shutdown_only):
     def getpid(args):
         return os.getpid()
 
@@ -69,14 +66,27 @@ def test_initialize_ray(cleanup_only):
     assert int(ray.state.cluster_resources()["CPU"]) == 1
     ray.shutdown()
 
+
+@pytest.mark.parametrize(
+    "ray_start_cluster", [{
+        "num_cpus": 1,
+        "num_nodes": 1,
+        "do_init": False,
+    }],
+    indirect=True)
+def test_connect_to_ray(ray_start_cluster):
+    def getpid(args):
+        return os.getpid()
+
+    def check_pool_size(pool, size):
+        args = [tuple() for _ in range(size)]
+        assert len(set(pool.map(getpid, args))) == size
+
+    address = ray_start_cluster.address
     # Use different numbers of CPUs to distinguish between starting a local
     # ray cluster and connecting to an existing one.
+    start_cpus = 1  # Set in fixture.
     init_cpus = 2
-    start_cpus = 3
-
-    # Start a ray cluster in the background.
-    subprocess.check_output(
-        ["ray", "start", "--head", "--num-cpus={}".format(start_cpus)])
 
     # Check that starting a pool still starts ray if RAY_ADDRESS not set.
     pool = Pool(processes=init_cpus)
@@ -87,14 +97,14 @@ def test_initialize_ray(cleanup_only):
 
     # Check that starting a pool connects to a running ray cluster if
     # ray_address is passed in.
-    pool = Pool(ray_address="auto")
+    pool = Pool(ray_address=address)
     assert ray.is_initialized()
     assert int(ray.state.cluster_resources()["CPU"]) == start_cpus
     check_pool_size(pool, start_cpus)
     ray.shutdown()
 
     # Set RAY_ADDRESS, so pools should connect to the running ray cluster.
-    os.environ["RAY_ADDRESS"] = "auto"
+    os.environ["RAY_ADDRESS"] = address
 
     # Check that starting a pool connects to a running ray cluster if
     # RAY_ADDRESS is set.
@@ -111,11 +121,8 @@ def test_initialize_ray(cleanup_only):
     assert int(ray.state.cluster_resources()["CPU"]) == start_cpus
     ray.shutdown()
 
-    # Clean up the background ray cluster.
-    subprocess.check_output(["ray", "stop"])
 
-
-def test_initializer(cleanup_only):
+def test_initializer(shutdown_only):
     def init(dirname):
         with open(os.path.join(dirname, str(os.getpid())), "w") as f:
             print("hello", file=f)
@@ -144,6 +151,7 @@ def test_close(pool_4_processes):
     pool_4_processes.join()
 
     # close() shouldn't interrupt pending tasks, so check that they succeeded.
+    result.wait(timeout=10)
     assert result.ready()
     assert result.successful()
     assert result.get() == ["hello"] * 4
@@ -164,6 +172,8 @@ def test_terminate(pool_4_processes):
     result.wait(timeout=10)
     assert result.ready()
     assert not result.successful()
+    with pytest.raises(ray.exceptions.RayError):
+        result.get()
 
 
 def test_apply(pool):
@@ -478,9 +488,14 @@ def test_imap_timeout(pool_4_processes):
         result_iter.next()
 
 
-def test_maxtasksperchild(cleanup_only):
+def test_maxtasksperchild(shutdown_only):
     def f(args):
         return os.getpid()
 
     pool = Pool(5, maxtasksperchild=1)
     assert len(set(pool.map(f, range(20)))) == 20
+
+
+if __name__ == "__main__":
+    import pytest
+    sys.exit(pytest.main(["-v", __file__]))