chore: import upstream snapshot with attribution

python/ray/train/v2/tests/conftest.py

@@ -0,0 +1,90 @@
+import logging
+
+import pytest
+
+import ray
+from ray import runtime_context
+from ray._common import utils as ray_utils
+from ray.cluster_utils import Cluster
+from ray.train.v2._internal.constants import (
+    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+)
+
+
+@pytest.fixture()
+def ray_start_4_cpus():
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture()
+def ray_start_4_cpus_2_gpus():
+    ray.init(num_cpus=4, num_gpus=2)
+    yield
+    ray.shutdown()
+
+
+@pytest.fixture
+def ray_start_2x2_gpu_cluster():
+    cluster = Cluster()
+    for _ in range(2):
+        cluster.add_node(num_cpus=4, num_gpus=2)
+
+    ray.init(address=cluster.address)
+
+    yield
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def setup_logging():
+    logger = logging.getLogger("ray.train")
+    orig_level = logger.getEffectiveLevel()
+    logger.setLevel(logging.INFO)
+    yield
+    logger.setLevel(orig_level)
+
+
+@pytest.fixture(autouse=True)
+def reset_ray_address(monkeypatch):
+    ray_utils.reset_ray_address()
+    yield
+    ray_utils.reset_ray_address()
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def disable_state_actor_polling(monkeypatch):
+    monkeypatch.setenv(ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR, "0")
+    yield
+
+
+@pytest.fixture
+def mock_runtime_context(monkeypatch):
+    @ray.remote(num_cpus=0)
+    class DummyActor:
+        pass
+
+    # Must return real actor handle so it can get passed to other actors
+    # Cannot create actor here since ray has not been initialized yet
+    def mock_current_actor(self):
+        return DummyActor.remote()
+
+    # In unit tests where the controller is not an actor, current_actor is
+    # a DummyActor, which is ok because it won't be called in those tests.
+    # In unit tests where the controller is an actor, current_actor is the
+    # controller actor because monkeypatch doesn't propagate to the actor
+    # process. Those tests can successfully test methods on that actor.
+    monkeypatch.setattr(
+        runtime_context.RuntimeContext, "current_actor", property(mock_current_actor)
+    )
+
+    yield