chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_elastic_e2e.py

@@ -0,0 +1,206 @@
+import sys
+import time
+from pathlib import Path
+from typing import List
+
+import pytest
+
+import ray
+import ray.train
+from ray.cluster_utils import Cluster
+from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
+from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def cluster():
+    cluster = Cluster(initialize_head=True, head_node_args=dict(num_cpus=0))
+    cluster.wait_for_nodes()
+    ray.init(
+        address=cluster.address,
+        runtime_env={"working_dir": str(Path(__file__).parent)},
+    )
+    yield cluster
+    ray.shutdown()
+    cluster.shutdown()
+
+
+def train_fn(config: dict):
+    train_context = ray.train.get_context()
+    rank = train_context.get_world_rank()
+
+    start_epoch = 1
+    checkpoint = ray.train.get_checkpoint()
+    min_world_size = None
+    max_world_size = None
+    if checkpoint:
+        checkpoint_data = load_dict_checkpoint(checkpoint)
+        start_epoch = checkpoint_data["epoch"] + 1
+        min_world_size = checkpoint_data.get("min_world_size")
+        max_world_size = checkpoint_data.get("max_world_size")
+        if rank == 0:
+            print("Restoring from epoch: ", start_epoch)
+
+    for epoch in range(start_epoch, config.get("num_epochs", 60) + 1):
+        world_size = train_context.get_world_size()
+        if min_world_size is None:
+            min_world_size = world_size
+        if max_world_size is None:
+            max_world_size = world_size
+        min_world_size = min(min_world_size, world_size)
+        max_world_size = max(max_world_size, world_size)
+        # TODO: This test injects errors by "killing nodes," which ungracefully
+        # kills processes. This means that any backlog in the checkpoint queue
+        # will not be flushed to the controller.
+        # This means that the checkpoint populated on restore may not be
+        # the most recent one.
+        # Set the poll interval < health check interval to reduce the
+        # backlog size to mitigate the issue.
+        time.sleep(2 * config.get("health_check_interval_s", 1))
+
+        with create_dict_checkpoint(
+            {
+                "epoch": epoch,
+                "min_world_size": min_world_size,
+                "max_world_size": max_world_size,
+            }
+        ) as checkpoint:
+            ray.train.report(
+                {
+                    "epoch": epoch,
+                    "world_size": world_size,
+                    "min_world_size": min_world_size,
+                    "max_world_size": max_world_size,
+                },
+                checkpoint=checkpoint if rank == 0 else None,
+                checkpoint_dir_name=f"checkpoint-epoch={epoch}",
+            )
+        if rank == 0:
+            print("Finished epoch: ", epoch)
+
+
+def test_elastic_training(monkeypatch, tmp_path, cluster):
+    """End to end test for elastic training.
+
+    This test covers:
+    * Elastic startup (0 resources -> min resources)
+    * Elastic scale up while running (min resources -> max resources)
+    * Elastic scale down due to failure while running
+    * Checkpointing + restoration
+    * Preemption failure handling
+    """
+    unit_time_s = 0.1
+    health_check_interval_s = unit_time_s
+    elastic_resize_monitor_interval_s = unit_time_s * 10
+    num_epochs = 30
+
+    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, str(health_check_interval_s))
+
+    @ray.remote(num_cpus=0)
+    def run_training():
+        trainer = DataParallelTrainer(
+            train_fn,
+            train_loop_config={
+                "num_epochs": num_epochs,
+                "health_check_interval_s": health_check_interval_s,
+            },
+            scaling_config=ray.train.ScalingConfig(
+                num_workers=(4, 32),
+                use_gpu=True,
+                elastic_resize_monitor_interval_s=elastic_resize_monitor_interval_s,
+            ),
+            run_config=ray.train.RunConfig(
+                storage_path=str(tmp_path),
+                checkpoint_config=ray.train.CheckpointConfig(num_to_keep=2),
+                # NOTE: The outer test script will inject 2 node failures.
+                failure_config=ray.train.FailureConfig(max_failures=2),
+            ),
+        )
+        return trainer.fit()
+
+    run_training_future = run_training.remote()
+
+    start = time.time()
+    ALL_NODES = []
+
+    def print_status(message):
+        elapsed = time.time() - start
+        print()
+        print("-" * 80)
+        cluster_resources = {
+            resource: value
+            for resource, value in ray.cluster_resources().items()
+            if resource in ("CPU", "GPU")
+        }
+        print(f"[elapsed={elapsed:.1f}s] {cluster_resources=}")
+        print(message)
+        print("-" * 80)
+        print()
+
+    def sleep(num_units):
+        time.sleep(unit_time_s * num_units)
+
+    def add_nodes(gpus: List[int]) -> List:
+        added_nodes = []
+        for num_gpus in gpus:
+            node = cluster.add_node(num_gpus=num_gpus, wait=False)
+            added_nodes.append(node)
+
+        print_status(f"Added {len(gpus)} node(s) with num_gpus: {gpus}")
+        cluster.wait_for_nodes()
+        return added_nodes
+
+    def remove_nodes(nodes: List):
+        for node in nodes:
+            cluster.remove_node(node)
+
+        cluster.wait_for_nodes()
+        print_status(f"Removed nodes: {nodes}")
+
+    # Wait a bit before adding resources.
+    print_status("Waiting for training to start...")
+    sleep(8)
+
+    # Add a node with 4 GPUs
+    ALL_NODES.extend(add_nodes([4]))
+
+    # Wait a bit before adding more resources.
+    sleep(8)
+    print("Adding 4 GPU node.")
+    ALL_NODES.extend(add_nodes([4]))
+    sleep(1)
+    ALL_NODES.extend(add_nodes([4]))
+    # Should not upscale here due to the elastic resize monitor interval.
+
+    # Should upscale to 12 during this sleep.
+    sleep(20)
+
+    # Kill a node.
+    remove_nodes([ALL_NODES.pop(0)])
+    sleep(12)
+
+    # Kill all worker nodes.
+    remove_nodes(ALL_NODES)
+    ALL_NODES = []
+
+    sleep(8)
+    ALL_NODES.extend(add_nodes(gpus=[1] * 16))
+
+    sleep(12)
+    # 4 extra GPUs shouldn't be used.
+    ALL_NODES.extend(add_nodes(gpus=[4] * 4 + [1] * 4))
+
+    result: ray.train.Result = ray.get(run_training_future)
+
+    print_status(f"Training finished with result: {result}")
+    assert not result.error
+    assert result.metrics["min_world_size"] >= 4
+    assert result.metrics["max_world_size"] <= 32
+    assert result.metrics["max_world_size"] >= result.metrics["min_world_size"]
+    assert result.checkpoint
+    assert Path(result.checkpoint.path).name == f"checkpoint-epoch={num_epochs}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-x", __file__]))