wehub-resource-sync
99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
import functools
|
|
import importlib
|
|
import sys
|
|
import warnings
|
|
|
|
import pytest
|
|
|
|
import ray.train
|
|
import ray.tune
|
|
from ray.train.constants import ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR
|
|
from ray.util.annotations import RayDeprecationWarning
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def enable_v2(monkeypatch):
|
|
monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", "1")
|
|
importlib.reload(ray.train)
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def enable_v2_migration_deprecation_messages(monkeypatch):
|
|
monkeypatch.setenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR, "1")
|
|
yield
|
|
monkeypatch.delenv(ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR)
|
|
|
|
|
|
@pytest.mark.parametrize("v2_enabled", [False, True])
|
|
def test_trainable_fn_utils(tmp_path, monkeypatch, v2_enabled):
|
|
monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", str(int(v2_enabled)))
|
|
importlib.reload(ray.train)
|
|
|
|
dummy_checkpoint_dir = tmp_path.joinpath("dummy")
|
|
dummy_checkpoint_dir.mkdir()
|
|
|
|
asserting_context = (
|
|
functools.partial(pytest.raises, DeprecationWarning)
|
|
if v2_enabled
|
|
else functools.partial(pytest.warns, RayDeprecationWarning)
|
|
)
|
|
|
|
def tune_fn(config):
|
|
with asserting_context(match="get_checkpoint"):
|
|
ray.train.get_checkpoint()
|
|
|
|
with warnings.catch_warnings():
|
|
ray.tune.get_checkpoint()
|
|
|
|
with asserting_context(match="get_context"):
|
|
ray.train.get_context()
|
|
|
|
with warnings.catch_warnings():
|
|
ray.tune.get_context()
|
|
|
|
with asserting_context(match="report"):
|
|
ray.train.report({"a": 1})
|
|
|
|
with warnings.catch_warnings():
|
|
ray.tune.report({"a": 1})
|
|
|
|
with pytest.warns(RayDeprecationWarning, match="update your imports"):
|
|
ray.tune.report(
|
|
{"a": 1},
|
|
checkpoint=ray.train.Checkpoint.from_directory(dummy_checkpoint_dir),
|
|
)
|
|
|
|
with warnings.catch_warnings():
|
|
ray.tune.report(
|
|
{"a": 1},
|
|
checkpoint=ray.tune.Checkpoint.from_directory(dummy_checkpoint_dir),
|
|
)
|
|
|
|
tuner = ray.tune.Tuner(
|
|
tune_fn, run_config=ray.tune.RunConfig(storage_path=tmp_path)
|
|
)
|
|
results = tuner.fit()
|
|
assert not results.errors
|
|
|
|
|
|
def test_configs():
|
|
with pytest.warns(RayDeprecationWarning, match="update your imports"):
|
|
ray.tune.Tuner(lambda c: None, run_config=ray.train.RunConfig())
|
|
|
|
with pytest.warns(RayDeprecationWarning, match="update your imports"):
|
|
ray.tune.Tuner(
|
|
lambda c: None,
|
|
run_config=ray.tune.RunConfig(failure_config=ray.train.FailureConfig()),
|
|
)
|
|
|
|
with warnings.catch_warnings():
|
|
ray.tune.Tuner(
|
|
lambda c: None,
|
|
run_config=ray.tune.RunConfig(failure_config=ray.tune.FailureConfig()),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|