import json import os import time import pytest import ray from ray.cluster_utils import Cluster from ray.train import RunConfig, ScalingConfig from ray.train._internal.state.schema import ( ActorStatusEnum, RunStatusEnum, TrainDatasetInfo, TrainRunInfo, TrainWorkerInfo, ) from ray.train._internal.state.state_actor import ( TRAIN_STATE_ACTOR_NAME, TRAIN_STATE_ACTOR_NAMESPACE, get_or_create_state_actor, ) from ray.train._internal.state.state_manager import TrainRunStateManager from ray.train._internal.worker_group import WorkerGroup from ray.train.data_parallel_trainer import DataParallelTrainer @pytest.fixture def ray_start_gpu_cluster(): cluster = Cluster() cluster.add_node(num_gpus=8, num_cpus=9) ray.shutdown() ray.init( address=cluster.address, runtime_env={"env_vars": {"RAY_TRAIN_ENABLE_STATE_TRACKING": "1"}}, ignore_reinit_error=True, ) yield ray.shutdown() cluster.shutdown() RUN_INFO_JSON_SAMPLE = """{ "name": "default_run", "id": "ad5256bc64c04c83833a8b006f531799", "job_id": "0000000001", "controller_actor_id": "3abd1972a19148d78acc78dd9414736e", "start_time_ms": 1717448423000, "run_status": "RUNNING", "status_detail": "", "end_time_ms": null, "resources": [{"CPU": 1}, {"CPU": 1}], "workers": [ { "actor_id": "3d86c25634a71832dac32c8802000000", "world_rank": 0, "local_rank": 0, "node_rank": 0, "node_id": "b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825", "node_ip": "10.0.208.100", "pid": 76071, "gpu_ids": [0], "status": "ALIVE", "resources": {"CPU": 1} }, { "actor_id": "8f162dd8365346d1b5c98ebd7338c4f9", "world_rank": 1, "local_rank": 1, "node_rank": 0, "node_id": "b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825", "node_ip": "10.0.208.100", "pid": 76072, "gpu_ids": [1], "status": "ALIVE", "resources": {"CPU": 1} } ], "datasets": [ { "name": "train", "dataset_name": "train_dataset", "dataset_uuid": "1" } ] }""" def _get_run_info_sample(run_id=None, run_name=None) -> TrainRunInfo: dataset_info = TrainDatasetInfo( name="train", dataset_name="train_dataset", dataset_uuid="1" ) worker_info_0 = TrainWorkerInfo( actor_id="3d86c25634a71832dac32c8802000000", world_rank=0, local_rank=0, node_rank=0, node_id="b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825", node_ip="10.0.208.100", pid=76071, gpu_ids=[0], status=ActorStatusEnum.ALIVE, resources={"CPU": 1}, ) worker_info_1 = TrainWorkerInfo( actor_id="8f162dd8365346d1b5c98ebd7338c4f9", world_rank=1, local_rank=1, node_rank=0, node_id="b1e6cbed8533ae2def4e7e7ced9d19858ceb1ed8ab9ba81ab9c07825", node_ip="10.0.208.100", pid=76072, gpu_ids=[1], status=ActorStatusEnum.ALIVE, resources={"CPU": 1}, ) run_info = TrainRunInfo( name=run_name if run_name else "default_run", id=run_id if run_id else "ad5256bc64c04c83833a8b006f531799", job_id="0000000001", controller_actor_id="3abd1972a19148d78acc78dd9414736e", workers=[worker_info_0, worker_info_1], datasets=[dataset_info], start_time_ms=1717448423000, run_status=RunStatusEnum.RUNNING, status_detail="", resources=[{"CPU": 1}, {"CPU": 1}], ) return run_info def test_schema_equivalance(): json_sample = RUN_INFO_JSON_SAMPLE dict_sample = json.loads(RUN_INFO_JSON_SAMPLE) run_info_from_json = TrainRunInfo.parse_raw(json_sample) run_info_from_obj = TrainRunInfo.parse_obj(dict_sample) # Test serialization equivalence assert run_info_from_json == run_info_from_obj # Test dict deserialization equivalence assert run_info_from_json.dict() == dict_sample # Test json deserialization equivalence assert json.loads(run_info_from_json.json()) == json.loads(json_sample) # Test constructors equivalence assert _get_run_info_sample() == run_info_from_json def test_state_actor_api(ray_start_4_cpus): state_actor = get_or_create_state_actor() named_actors = ray.util.list_named_actors(all_namespaces=True) assert { "name": TRAIN_STATE_ACTOR_NAME, "namespace": TRAIN_STATE_ACTOR_NAMESPACE, } in named_actors # Concurrently register 100 runs num_runs = 100 info_list = [_get_run_info_sample(run_id=str(i)) for i in range(num_runs)] ray.get([state_actor.register_train_run.remote(run) for run in info_list]) # Test get all runs train_runs = ray.get(state_actor.get_all_train_runs.remote()) assert len(train_runs) == num_runs # Test get a single run by run_id for i in range(num_runs): run_info = ray.get(state_actor.get_train_run.remote(run_id=str(i))) assert run_info == info_list[i] def test_state_manager(ray_start_gpu_cluster): worker_group = WorkerGroup(num_workers=4, resources_per_worker={"GPU": 1}) # No errors raised if TrainStateActor is not started state_manager = TrainRunStateManager(state_actor=None) state_manager.register_train_run( run_id="run_id", run_name="run_name", job_id="0000000001", controller_actor_id="3abd1972a19148d78acc78dd9414736e", datasets={}, worker_group=worker_group, start_time_ms=int(time.time() * 1000), run_status=RunStatusEnum.RUNNING, resources=[{"CPU": 1}, {"CPU": 1}], ) # Register 100 runs with 10 TrainRunStateManagers state_actor = get_or_create_state_actor() for i in range(10): state_manager = TrainRunStateManager(state_actor=state_actor) for j in range(10): run_id = i * 10 + j state_manager.register_train_run( run_id=str(run_id), run_name="run_name", job_id="0000000001", controller_actor_id="3abd1972a19148d78acc78dd9414736e", datasets={ "train": ray.data.from_items(list(range(4))), "eval": ray.data.from_items(list(range(4))), }, worker_group=worker_group, start_time_ms=int(time.time() * 1000), run_status=RunStatusEnum.RUNNING, resources=[{"CPU": 1}, {"CPU": 1}], ) runs = ray.get(state_actor.get_all_train_runs.remote()) assert len(runs) == 100 for i in range(100): run_id = str(i) run_info = ray.get(state_actor.get_train_run.remote(run_id=run_id)) assert run_info and run_info.id == run_id @pytest.mark.parametrize("gpus_per_worker", [0, 1, 2]) def test_track_e2e_training(ray_start_gpu_cluster, gpus_per_worker): os.environ["RAY_TRAIN_ENABLE_STATE_TRACKING"] = "1" num_workers = 4 run_name = "test" datasets = { "train": ray.data.from_items(list(range(4))), "eval": ray.data.from_items(list(range(4))), } if gpus_per_worker == 0: use_gpu = False resources_per_worker = {"CPU": 1} else: use_gpu = True resources_per_worker = {"GPU": gpus_per_worker} trainer = DataParallelTrainer( train_loop_per_worker=lambda: None, run_config=RunConfig(name=run_name), scaling_config=ScalingConfig( num_workers=num_workers, use_gpu=use_gpu, resources_per_worker=resources_per_worker, ), datasets=datasets, ) trainer.fit() state_actor = ray.get_actor( name=TRAIN_STATE_ACTOR_NAME, namespace=TRAIN_STATE_ACTOR_NAMESPACE ) runs = ray.get(state_actor.get_all_train_runs.remote()) run_id = next(iter(runs.keys())) run = next(iter(runs.values())) # Check Run Info assert run.id == run_id assert run.name == run_name assert len(run.workers) == num_workers assert run.controller_actor_id and run.job_id world_ranks = [worker.world_rank for worker in run.workers] local_ranks = [worker.local_rank for worker in run.workers] node_ranks = [worker.node_rank for worker in run.workers] # Ensure that the workers are sorted by global rank assert world_ranks == [0, 1, 2, 3] assert local_ranks == [0, 1, 2, 3] assert node_ranks == [0, 0, 0, 0] # Check GPU ids gpu_ids = [worker.gpu_ids for worker in run.workers] if gpus_per_worker == 0: assert gpu_ids == [[], [], [], []] elif gpus_per_worker == 1: assert gpu_ids == [[0], [1], [2], [3]] elif gpus_per_worker == 2: flat_gpu_ids = set() for ids in gpu_ids: flat_gpu_ids.update(ids) assert flat_gpu_ids == set(range(8)) # Check Datasets for dataset_info in run.datasets: dataset = datasets[dataset_info.name] # DataConfig will automatically set the dataset_name to the key of the dataset dict. assert dataset_info.dataset_name == dataset_info.name assert dataset_info.dataset_uuid == dataset._uuid @pytest.mark.parametrize("raise_error", [True, False]) def test_train_run_status(ray_start_gpu_cluster, raise_error): os.environ["RAY_TRAIN_ENABLE_STATE_TRACKING"] = "1" def get_train_run(): state_actor = ray.get_actor( name=TRAIN_STATE_ACTOR_NAME, namespace=TRAIN_STATE_ACTOR_NAMESPACE ) runs = ray.get(state_actor.get_all_train_runs.remote()) return next(iter(runs.values())) def check_run_status(expected_status): run = get_train_run() assert run.run_status == expected_status def check_run_error(failed_rank, error_message): run = get_train_run() assert run.status_detail assert f"Rank {failed_rank} worker raised an error" in run.status_detail assert error_message in run.status_detail failed_rank = 0 error_message = "User Application Error" def train_func(): check_run_status(expected_status=RunStatusEnum.RUNNING) if raise_error and ray.train.get_context().get_world_rank() == failed_rank: raise RuntimeError(error_message) trainer = DataParallelTrainer( train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=4, use_gpu=False), ) try: trainer.fit() except Exception: pass if raise_error: check_run_status(expected_status=RunStatusEnum.ERRORED) check_run_error(failed_rank=failed_rank, error_message=error_message) else: check_run_status(expected_status=RunStatusEnum.FINISHED) ray.shutdown() if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-x", __file__]))