wehub-resource-sync
127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
import uuid
|
|
|
|
import pytest
|
|
|
|
from ray.core.generated.export_train_state_pb2 import (
|
|
ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
|
|
ExportTrainRunEventData as ProtoTrainRun,
|
|
)
|
|
from ray.train._internal.state.export import (
|
|
train_run_info_to_proto_attempt,
|
|
train_run_info_to_proto_run,
|
|
)
|
|
from ray.train._internal.state.schema import (
|
|
ActorStatusEnum,
|
|
RunStatusEnum,
|
|
TrainRunInfo,
|
|
TrainWorkerInfo,
|
|
)
|
|
|
|
|
|
def create_mock_train_run_info(run_status=RunStatusEnum.RUNNING):
|
|
"""Create a minimal mock TrainRunInfo."""
|
|
worker = TrainWorkerInfo(
|
|
actor_id=uuid.uuid4().hex,
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
node_id=uuid.uuid4().hex,
|
|
node_ip="127.0.0.1",
|
|
pid=1234,
|
|
gpu_ids=[0],
|
|
status=ActorStatusEnum.ALIVE,
|
|
resources={"CPU": 1},
|
|
)
|
|
|
|
return TrainRunInfo(
|
|
name="test_run",
|
|
id=uuid.uuid4().hex,
|
|
job_id=uuid.uuid4().hex,
|
|
controller_actor_id=uuid.uuid4().hex,
|
|
workers=[worker],
|
|
datasets=[],
|
|
run_status=run_status,
|
|
status_detail="Error details" if run_status == RunStatusEnum.ERRORED else "",
|
|
start_time_ms=1000,
|
|
end_time_ms=2000
|
|
if run_status in [RunStatusEnum.FINISHED, RunStatusEnum.ERRORED]
|
|
else None,
|
|
resources=[{"CPU": 1}],
|
|
)
|
|
|
|
|
|
def test_train_run_info_to_proto_run():
|
|
"""Test that run info is correctly exported to proto run."""
|
|
run_info = create_mock_train_run_info()
|
|
proto_run = train_run_info_to_proto_run(run_info)
|
|
|
|
assert proto_run.ray_train_version == 1
|
|
assert proto_run.id == run_info.id
|
|
assert proto_run.name == run_info.name
|
|
assert proto_run.job_id == bytes.fromhex(run_info.job_id)
|
|
assert proto_run.status == ProtoTrainRun.RunStatus.RUNNING
|
|
|
|
|
|
def test_train_run_info_to_proto_attempt():
|
|
"""Test that run info is correctly exported to proto attempt."""
|
|
run_info = create_mock_train_run_info()
|
|
proto_attempt = train_run_info_to_proto_attempt(run_info)
|
|
|
|
assert proto_attempt.ray_train_version == 1
|
|
assert proto_attempt.run_id == run_info.id
|
|
assert proto_attempt.status == ProtoTrainRunAttempt.RunAttemptStatus.RUNNING
|
|
|
|
# Verify worker fields
|
|
assert len(proto_attempt.workers) == 1
|
|
proto_worker = proto_attempt.workers[0]
|
|
worker_info = run_info.workers[0]
|
|
|
|
assert proto_worker.world_rank == worker_info.world_rank
|
|
assert proto_worker.actor_id == bytes.fromhex(worker_info.actor_id)
|
|
|
|
|
|
def test_status_mapping():
|
|
"""Test that status is correctly mapped between schemas."""
|
|
status_pairs = [
|
|
(
|
|
RunStatusEnum.STARTED,
|
|
ProtoTrainRun.RunStatus.INITIALIZING,
|
|
ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
|
|
),
|
|
(
|
|
RunStatusEnum.RUNNING,
|
|
ProtoTrainRun.RunStatus.RUNNING,
|
|
ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
|
|
),
|
|
(
|
|
RunStatusEnum.FINISHED,
|
|
ProtoTrainRun.RunStatus.FINISHED,
|
|
ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
|
|
),
|
|
(
|
|
RunStatusEnum.ERRORED,
|
|
ProtoTrainRun.RunStatus.ERRORED,
|
|
ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
|
|
),
|
|
(
|
|
RunStatusEnum.ABORTED,
|
|
ProtoTrainRun.RunStatus.ABORTED,
|
|
ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
|
|
),
|
|
]
|
|
|
|
for run_status, expected_run_status, expected_attempt_status in status_pairs:
|
|
run_info = create_mock_train_run_info(run_status=run_status)
|
|
|
|
proto_run = train_run_info_to_proto_run(run_info)
|
|
assert proto_run.status == expected_run_status
|
|
|
|
proto_attempt = train_run_info_to_proto_attempt(run_info)
|
|
assert proto_attempt.status == expected_attempt_status
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|