chore: import upstream snapshot with attribution

python/ray/tests/test_multiprocessing_standalone.py

@@ -0,0 +1,182 @@
+"""Tests for ray.util.multiprocessing that require a standalone Ray cluster per test.
+
+Tests that can run on a shared Ray cluster fixture should go in test_multiprocessing.py
+"""
+import math
+import multiprocessing as mp
+import os
+import sys
+
+import pytest
+
+import ray
+from ray._private.test_utils import persistent_gcs_test_enabled
+from ray.util.multiprocessing import Pool
+
+
+@pytest.fixture(scope="module")
+def ray_init_4_cpu_shared():
+    yield ray.init(num_cpus=4)
+
+
+@pytest.fixture
+def pool_4_processes(ray_init_4_cpu_shared):
+    pool = Pool(processes=4)
+    yield pool
+    pool.terminate()
+    pool.join()
+
+
+@pytest.fixture
+def pool_4_processes_python_multiprocessing_lib():
+    pool = mp.Pool(processes=4)
+    yield pool
+    pool.terminate()
+    pool.join()
+
+
+@pytest.mark.skipif(
+    persistent_gcs_test_enabled(),
+    reason="Starts multiple Ray instances in parallel with the same namespace.",
+)
+def test_ray_init(shutdown_only):
+    def getpid(i: int):
+        return os.getpid()
+
+    def check_pool_size(pool, size: int):
+        assert len(set(pool.map(getpid, range(size)))) == size
+
+    # Check that starting a pool starts ray if not initialized.
+    assert not ray.is_initialized()
+    with Pool(processes=4) as pool:
+        assert ray.is_initialized()
+        check_pool_size(pool, 4)
+        assert int(ray.cluster_resources()["CPU"]) == 4
+    pool.join()
+
+    # Check that starting a pool doesn't affect ray if there is a local
+    # ray cluster running.
+    assert ray.is_initialized()
+    assert int(ray.cluster_resources()["CPU"]) == 4
+    with Pool(processes=2) as pool:
+        assert ray.is_initialized()
+        check_pool_size(pool, 2)
+        assert int(ray.cluster_resources()["CPU"]) == 4
+    pool.join()
+
+    # Check that trying to start a pool on an existing ray cluster throws an
+    # error if there aren't enough CPUs for the number of processes.
+    assert ray.is_initialized()
+    assert int(ray.cluster_resources()["CPU"]) == 4
+    with pytest.raises(ValueError):
+        Pool(processes=8)
+        assert int(ray.cluster_resources()["CPU"]) == 4
+
+
+@pytest.mark.skipif(
+    persistent_gcs_test_enabled(),
+    reason="Starts multiple Ray instances in parallel with the same namespace.",
+)
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 1,
+            "num_nodes": 1,
+            "do_init": False,
+        }
+    ],
+    indirect=True,
+)
+def test_connect_to_ray(monkeypatch, 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
+
+    # Use different numbers of CPUs to distinguish between starting a local
+    # ray cluster and connecting to an existing one.
+    ray.init(address=ray_start_cluster.address)
+    existing_cluster_cpus = int(ray.cluster_resources()["CPU"])
+    local_cluster_cpus = existing_cluster_cpus + 1
+    ray.shutdown()
+
+    # Check that starting a pool connects to the running ray cluster by default.
+    assert not ray.is_initialized()
+    with Pool() as pool:
+        assert ray.is_initialized()
+        check_pool_size(pool, existing_cluster_cpus)
+        assert int(ray.cluster_resources()["CPU"]) == existing_cluster_cpus
+    pool.join()
+    ray.shutdown()
+
+    # Check that starting a pool connects to a running ray cluster if
+    # ray_address is set to the cluster address.
+    assert not ray.is_initialized()
+    with Pool(ray_address=ray_start_cluster.address) as pool:
+        check_pool_size(pool, existing_cluster_cpus)
+        assert int(ray.cluster_resources()["CPU"]) == existing_cluster_cpus
+    pool.join()
+    ray.shutdown()
+
+    # Check that starting a pool connects to a running ray cluster if
+    # RAY_ADDRESS is set to the cluster address.
+    assert not ray.is_initialized()
+    monkeypatch.setenv("RAY_ADDRESS", ray_start_cluster.address)
+    with Pool() as pool:
+        check_pool_size(pool, existing_cluster_cpus)
+        assert int(ray.cluster_resources()["CPU"]) == existing_cluster_cpus
+    pool.join()
+    ray.shutdown()
+
+    # Check that trying to start a pool on an existing ray cluster throws an
+    # error if there aren't enough CPUs for the number of processes.
+    assert not ray.is_initialized()
+    with pytest.raises(Exception):
+        Pool(processes=existing_cluster_cpus + 1)
+    assert int(ray.cluster_resources()["CPU"]) == existing_cluster_cpus
+    ray.shutdown()
+
+    # Check that starting a pool starts a local ray cluster if ray_address="local".
+    assert not ray.is_initialized()
+    with Pool(processes=local_cluster_cpus, ray_address="local") as pool:
+        check_pool_size(pool, local_cluster_cpus)
+        assert int(ray.cluster_resources()["CPU"]) == local_cluster_cpus
+    pool.join()
+    ray.shutdown()
+
+    # Check that starting a pool starts a local ray cluster if RAY_ADDRESS="local".
+    assert not ray.is_initialized()
+    monkeypatch.setenv("RAY_ADDRESS", "local")
+    with Pool(processes=local_cluster_cpus) as pool:
+        check_pool_size(pool, local_cluster_cpus)
+        assert int(ray.cluster_resources()["CPU"]) == local_cluster_cpus
+    pool.join()
+    ray.shutdown()
+
+
+def test_maxtasksperchild(shutdown_only):
+    with Pool(processes=5, maxtasksperchild=1) as pool:
+        assert len(set(pool.map(lambda _: os.getpid(), range(20)))) == 20
+    pool.join()
+
+
+def test_deadlock_avoidance_in_recursive_tasks(shutdown_only):
+    ray.init(num_cpus=1)
+
+    def poolit_a(_):
+        with Pool() as pool:
+            return list(pool.map(math.sqrt, range(0, 2, 1)))
+
+    def poolit_b():
+        with Pool() as pool:
+            return list(pool.map(poolit_a, range(2, 4, 1)))
+
+    result = poolit_b()
+    assert result == [[0.0, 1.0], [0.0, 1.0]]
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))