chore: import upstream snapshot with attribution

python/ray/train/tests/test_session.py

@@ -0,0 +1,437 @@
+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__]))