import os import sys import time from pathlib import Path from typing import List import pytest import ray import ray.train from ray._common.test_utils import wait_for_condition 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.jax import JaxTrainer @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) @pytest.mark.skipif( sys.version_info >= (3, 12), reason="Current jax version (0.4.13) is not supported in python 3.12+", ) def test_elastic_training_tpu(monkeypatch, tmp_path, cluster): """End to end test for TPU elastic training with the JaxTrainer.""" unit_time_s = 1.0 health_check_interval_s = unit_time_s elastic_resize_monitor_interval_s = unit_time_s * 5 num_epochs = 30 monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, str(health_check_interval_s)) monkeypatch.setenv("JAX_PLATFORMS", "cpu") @ray.remote(num_cpus=0) def run_training(): trainer = JaxTrainer( train_fn, train_loop_config={ "num_epochs": num_epochs, "health_check_interval_s": health_check_interval_s, }, scaling_config=ray.train.ScalingConfig( use_tpu=True, accelerator_type="TPU-V6E", topology="2x4", resources_per_worker={"TPU": 4, "CPU": 1}, num_workers=(2, 6), # Scale between 1 and 3 slices. 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), failure_config=ray.train.FailureConfig(max_failures=3), ), ) return trainer.fit() run_training_future = run_training.remote() start = time.time() ALL_NODES = [] def print_status(message): elapsed = time.time() - start print(f"\n{'-' * 80}") cluster_resources = { resource: value for resource, value in ray.cluster_resources().items() if "TPU" in resource or "CPU" in resource } print(f"[elapsed={elapsed:.1f}s] {cluster_resources=}") print(message) print(f"{'-' * 80}\n") def provision_tpu_node(slice_name: str, worker_id: int, is_head: bool = False): pod_type = "v6e-8" topology = "2x4" node_env = { "PATH": os.environ.get("PATH", ""), "TPU_NAME": slice_name, "TPU_WORKER_ID": str(worker_id), "TPU_ACCELERATOR_TYPE": pod_type, "TPU_TOPOLOGY": topology, "JAX_PLATFORMS": "cpu", HEALTH_CHECK_INTERVAL_S_ENV_VAR: str(health_check_interval_s), } labels = { "ray.io/tpu-slice-name": slice_name, "ray.io/tpu-worker-id": str(worker_id), "ray.io/tpu-pod-type": pod_type, } resources = {"TPU": 4, "accelerator_type:TPU-V6E": 1} if is_head: resources[f"TPU-{pod_type}-head"] = 1 node = cluster.add_node( num_cpus=8, resources=resources, labels=labels, env_vars=node_env, wait=True, ) return node def verify_active_workers(expected_count: int) -> bool: try: from ray.util.state import list_actors workers = list_actors( filters=[("class_name", "=", "RayTrainWorker"), ("state", "=", "ALIVE")] ) return len(workers) == expected_count except Exception: # Ignore transient State API errors during cluster churn return False def remove_nodes(nodes: List): for node in nodes: cluster.remove_node(node) cluster.wait_for_nodes() print_status(f"Removed {len(nodes)} node(s).") print_status( "Adding 1 TPU node. Waiting for training to ignore it since it's not a full slice." ) ALL_NODES.append( provision_tpu_node(slice_name="slice-A", worker_id=0, is_head=True) ) print_status("Adding 2nd TPU node to complete slice-A. Training should start.") ALL_NODES.append( provision_tpu_node(slice_name="slice-A", worker_id=1, is_head=False) ) print_status("Waiting for initial scale-up to 2 workers...") wait_for_condition(lambda: verify_active_workers(2), timeout=120) print_status("Adding full second TPU slice. Policy should upscale.") ALL_NODES.append( provision_tpu_node(slice_name="slice-B", worker_id=0, is_head=True) ) ALL_NODES.append( provision_tpu_node(slice_name="slice-B", worker_id=1, is_head=False) ) print_status("Waiting for elastic scale-up to 4 workers.") wait_for_condition(lambda: verify_active_workers(4), timeout=120) # Run a couple of epochs at max capacity so the metrics reflect the max_world_size. time.sleep(8) # Multi-host TPUs on GKE with KubeRay are scaled atomically in slices. print_status("Killing second TPU slice to simulate full slice preemption.") node_b_worker = ALL_NODES.pop() node_b_head = ALL_NODES.pop() remove_nodes([node_b_worker, node_b_head]) print_status("Waiting for policy to scale down and recover with 2 workers.") wait_for_condition(lambda: verify_active_workers(2), timeout=120) # Run a couple of epochs after the recovery. time.sleep(8) 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"] == 2 assert result.metrics["max_world_size"] == 4 assert result.checkpoint assert Path(result.checkpoint.path).name == f"checkpoint-epoch={num_epochs}" if __name__ == "__main__": sys.exit(pytest.main(["-v", "-x", __file__]))