wehub-resource-sync
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
import sys
|
|
from typing import Dict, Optional
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune import Callback, ResumeConfig
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def ray_start_4_cpus_2_gpus_extra():
|
|
address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def register_test_trainable():
|
|
register_mock_trainable()
|
|
|
|
|
|
class StatefulCallback(Callback):
|
|
CKPT_FILE_TMPL = "test-callback-state-{}.json"
|
|
|
|
def __init__(self):
|
|
self.counter = 0
|
|
|
|
def on_trial_result(self, iteration, trials, trial, result, **info):
|
|
self.counter += 1
|
|
|
|
def get_state(self) -> Optional[Dict]:
|
|
return {"counter": self.counter}
|
|
|
|
def set_state(self, state: Dict):
|
|
self.counter = state["counter"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_callback_save_restore(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmpdir
|
|
):
|
|
"""Check that callback state is restored correctly.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCallbackSaveRestore
|
|
"""
|
|
storage = mock_storage_context()
|
|
runner = TuneController(callbacks=[StatefulCallback()], storage=storage)
|
|
runner.add_trial(Trial(MOCK_TRAINABLE_NAME, stub=True, storage=storage))
|
|
for i in range(3):
|
|
runner._callbacks.on_trial_result(
|
|
iteration=i, trials=None, trial=None, result=None
|
|
)
|
|
runner.checkpoint(force=True, wait=True)
|
|
callback = StatefulCallback()
|
|
runner2 = TuneController(callbacks=[callback], storage=storage)
|
|
assert callback.counter == 0
|
|
runner2.resume(resume_config=ResumeConfig())
|
|
assert callback.counter == 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|