wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
dfc4684ba4becbf1b223029defa5db0711b9e3a4
ray-project--ray/python/ray/train/tests/test_state.py
T
wehub-resource-sync f1825c8ceb chore: import upstream snapshot with attribution
2026-07-13 13:17:40 +08:00

352 lines
11 KiB
Python
Raw Blame History

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__]))
Reference in New Issue View Git Blame Copy Permalink