wehub-resource-sync
91 lines
2.1 KiB
Python
91 lines
2.1 KiB
Python
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
|