wehub-resource-sync
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
from typing import Optional
|
|
|
|
from ray.core.generated.export_train_state_pb2 import (
|
|
ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
|
|
ExportTrainRunEventData as ProtoTrainRun,
|
|
)
|
|
from ray.train._internal.state.schema import (
|
|
ActorStatusEnum,
|
|
RunStatusEnum,
|
|
TrainRunInfo,
|
|
TrainWorkerInfo,
|
|
)
|
|
|
|
TRAIN_SCHEMA_VERSION = 1
|
|
RAY_TRAIN_VERSION = 1
|
|
|
|
# Status mapping dictionaries
|
|
_ACTOR_STATUS_MAP = {
|
|
ActorStatusEnum.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
|
|
ActorStatusEnum.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
|
|
}
|
|
|
|
_RUN_ATTEMPT_STATUS_MAP = {
|
|
RunStatusEnum.STARTED: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
|
|
RunStatusEnum.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
|
|
RunStatusEnum.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
|
|
RunStatusEnum.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
|
|
RunStatusEnum.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
|
|
}
|
|
|
|
_RUN_STATUS_MAP = {
|
|
RunStatusEnum.STARTED: ProtoTrainRun.RunStatus.INITIALIZING,
|
|
RunStatusEnum.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
|
|
RunStatusEnum.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
|
|
RunStatusEnum.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
|
|
RunStatusEnum.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
|
|
}
|
|
|
|
|
|
def _ms_to_ns(ms: Optional[int]) -> Optional[int]:
|
|
if ms is None:
|
|
return None
|
|
return ms * 1000000
|
|
|
|
|
|
# Helper conversion functions
|
|
def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
|
|
"""Convert resources dictionary to protobuf TrainResources."""
|
|
return ProtoTrainRunAttempt.TrainResources(resources=resources)
|
|
|
|
|
|
def _to_proto_worker(worker: TrainWorkerInfo) -> ProtoTrainRunAttempt.TrainWorker:
|
|
"""Convert TrainWorker to protobuf format."""
|
|
proto_worker = ProtoTrainRunAttempt.TrainWorker(
|
|
world_rank=worker.world_rank,
|
|
local_rank=worker.local_rank,
|
|
node_rank=worker.node_rank,
|
|
actor_id=bytes.fromhex(worker.actor_id),
|
|
node_id=bytes.fromhex(worker.node_id),
|
|
node_ip=worker.node_ip,
|
|
pid=worker.pid,
|
|
gpu_ids=worker.gpu_ids,
|
|
status=_ACTOR_STATUS_MAP[worker.status],
|
|
resources=_to_proto_resources(worker.resources),
|
|
)
|
|
|
|
return proto_worker
|
|
|
|
|
|
# Main conversion functions
|
|
def train_run_info_to_proto_run(run_info: TrainRunInfo) -> ProtoTrainRun:
|
|
"""Convert TrainRunInfo to TrainRun protobuf format."""
|
|
proto_run = ProtoTrainRun(
|
|
schema_version=TRAIN_SCHEMA_VERSION,
|
|
ray_train_version=RAY_TRAIN_VERSION,
|
|
id=run_info.id,
|
|
name=run_info.name,
|
|
job_id=bytes.fromhex(run_info.job_id),
|
|
controller_actor_id=bytes.fromhex(run_info.controller_actor_id),
|
|
status=_RUN_STATUS_MAP[run_info.run_status],
|
|
status_detail=run_info.status_detail,
|
|
start_time_ns=_ms_to_ns(run_info.start_time_ms),
|
|
end_time_ns=_ms_to_ns(run_info.end_time_ms),
|
|
)
|
|
|
|
return proto_run
|
|
|
|
|
|
def train_run_info_to_proto_attempt(run_info: TrainRunInfo) -> ProtoTrainRunAttempt:
|
|
"""Convert TrainRunInfo to TrainRunAttempt protobuf format."""
|
|
|
|
proto_attempt = ProtoTrainRunAttempt(
|
|
schema_version=TRAIN_SCHEMA_VERSION,
|
|
ray_train_version=RAY_TRAIN_VERSION,
|
|
run_id=run_info.id,
|
|
attempt_id=run_info.id, # Same as run_id
|
|
status=_RUN_ATTEMPT_STATUS_MAP[run_info.run_status],
|
|
status_detail=run_info.status_detail,
|
|
start_time_ns=_ms_to_ns(run_info.start_time_ms),
|
|
end_time_ns=_ms_to_ns(run_info.end_time_ms),
|
|
resources=[_to_proto_resources(r) for r in run_info.resources],
|
|
workers=[_to_proto_worker(worker) for worker in run_info.workers],
|
|
)
|
|
|
|
return proto_attempt
|