wehub-resource-sync
438 lines
11 KiB
Python
438 lines
11 KiB
Python
import tempfile
|
|
import time
|
|
import warnings
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.air._internal.util import StartTraceback
|
|
from ray.air.constants import SESSION_MISUSE_LOG_ONCE_KEY
|
|
from ray.air.session import (
|
|
get_checkpoint,
|
|
get_dataset_shard,
|
|
get_local_rank,
|
|
get_world_rank,
|
|
get_world_size,
|
|
report,
|
|
)
|
|
from ray.train._internal.accelerator import Accelerator
|
|
from ray.train._internal.session import (
|
|
_TrainingResult,
|
|
get_accelerator,
|
|
get_session,
|
|
init_session,
|
|
set_accelerator,
|
|
shutdown_session,
|
|
)
|
|
from ray.train._internal.storage import StorageContext
|
|
from ray.train.error import SessionMisuseError
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
|
|
storage = StorageContext(
|
|
storage_path=tempfile.mkdtemp(),
|
|
experiment_dir_name="exp_name",
|
|
trial_dir_name="trial_name",
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def ray_start_4_cpus():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def session():
|
|
def f():
|
|
return 1
|
|
|
|
init_session(
|
|
training_func=f,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
yield get_session()
|
|
shutdown_session()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def shutdown():
|
|
if get_session():
|
|
shutdown_session()
|
|
|
|
|
|
def test_init_fail(session):
|
|
with pytest.raises(ValueError):
|
|
init_session(lambda: 1, 0)
|
|
|
|
|
|
def test_shutdown(session):
|
|
shutdown_session()
|
|
assert not get_session()
|
|
|
|
|
|
def test_world_rank(session):
|
|
assert get_world_rank() == 0
|
|
shutdown_session()
|
|
# Make sure default to 0.
|
|
assert get_world_rank() == 0
|
|
|
|
|
|
def test_local_rank(session):
|
|
assert get_local_rank() == 0
|
|
shutdown_session()
|
|
# Make sure default to 0.
|
|
assert get_local_rank() == 0
|
|
|
|
|
|
def test_world_size(session):
|
|
assert get_world_size() == 1
|
|
shutdown_session()
|
|
# Make sure default to 1.
|
|
assert get_world_size() == 1
|
|
|
|
|
|
def test_train(session):
|
|
session.start()
|
|
session.finish()
|
|
|
|
|
|
def test_get_dataset_shard():
|
|
dataset = ray.data.from_items([1, 2, 3])
|
|
init_session(
|
|
training_func=lambda: 1,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
dataset_shard=dataset,
|
|
storage=storage,
|
|
)
|
|
assert get_dataset_shard() == dataset
|
|
shutdown_session()
|
|
|
|
|
|
def test_report():
|
|
def train_func():
|
|
for i in range(2):
|
|
report(dict(loss=i))
|
|
|
|
init_session(
|
|
training_func=train_func,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
assert session.get_next().metrics["loss"] == 0
|
|
assert session.get_next().metrics["loss"] == 1
|
|
shutdown_session()
|
|
|
|
|
|
def test_report_fail():
|
|
def train_func():
|
|
for i in range(2):
|
|
report(i)
|
|
return 1
|
|
|
|
init_session(
|
|
training_func=train_func,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
with pytest.raises(StartTraceback):
|
|
session.get_next()
|
|
shutdown_session()
|
|
|
|
|
|
def test_report_after_finish(session):
|
|
session.start()
|
|
session.pause_reporting()
|
|
session.finish()
|
|
for _ in range(2):
|
|
report(dict(loss=1))
|
|
assert session.get_next() is None
|
|
shutdown_session()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"block,put_result_queue,put_actor_queue",
|
|
[
|
|
(False, False, False),
|
|
(False, False, True),
|
|
(False, True, False),
|
|
(True, False, False),
|
|
(True, False, True),
|
|
(True, True, False),
|
|
],
|
|
)
|
|
def test_get_result_from_queues(session, block, put_result_queue, put_actor_queue):
|
|
"""Verify that we get the expected _TrainingResult from each result queue.
|
|
|
|
`block` describes whether we wait for a result or return after a timeout.
|
|
This argument should have no impact on this unit test.
|
|
`put_result_queue` and `put_actor_queue` are mutually exclusive and describe
|
|
which queue has results to process. The returned _TrainingResult should be
|
|
from the expected queue.
|
|
"""
|
|
result_queue_training_result = _TrainingResult(
|
|
checkpoint=None,
|
|
metrics={"result_queue_metric_key": "result_queue_metric_value"},
|
|
)
|
|
if put_result_queue:
|
|
session.result_queue.put(result_queue_training_result, block=True)
|
|
inter_actor_result = {"inter_actor_metric_key": "inter_actor_metric_value"}
|
|
if put_actor_queue:
|
|
session._get_or_create_inter_actor_queue().put(inter_actor_result, block=True)
|
|
result = session._get_result_from_queues(block=block)
|
|
if put_result_queue:
|
|
assert result == result_queue_training_result
|
|
elif put_actor_queue:
|
|
assert (
|
|
result.metrics["inter_actor_metric_key"]
|
|
== inter_actor_result["inter_actor_metric_key"]
|
|
)
|
|
else:
|
|
assert result is None
|
|
|
|
|
|
def test_no_start(session):
|
|
with pytest.raises(RuntimeError):
|
|
session.get_next()
|
|
shutdown_session()
|
|
|
|
|
|
def test_checkpoint():
|
|
def train_func():
|
|
for i in range(2):
|
|
with create_dict_checkpoint(dict(epoch=i)) as checkpoint:
|
|
report({}, checkpoint=checkpoint)
|
|
|
|
def validate_zero(expected):
|
|
next = session.get_next()
|
|
assert next is not None and next.checkpoint is not None
|
|
assert load_dict_checkpoint(next.checkpoint)["epoch"] == expected
|
|
|
|
init_session(
|
|
training_func=train_func,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
validate_zero(0)
|
|
validate_zero(1)
|
|
session.finish()
|
|
shutdown_session()
|
|
|
|
|
|
def test_load_checkpoint_after_save():
|
|
def train_func():
|
|
for i in range(2):
|
|
with create_dict_checkpoint(dict(epoch=i)) as checkpoint:
|
|
report(dict(epoch=i), checkpoint=checkpoint)
|
|
checkpoint = get_checkpoint()
|
|
assert load_dict_checkpoint(checkpoint)["epoch"] == i
|
|
|
|
init_session(
|
|
training_func=train_func,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
for i in range(2):
|
|
session.get_next()
|
|
session.finish()
|
|
shutdown_session()
|
|
|
|
|
|
def test_locking():
|
|
"""Tests that report pauses training until fetch_next or finish."""
|
|
|
|
def train_1():
|
|
import _thread
|
|
|
|
_thread.interrupt_main()
|
|
|
|
init_session(
|
|
training_func=train_1,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
with pytest.raises(KeyboardInterrupt):
|
|
session.start()
|
|
shutdown_session()
|
|
|
|
def train_2():
|
|
for i in range(2):
|
|
report(dict(loss=i))
|
|
train_1()
|
|
|
|
init_session(
|
|
training_func=train_2,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
time.sleep(3)
|
|
|
|
session.pause_reporting()
|
|
# Releases session.continue_lock to resume the training thread.
|
|
session.get_next()
|
|
|
|
with pytest.raises(KeyboardInterrupt):
|
|
session.get_next()
|
|
session.finish()
|
|
shutdown_session()
|
|
|
|
|
|
def reset_log_once_with_str(str_to_append=None):
|
|
key = SESSION_MISUSE_LOG_ONCE_KEY
|
|
if str_to_append:
|
|
key += f"-{str_to_append}"
|
|
ray.util.debug.reset_log_once(key)
|
|
|
|
|
|
@pytest.mark.parametrize("fn", [get_checkpoint, get_dataset_shard])
|
|
def test_warn(fn):
|
|
"""Checks if calling session functions outside of session raises warning."""
|
|
|
|
with warnings.catch_warnings(record=True) as record:
|
|
warnings.simplefilter("always")
|
|
# Ignore Deprecation warnings.
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
assert not fn()
|
|
|
|
assert fn.__name__ in record[0].message.args[0]
|
|
|
|
reset_log_once_with_str(fn.__name__)
|
|
|
|
|
|
def test_warn_report():
|
|
"""Checks if calling session.report function outside of session raises warning."""
|
|
|
|
fn = report
|
|
|
|
with warnings.catch_warnings(record=True) as record:
|
|
warnings.simplefilter("always")
|
|
# Ignore Deprecation warnings.
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
assert not fn(dict())
|
|
|
|
assert fn.__name__ in record[0].message.args[0]
|
|
|
|
reset_log_once_with_str(fn.__name__)
|
|
|
|
|
|
def test_warn_once():
|
|
"""Checks if session misuse warning is only shown once per function."""
|
|
|
|
with warnings.catch_warnings(record=True) as record:
|
|
# Ignore Deprecation warnings.
|
|
warnings.simplefilter("always")
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
assert not get_checkpoint()
|
|
assert not get_checkpoint()
|
|
assert not report(dict(x=2))
|
|
assert not report(dict(x=2))
|
|
assert not get_dataset_shard()
|
|
assert not get_dataset_shard()
|
|
|
|
# Should only warn once.
|
|
assert len(record) == 3
|
|
|
|
|
|
class FakeAccelerator(Accelerator):
|
|
pass
|
|
|
|
|
|
def test_set_and_get_accelerator(session):
|
|
accelerator = FakeAccelerator()
|
|
set_accelerator(accelerator)
|
|
assert get_accelerator(FakeAccelerator) is accelerator
|
|
|
|
|
|
def test_get_accelerator_constructs_default_accelerator(session):
|
|
assert isinstance(get_accelerator(FakeAccelerator), FakeAccelerator)
|
|
|
|
|
|
def test_get_accelerator_raises_error_outside_session():
|
|
with pytest.raises(SessionMisuseError):
|
|
get_accelerator(FakeAccelerator)
|
|
|
|
|
|
def test_set_accelerator_raises_error_if_accelerator_already_set(session):
|
|
accelerator1, accelerator2 = FakeAccelerator(), FakeAccelerator()
|
|
set_accelerator(accelerator1)
|
|
with pytest.raises(RuntimeError):
|
|
set_accelerator(accelerator2)
|
|
|
|
|
|
def test_set_accelerator_raises_error_outside_session():
|
|
accelerator = FakeAccelerator()
|
|
with pytest.raises(SessionMisuseError):
|
|
set_accelerator(accelerator)
|
|
|
|
|
|
def test_application_error_raised():
|
|
def f():
|
|
raise ValueError
|
|
|
|
init_session(
|
|
training_func=f,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
local_world_size=1,
|
|
world_size=1,
|
|
storage=storage,
|
|
)
|
|
session = get_session()
|
|
session.start()
|
|
with pytest.raises(StartTraceback):
|
|
session.get_next()
|
|
shutdown_session()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|