chore: import upstream snapshot with attribution

python/ray/util/collective/tests/conftest.py

@@ -0,0 +1,106 @@
+"""Some fixtures for collective tests."""
+import logging
+
+import pytest
+
+import ray
+
+try:
+    from ray.util.collective.collective_group.nccl_collective_group import (
+        _get_comm_key_from_devices,
+        _get_comm_key_send_recv,
+    )
+except Exception:  # Cupy/NCCL may be unavailable on CPU-only setups
+    _get_comm_key_from_devices = None
+    _get_comm_key_send_recv = None
+from ray.util.collective.const import get_store_name
+
+logger = logging.getLogger(__name__)
+logger.setLevel("INFO")
+
+
+# TODO (Hao): remove this clean_up function as it sometimes crashes Ray.
+def clean_up():
+    # If NCCL helpers are unavailable (e.g., no cupy), skip cleanup.
+    if _get_comm_key_from_devices is None or _get_comm_key_send_recv is None:
+        return
+    group_names = ["default", "test", "123?34!", "default2", "random"]
+    group_names.extend([str(i) for i in range(10)])
+    max_world_size = 4
+    all_keys = []
+    for name in group_names:
+        devices = [[0], [0, 1], [1, 0]]
+        for d in devices:
+            collective_communicator_key = _get_comm_key_from_devices(d)
+            all_keys.append(collective_communicator_key + "@" + name)
+        for i in range(max_world_size):
+            for j in range(max_world_size):
+                if i < j:
+                    p2p_communicator_key = _get_comm_key_send_recv(i, 0, j, 0)
+                    all_keys.append(p2p_communicator_key + "@" + name)
+    for group_key in all_keys:
+        store_name = get_store_name(group_key)
+        try:
+            actor = ray.get_actor(store_name)
+        except ValueError:
+            actor = None
+        if actor:
+            logger.debug(
+                "Killing actor with group_key: '{}' and store: '{}'.".format(
+                    group_key, store_name
+                )
+            )
+            ray.kill(actor)
+
+
+@pytest.fixture
+def ray_start_single_node_2_gpus():
+    # Please start this fixture in a cluster with 2 GPUs.
+    address_info = ray.init(num_gpus=2)
+    yield address_info
+    ray.shutdown()
+
+
+# Hao: this fixture is a bit tricky.
+# I use a bash script to start a ray cluster on
+# my own on-premise cluster before run this fixture.
+@pytest.fixture
+def ray_start_distributed_2_nodes_4_gpus():
+    # The cluster has a setup of 2 nodes, each node with 2
+    # GPUs. Each actor will be allocated 1 GPU.
+    ray.init("auto")
+    yield
+    clean_up()
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_distributed_multigpu_2_nodes_4_gpus():
+    # The cluster has a setup of 2 nodes, each node with 2
+    # GPUs. Each actor will be allocated 2 GPUs.
+    ray.init("auto")
+    yield
+    clean_up()
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_single_node():
+    address_info = ray.init(num_cpus=8)
+    yield address_info
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_distributed_2_nodes():
+    # The cluster has a setup of 2 nodes.
+    # no GPUs!
+    ray.init("auto")
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()