chore: import upstream snapshot with attribution

python/ray/train/tests/test_training_iterator.py

@@ -0,0 +1,395 @@
+import functools
+import sys
+import time
+from contextlib import nullcontext
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray import train
+from ray.air._internal.util import StartTraceback
+from ray.train import DataConfig
+from ray.train._internal.backend_executor import BackendExecutor
+from ray.train._internal.session import get_session, init_session
+from ray.train._internal.utils import construct_train_func
+from ray.train._internal.worker_group import WorkerGroup
+from ray.train.backend import BackendConfig
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.tests.util import mock_storage_context
+from ray.train.trainer import TrainingIterator
+
+MAX_RETRIES = 3
+
+
+@pytest.fixture(autouse=True, scope="module")
+def patch_tune_session():
+    if not get_session():
+        init_session(
+            training_func=None,
+            world_rank=None,
+            local_rank=None,
+            node_rank=None,
+            local_world_size=None,
+            world_size=None,
+            storage=mock_storage_context(),
+        )
+    yield
+
+
+@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()
+
+
+def gen_execute_single_async_special(special_f):
+    def execute_single_async_special(self, i, f, *args, **kwargs):
+        assert len(self.workers) == 2
+        if i == 0 and hasattr(self, "should_fail") and self.should_fail:
+            kwargs["train_func"] = special_f
+        return (
+            self.workers[i]
+            .actor._RayTrainWorker__execute.options(name=f.__name__)
+            .remote(f, *args, **kwargs)
+        )
+
+    return execute_single_async_special
+
+
+def gen_new_backend_executor(special_f):
+    """Returns a BackendExecutor that runs special_f on worker 0 once."""
+
+    class TestBackendExecutor(BackendExecutor):
+        _has_failed = False
+
+        def start_training(self, *args, **kwargs):
+            special_execute = gen_execute_single_async_special(special_f)
+            if not self._has_failed:
+                self.worker_group.should_fail = True
+                self._has_failed = True
+            else:
+                self.worker_group.should_fail = False
+            with patch.object(WorkerGroup, "execute_single_async", special_execute):
+                super().start_training(*args, **kwargs)
+
+    return TestBackendExecutor
+
+
+def create_iterator(
+    train_func,
+    backend_config,
+    *,
+    num_workers=2,
+    backend_executor_cls=BackendExecutor,
+    init_hook=None,
+):
+    # Similar logic to the old Trainer.run_iterator().
+
+    train_func = construct_train_func(train_func, None, train_func_context=nullcontext)
+
+    backend_executor = backend_executor_cls(
+        backend_config=backend_config, num_workers=num_workers, max_retries=MAX_RETRIES
+    )
+    backend_executor.start(init_hook)
+
+    return TrainingIterator(
+        backend_executor=backend_executor,
+        backend_config=backend_config,
+        train_func=train_func,
+        datasets={},
+        metadata={},
+        data_config=DataConfig(),
+        checkpoint=None,
+    )
+
+
+def test_run_iterator(ray_start_4_cpus):
+    config = BackendConfig()
+
+    def train_func():
+        for i in range(3):
+            train.report(dict(index=i))
+        return 1
+
+    iterator = create_iterator(train_func, config)
+
+    count = 0
+    for results in iterator:
+        assert all(value.metrics["index"] == count for value in results)
+        count += 1
+
+    assert count == 3
+    assert iterator.is_finished()
+
+    with pytest.raises(StopIteration):
+        next(iterator)
+
+
+def test_run_iterator_error(ray_start_4_cpus):
+    config = BackendConfig()
+
+    def fail_train():
+        raise NotImplementedError
+
+    iterator = create_iterator(fail_train, config)
+
+    with pytest.raises(StartTraceback) as exc:
+        next(iterator)
+
+    assert isinstance(exc.value.__cause__, NotImplementedError), (
+        exc.value,
+        exc.value.__cause__,
+    )
+
+    assert iterator.is_finished()
+
+
+def test_worker_failure_1(ray_start_4_cpus):
+    def train_func():
+        train.report({"test": 1})
+
+    def train_actor_failure():
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert all(result.metrics["test"] == 1 for result in worker_results)
+
+
+def test_worker_failure_2(ray_start_4_cpus):
+    def train_func():
+        for _ in range(2):
+            train.report(dict(loss=1))
+
+    def train_actor_failure():
+        for _ in range(2):
+            train.report(dict(loss=1))
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert all(result.metrics["loss"] == 1 for result in worker_results)
+
+
+def test_worker_failure_local_rank(ray_start_4_cpus):
+    def train_func():
+        train.report({"rank": train.get_context().get_local_rank()})
+
+    def train_actor_failure():
+        import sys
+
+        sys.exit(1)
+
+    new_backend_executor_cls = gen_new_backend_executor(train_actor_failure)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        train_func, config, backend_executor_cls=new_backend_executor_cls
+    )
+    for worker_results in iterator:
+        assert {result.metrics["rank"] for result in worker_results} == {0, 1}
+
+
+def test_worker_start_failure(ray_start_4_cpus):
+    def init_hook():
+        pass
+
+    def init_hook_fail():
+        ray.actor.exit_actor()
+
+    class TestBackendExecutor(BackendExecutor):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def _restart(self):
+            self._initialization_hook = init_hook
+            super()._restart()
+
+    config = BackendConfig()
+
+    iterator = create_iterator(
+        lambda x: 1,
+        config,
+        backend_executor_cls=TestBackendExecutor,
+        init_hook=init_hook_fail,
+    )
+    assert len(iterator._backend_executor.get_worker_group()) == 2
+
+
+def test_max_failures(ray_start_4_cpus):
+    def train_func():
+        import sys
+
+        sys.exit(1)
+
+    config = BackendConfig()
+
+    iterator = create_iterator(train_func, config)
+    with pytest.raises(RuntimeError):
+        for _ in iterator:
+            pass
+    assert iterator._backend_executor._get_num_failures() == MAX_RETRIES
+
+
+def test_start_max_failures(ray_start_4_cpus):
+    def init_hook_fail():
+        import sys
+
+        sys.exit(1)
+
+    config = BackendConfig()
+
+    with pytest.raises(RuntimeError):
+        create_iterator(lambda x: 1, config, init_hook=init_hook_fail)
+
+
+class KillCallback:
+    def __init__(self, fail_on, backend_executor):
+        self.counter = 0
+        self.fail_on = fail_on
+        self.worker_group = backend_executor.get_worker_group()
+        self.results = []
+
+    def handle_result(self, intermiedate_results=None):
+        if intermiedate_results:
+            self.results.append(intermiedate_results)
+        if self.counter == self.fail_on:
+            print("killing")
+            self.results = []
+            ray.kill(self.worker_group.workers[0].actor)
+            time.sleep(3)
+        self.counter += 1
+
+
+@pytest.mark.parametrize(
+    "backend",
+    ["test", "torch", "tf"] if sys.version_info < (3, 12) else ["test", "torch"],
+)
+def test_worker_kill(ray_start_4_cpus, backend):
+    if backend == "test":
+        test_config = BackendConfig()
+    elif backend == "torch":
+        from ray.train.torch import TorchConfig
+
+        test_config = TorchConfig()
+    elif backend == "tf":
+        from ray.train.tensorflow import TensorflowConfig
+
+        test_config = TensorflowConfig()
+
+    def train_func():
+        for i in range(2):
+            train.report(dict(loss=1, iter=i))
+
+    iterator = create_iterator(train_func, test_config)
+    kill_callback = KillCallback(fail_on=0, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        # Run 1: iter=0, counter=1, Successful
+        # Run 2: iter=1, counter=1, Unsuccessful, starts training from beginning
+        # Run 3: iter=0, counter=2, Successful
+        # Run 4: iter=1, counter=3, Successful
+        kill_callback.handle_result()
+    assert kill_callback.counter == 3
+
+    iterator = create_iterator(train_func, test_config)
+    kill_callback = KillCallback(fail_on=1, backend_executor=iterator._backend_executor)
+    for intermediate_result in iterator:
+        # Run 1: iter=0, counter=1, Successful
+        # Run 2: iter=1, counter=2, Successful
+        # Run 3: None, counter=2, Unsuccessful, starts training from beginning.
+        # Run 4: iter=0, counter=3, Successful
+        # Run 5: iter=1, counter=4, Successful
+        kill_callback.handle_result()
+    assert kill_callback.counter == 4
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
+)
+def test_tensorflow_mnist_fail(ray_start_4_cpus):
+    """Tests if tensorflow example works even with worker failure."""
+    epochs = 3
+    num_workers = 2
+
+    from ray.train.examples.tf.tensorflow_mnist_example import (
+        train_func as tensorflow_mnist_train_func,
+    )
+    from ray.train.tensorflow import TensorflowConfig
+
+    test_config = TensorflowConfig()
+
+    train_func = functools.partial(
+        tensorflow_mnist_train_func, {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    )
+    iterator = create_iterator(train_func, test_config, num_workers=num_workers)
+    kill_callback = KillCallback(fail_on=0, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        assert len(intermediate_result) == num_workers
+        kill_callback.handle_result(intermediate_result)
+
+    results = kill_callback.results
+    assert len(results) == epochs
+    last_iter_result = results[-1][0].metrics
+    first_iter_result = results[0][0].metrics
+
+    assert last_iter_result["loss"] < first_iter_result["loss"]
+    assert last_iter_result["accuracy"] > first_iter_result["accuracy"]
+
+
+def test_torch_linear_failure(ray_start_4_cpus):
+    num_workers = 2
+    epochs = 3
+
+    from ray.train.torch import TorchConfig
+
+    test_config = TorchConfig()
+
+    train_func = functools.partial(
+        linear_train_func, {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    )
+
+    iterator = create_iterator(train_func, test_config, num_workers=num_workers)
+    kill_callback = KillCallback(fail_on=1, backend_executor=iterator._backend_executor)
+
+    for intermediate_result in iterator:
+        assert len(intermediate_result) == num_workers
+        kill_callback.handle_result(intermediate_result)
+
+    results = kill_callback.results
+    assert len(results) == epochs
+    for i in range(num_workers):
+        last_result = results[-1][i].metrics
+        first_result = results[0][i].metrics
+        assert last_result["loss"] < first_result["loss"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(sys.argv[1:] + ["-v", "-x", __file__]))