wehub-resource-sync
91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
import pytest
|
|
|
|
import ray
|
|
from ray import train
|
|
from ray.train import ScalingConfig
|
|
from ray.train._internal.worker_group import WorkerGroup
|
|
from ray.train.backend import Backend, BackendConfig
|
|
from ray.train.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
class TestConfig(BackendConfig):
|
|
@property
|
|
def backend_cls(self):
|
|
return TestBackend
|
|
|
|
|
|
class TestBackend(Backend):
|
|
def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
|
|
pass
|
|
|
|
def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
|
|
pass
|
|
|
|
|
|
def test_run(ray_start_4_cpus):
|
|
"""Tests that Train can be run without any specific backends."""
|
|
num_workers = 2
|
|
key = "value"
|
|
value = 1
|
|
config = TestConfig()
|
|
|
|
def train_func():
|
|
checkpoint = train.get_checkpoint()
|
|
checkpoint_dict = load_dict_checkpoint(checkpoint)
|
|
if train.get_context().get_world_rank() == 0:
|
|
train.report(metrics=checkpoint_dict, checkpoint=checkpoint)
|
|
else:
|
|
train.report(metrics=checkpoint_dict)
|
|
return checkpoint_dict[key]
|
|
|
|
with create_dict_checkpoint({key: value}) as checkpoint:
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=config,
|
|
resume_from_checkpoint=checkpoint,
|
|
scaling_config=ScalingConfig(num_workers=num_workers),
|
|
)
|
|
results = trainer.fit()
|
|
|
|
assert load_dict_checkpoint(results.checkpoint) == load_dict_checkpoint(
|
|
checkpoint
|
|
)
|
|
|
|
|
|
def test_failure():
|
|
"""Tests that backend frameworks and non-critical libraries are not imported."""
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import torch # noqa: F401
|
|
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import tensorflow # noqa: F401
|
|
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import horovod # noqa: F401
|
|
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import accelerate # noqa: F401
|
|
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import transformers # noqa: F401
|
|
|
|
with pytest.raises(ModuleNotFoundError):
|
|
import xgboost # noqa: F401
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|