chore: import upstream snapshot with attribution

python/ray/train/_internal/state/export.py

@@ -0,0 +1,105 @@
+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