chore: import upstream snapshot with attribution

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

@@ -0,0 +1,313 @@
+import logging
+
+from google.protobuf.struct_pb2 import Struct
+
+from ray.core.generated.export_train_state_pb2 import (
+    ExportTrainRunAttemptEventData as ProtoTrainRunAttempt,
+    ExportTrainRunEventData as ProtoTrainRun,
+)
+from ray.dashboard.modules.metrics.dashboards.common import Panel
+from ray.dashboard.modules.metrics.dashboards.train_dashboard_panels import (
+    TRAIN_RUN_PANELS,
+    TRAIN_WORKER_PANELS,
+)
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    BackendConfig,
+    DataConfig,
+    ExecutionOptions,
+    RunAttemptStatus,
+    RunConfig,
+    RunSettings,
+    RunStatus,
+    ScalingConfig,
+    TrainRun,
+    TrainRunAttempt,
+    TrainWorker,
+)
+from ray.train.v2._internal.util import TrainingFramework
+
+# Increment each time the exported Train schema changes (proto, pydantic, or
+# exported json) so downstream consumers can distinguish schema versions.
+TRAIN_SCHEMA_VERSION = 4
+RAY_TRAIN_VERSION = 2
+
+# Status mapping dictionaries
+_ACTOR_STATUS_MAP = {
+    ActorStatus.ALIVE: ProtoTrainRunAttempt.ActorStatus.ALIVE,
+    ActorStatus.DEAD: ProtoTrainRunAttempt.ActorStatus.DEAD,
+}
+
+_RUN_ATTEMPT_STATUS_MAP = {
+    RunAttemptStatus.PENDING: ProtoTrainRunAttempt.RunAttemptStatus.PENDING,
+    RunAttemptStatus.RUNNING: ProtoTrainRunAttempt.RunAttemptStatus.RUNNING,
+    RunAttemptStatus.FINISHED: ProtoTrainRunAttempt.RunAttemptStatus.FINISHED,
+    RunAttemptStatus.ERRORED: ProtoTrainRunAttempt.RunAttemptStatus.ERRORED,
+    RunAttemptStatus.ABORTED: ProtoTrainRunAttempt.RunAttemptStatus.ABORTED,
+}
+
+_RUN_STATUS_MAP = {
+    RunStatus.INITIALIZING: ProtoTrainRun.RunStatus.INITIALIZING,
+    RunStatus.SCHEDULING: ProtoTrainRun.RunStatus.SCHEDULING,
+    RunStatus.RUNNING: ProtoTrainRun.RunStatus.RUNNING,
+    RunStatus.RESTARTING: ProtoTrainRun.RunStatus.RESTARTING,
+    RunStatus.RESIZING: ProtoTrainRun.RunStatus.RESIZING,
+    RunStatus.FINISHED: ProtoTrainRun.RunStatus.FINISHED,
+    RunStatus.ERRORED: ProtoTrainRun.RunStatus.ERRORED,
+    RunStatus.ABORTED: ProtoTrainRun.RunStatus.ABORTED,
+}
+
+_TRAINING_FRAMEWORK_MAP = {
+    None: ProtoTrainRun.BackendConfig.TrainingFramework.TRAINING_FRAMEWORK_UNSPECIFIED,
+    TrainingFramework.TORCH: ProtoTrainRun.BackendConfig.TrainingFramework.TORCH,
+    TrainingFramework.JAX: ProtoTrainRun.BackendConfig.TrainingFramework.JAX,
+    TrainingFramework.TENSORFLOW: ProtoTrainRun.BackendConfig.TrainingFramework.TENSORFLOW,
+    TrainingFramework.XGBOOST: ProtoTrainRun.BackendConfig.TrainingFramework.XGBOOST,
+    TrainingFramework.LIGHTGBM: ProtoTrainRun.BackendConfig.TrainingFramework.LIGHTGBM,
+}
+
+logger = logging.getLogger(__name__)
+
+
+def _dict_to_struct(d: dict) -> Struct:
+    """Returns a protobuf Struct from a dictionary."""
+    s = Struct()
+    s.update(d)
+    return s
+
+
+def _to_proto_resources(resources: dict) -> ProtoTrainRunAttempt.TrainResources:
+    """Convert resources dictionary to protobuf TrainResources."""
+    return ProtoTrainRunAttempt.TrainResources(resources=resources)
+
+
+def _to_proto_worker(worker: TrainWorker) -> ProtoTrainRunAttempt.TrainWorker:
+    """Convert TrainWorker to protobuf format."""
+    status = None
+    if worker.status is not None:
+        status = _ACTOR_STATUS_MAP[worker.status]
+
+    return 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=status,
+        resources=_to_proto_resources(worker.resources.resources),
+        log_file_path=worker.log_file_path,
+    )
+
+
+# Main conversion functions
+def train_run_attempt_to_proto(attempt: TrainRunAttempt) -> ProtoTrainRunAttempt:
+    """Convert TrainRunAttempt to protobuf format."""
+    proto_attempt = ProtoTrainRunAttempt(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        run_id=attempt.run_id,
+        attempt_id=attempt.attempt_id,
+        status=_RUN_ATTEMPT_STATUS_MAP[attempt.status],
+        status_detail=attempt.status_detail,
+        start_time_ns=attempt.start_time_ns,
+        end_time_ns=attempt.end_time_ns,
+        resources=[_to_proto_resources(r.resources) for r in attempt.resources],
+        workers=[_to_proto_worker(w) for w in attempt.workers],
+    )
+
+    return proto_attempt
+
+
+def _to_proto_dashboard_panel(panel: Panel) -> ProtoTrainRun.DashboardPanelMetadata:
+    """Convert Dashboard Panel to protobuf format."""
+    proto_panel = ProtoTrainRun.DashboardPanelMetadata(
+        id=str(panel.id),
+        title=panel.title,
+    )
+
+    return proto_panel
+
+
+def to_proto_backend_config(
+    backend_config: BackendConfig,
+) -> ProtoTrainRun.BackendConfig:
+    """Convert BackendConfig to protobuf format."""
+    proto_backend_config = ProtoTrainRun.BackendConfig(
+        framework=_TRAINING_FRAMEWORK_MAP[backend_config.framework],
+    )
+
+    proto_backend_config.config.CopyFrom(_dict_to_struct(backend_config.config))
+
+    return proto_backend_config
+
+
+def to_proto_scaling_config(
+    scaling_config: ScalingConfig,
+) -> ProtoTrainRun.ScalingConfig:
+    """Convert ScalingConfig to protobuf format."""
+    proto_scaling_config = ProtoTrainRun.ScalingConfig(
+        use_gpu=scaling_config.use_gpu,
+        placement_strategy=scaling_config.placement_strategy,
+        use_tpu=scaling_config.use_tpu,
+    )
+
+    if isinstance(scaling_config.num_workers, tuple):
+        proto_scaling_config.num_workers_range.CopyFrom(
+            ProtoTrainRun.ScalingConfig.IntRange(
+                min=scaling_config.num_workers[0],
+                max=scaling_config.num_workers[1],
+            )
+        )
+    else:
+        proto_scaling_config.num_workers_fixed = scaling_config.num_workers
+
+    if scaling_config.resources_per_worker is not None:
+        proto_scaling_config.resources_per_worker.values.update(
+            scaling_config.resources_per_worker
+        )
+
+    if scaling_config.accelerator_type is not None:
+        proto_scaling_config.accelerator_type = scaling_config.accelerator_type
+    if scaling_config.topology is not None:
+        proto_scaling_config.topology = scaling_config.topology
+
+    if scaling_config.bundle_label_selector is not None:
+        selectors = scaling_config.bundle_label_selector
+        if isinstance(selectors, dict):
+            proto_scaling_config.label_selector_single.values.update(selectors)
+        else:
+            proto_scaling_config.label_selector_list.values.extend(
+                [ProtoTrainRun.ScalingConfig.StringMap(values=s) for s in selectors]
+            )
+
+    return proto_scaling_config
+
+
+def _to_proto_execution_options(
+    execution_options: ExecutionOptions,
+) -> ProtoTrainRun.ExecutionOptions:
+    """Convert a single ExecutionOptions schema model to protobuf."""
+    return ProtoTrainRun.ExecutionOptions(
+        resource_limits=_dict_to_struct(execution_options.resource_limits),
+        exclude_resources=_dict_to_struct(execution_options.exclude_resources),
+        preserve_order=execution_options.preserve_order,
+        actor_locality_enabled=execution_options.actor_locality_enabled,
+        verbose_progress=execution_options.verbose_progress,
+    )
+
+
+def to_proto_data_config(data_config: DataConfig) -> ProtoTrainRun.DataConfig:
+    """Convert DataConfig to protobuf format."""
+    data_execution_options = data_config.data_execution_options
+    proto_data_config = ProtoTrainRun.DataConfig(
+        enable_shard_locality=data_config.enable_shard_locality,
+        data_execution_options=ProtoTrainRun.DataExecutionOptions(
+            default=_to_proto_execution_options(data_execution_options.default),
+            per_dataset_execution_options={
+                name: _to_proto_execution_options(opts)
+                for name, opts in data_execution_options.per_dataset_execution_options.items()
+            },
+        ),
+    )
+
+    if data_config.datasets_to_split == "all":
+        proto_data_config.all.SetInParent()
+    else:
+        proto_data_config.datasets.values.extend(data_config.datasets_to_split)
+
+    return proto_data_config
+
+
+def _to_proto_failure_config(run_config: RunConfig) -> ProtoTrainRun.FailureConfig:
+    """Convert RunConfig.failure_config to protobuf format."""
+    return ProtoTrainRun.FailureConfig(
+        max_failures=run_config.failure_config.max_failures,
+        controller_failure_limit=run_config.failure_config.controller_failure_limit,
+    )
+
+
+def _to_proto_checkpoint_config(
+    run_config: RunConfig,
+) -> ProtoTrainRun.CheckpointConfig:
+    """Convert RunConfig.checkpoint_config to protobuf format."""
+    checkpoint_score_order = ProtoTrainRun.CheckpointConfig.CheckpointScoreOrder.Value(
+        run_config.checkpoint_config.checkpoint_score_order.upper()
+    )
+
+    proto_checkpoint_config = ProtoTrainRun.CheckpointConfig(
+        checkpoint_score_order=checkpoint_score_order
+    )
+    if run_config.checkpoint_config.num_to_keep is not None:
+        proto_checkpoint_config.num_to_keep = run_config.checkpoint_config.num_to_keep
+    if run_config.checkpoint_config.checkpoint_score_attribute is not None:
+        proto_checkpoint_config.checkpoint_score_attribute = (
+            run_config.checkpoint_config.checkpoint_score_attribute
+        )
+    return proto_checkpoint_config
+
+
+def to_proto_run_config(run_config: RunConfig) -> ProtoTrainRun.RunConfig:
+    """Convert RunConfig to protobuf format."""
+    proto_run_config = ProtoTrainRun.RunConfig(
+        name=run_config.name,
+        failure_config=_to_proto_failure_config(run_config),
+        worker_runtime_env=_dict_to_struct(run_config.worker_runtime_env),
+        checkpoint_config=_to_proto_checkpoint_config(run_config),
+        storage_path=run_config.storage_path,
+    )
+
+    if run_config.storage_filesystem is not None:
+        proto_run_config.storage_filesystem = run_config.storage_filesystem
+
+    return proto_run_config
+
+
+def _to_proto_run_settings(run_settings: RunSettings) -> ProtoTrainRun.RunSettings:
+    """Convert RunSettings to protobuf format."""
+
+    proto_run_settings = ProtoTrainRun.RunSettings(
+        backend_config=to_proto_backend_config(run_settings.backend_config),
+        scaling_config=to_proto_scaling_config(run_settings.scaling_config),
+        datasets=run_settings.datasets,
+        data_config=to_proto_data_config(run_settings.data_config),
+        run_config=to_proto_run_config(run_settings.run_config),
+    )
+
+    if run_settings.train_loop_config is not None:
+        proto_run_settings.train_loop_config.CopyFrom(
+            _dict_to_struct(run_settings.train_loop_config)
+        )
+
+    return proto_run_settings
+
+
+def train_run_to_proto(run: TrainRun) -> ProtoTrainRun:
+    """Convert TrainRun to protobuf format."""
+
+    proto_train_run_panels = [_to_proto_dashboard_panel(p) for p in TRAIN_RUN_PANELS]
+    proto_train_worker_panels = [
+        _to_proto_dashboard_panel(p) for p in TRAIN_WORKER_PANELS
+    ]
+
+    proto_train_run = ProtoTrainRun(
+        schema_version=TRAIN_SCHEMA_VERSION,
+        ray_train_version=RAY_TRAIN_VERSION,
+        id=run.id,
+        name=run.name,
+        job_id=bytes.fromhex(run.job_id),
+        controller_actor_id=bytes.fromhex(run.controller_actor_id),
+        status=_RUN_STATUS_MAP[run.status],
+        status_detail=run.status_detail,
+        start_time_ns=run.start_time_ns,
+        end_time_ns=run.end_time_ns,
+        controller_log_file_path=run.controller_log_file_path,
+        train_run_panels=proto_train_run_panels,
+        train_worker_panels=proto_train_worker_panels,
+        framework_versions=run.framework_versions,
+        run_settings=_to_proto_run_settings(run.run_settings),
+    )
+
+    return proto_train_run

python/ray/train/v2/_internal/state/schema.py

@@ -0,0 +1,532 @@
+import math
+from collections.abc import Mapping
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+from pydantic import field_validator
+
+from ray._common.pydantic_compat import BaseModel, Field
+from ray.dashboard.modules.job.pydantic_models import JobDetails
+from ray.train.v2._internal.util import TrainingFramework
+from ray.util.annotations import DeveloperAPI
+
+MAX_ERROR_STACK_TRACE_LENGTH = 50000
+
+
+def _to_json_serializable_value(value: Any, *, max_depth: int = 3) -> Any:
+    """Recursively coerce a value into a human-readable, JSON serializable representation.
+
+    If ``value`` is a list or dict, this function walks through it and replaces non-JSON
+    serializable fields (e.g. custom objects, modules, tensors, callables, etc.) with a
+    human-readable string representation.
+
+    Args:
+        value: Any Python value. Primitives pass through; collections recurse;
+            other types are stringified.
+        max_depth: Truncates dicts nested beyond ``max_depth`` to ``"..."``.
+            Lists do not consume depth.
+
+    Returns:
+        The JSON serializable representation of the value.
+    """
+    if max_depth <= 0:
+        raise ValueError("max_depth must be greater than 0")
+
+    def _safe_str(v):
+        try:
+            return str(v)
+        except Exception:
+            return type(v).__name__
+
+    def _walk(value, depth):
+        if value is None or isinstance(value, (bool, int, str)):
+            return value
+        if isinstance(value, float):
+            return str(value) if not math.isfinite(value) else value
+        if isinstance(value, Mapping):
+            if depth <= 0:
+                return "..."
+            try:
+                items = list(value.items())
+            except Exception:
+                # Custom Mapping subclass with a broken `.items()`.
+                return type(value).__name__
+            return {_safe_str(k): _walk(v, depth - 1) for k, v in items}
+
+        # Tuples, sets, and frozensets all become lists in JSON.
+        if isinstance(value, (list, tuple, set, frozenset)):
+            return [_walk(v, depth) for v in value]
+
+        cls = type(value)
+        # Use class name if no custom string representation is defined.
+        if cls.__str__ is object.__str__ and cls.__repr__ is object.__repr__:
+            return cls.__name__
+        return _safe_str(value)
+
+    return _walk(value, max_depth)
+
+
+@DeveloperAPI
+class RunStatus(str, Enum):
+    """Enumeration of the possible statuses for a Train run."""
+
+    # ====== Active States ======
+    # The Train run is currently in the process of initializing.
+    INITIALIZING = "INITIALIZING"
+    # The Train run is waiting to be scheduled.
+    SCHEDULING = "SCHEDULING"
+    # The Train run is currently in progress.
+    RUNNING = "RUNNING"
+    # The Train run is recovering from a failure or restart.
+    RESTARTING = "RESTARTING"
+    # The Train run is resizing.
+    RESIZING = "RESIZING"
+
+    # ===== Terminal States ======
+    # The Train run completed successfully.
+    FINISHED = "FINISHED"
+    # The Train run failed due to an error in the training workers.
+    ERRORED = "ERRORED"
+    # The Train run was terminated due to system or controller errors.
+    ABORTED = "ABORTED"
+
+    def is_terminal(self) -> bool:
+        return self in [RunStatus.FINISHED, RunStatus.ERRORED, RunStatus.ABORTED]
+
+
+@DeveloperAPI
+class RunAttemptStatus(str, Enum):
+    """Enumeration of the possible statuses for a Train run attempt."""
+
+    # ====== Active States ======
+    # The run attempt is waiting to be scheduled.
+    PENDING = "PENDING"
+    # The run attempt is currently in progress.
+    RUNNING = "RUNNING"
+
+    # ===== Terminal States =====
+    # The run attempt completed successfully.
+    FINISHED = "FINISHED"
+    # The run attempt failed due to an error in the training workers.
+    ERRORED = "ERRORED"
+    # The run attempt was terminated due to system or controller errors.
+    ABORTED = "ABORTED"
+
+    def is_terminal(self) -> bool:
+        return self in [
+            RunAttemptStatus.FINISHED,
+            RunAttemptStatus.ERRORED,
+            RunAttemptStatus.ABORTED,
+        ]
+
+
+@DeveloperAPI
+class ActorStatus(str, Enum):
+    """Enumeration of the statuses for a Train worker actor."""
+
+    # The actor is currently active.
+    ALIVE = "ALIVE"
+    # The actor is no longer active.
+    DEAD = "DEAD"
+
+
+@DeveloperAPI
+class TrainResources(BaseModel):
+    """Resources allocated for a Train worker or run."""
+
+    resources: Dict[str, float] = Field(
+        description="A dictionary specifying the types and amounts of resources "
+        "allocated (e.g., CPU, GPU)."
+    )
+
+
+@DeveloperAPI
+class TrainWorker(BaseModel):
+    """Metadata about a Ray Train worker."""
+
+    world_rank: int = Field(
+        description="The global rank of the worker in the training cluster."
+    )
+    local_rank: int = Field(description="The local rank of the worker on its node.")
+    node_rank: int = Field(description="The rank of the worker's node in the cluster.")
+    actor_id: str = Field(description="The unique ID of the worker's actor.")
+    node_id: str = Field(
+        description="The unique ID of the node where the worker is running."
+    )
+    node_ip: str = Field(
+        description="The IP address of the node where the worker is running."
+    )
+    pid: int = Field(description="The process ID of the worker.")
+    gpu_ids: List[int] = Field(description="A list of GPU IDs allocated to the worker.")
+    status: Optional[ActorStatus] = Field(
+        None, description="The current status of the worker actor."
+    )
+    resources: TrainResources = Field(
+        description="The resources allocated to this Train worker."
+    )
+    log_file_path: Optional[str] = Field(
+        None, description="The path to the log file for the Train worker."
+    )
+
+
+@DeveloperAPI
+class MemoryInfo(BaseModel):
+    """Memory usage information for a process."""
+
+    rss: int = Field(description="The resident set size (RSS) memory usage in bytes.")
+    vms: int = Field(description="The virtual memory size (VMS) usage in bytes.")
+    pfaults: Optional[int] = Field(None, description="The number of page faults.")
+    pageins: Optional[int] = Field(None, description="The number of page-ins.")
+
+
+@DeveloperAPI
+class ProcessStats(BaseModel):
+    """CPU and memory statistics for a process."""
+
+    cpuPercent: float = Field(description="The percentage of CPU usage.")
+    mem: Optional[List[int]] = Field(
+        None,
+        description="Memory statistics, including total memory, free memory, "
+        "and memory usage ratio.",
+    )
+    memoryInfo: MemoryInfo = Field(description="Detailed memory usage information.")
+
+
+class ProcessGPUUsage(BaseModel):
+    """GPU usage statistics for a process."""
+
+    pid: int = Field(description="The process ID.")
+    gpuMemoryUsage: int = Field(description="The GPU memory usage in bytes.")
+
+
+@DeveloperAPI
+class GPUStats(BaseModel):
+    """Statistics for a GPU."""
+
+    uuid: str = Field(description="The unique identifier of the GPU.")
+    index: int = Field(description="The index of the GPU.")
+    name: str = Field(description="The name of the GPU.")
+    utilizationGpu: Optional[float] = Field(
+        None, description="The percentage utilization of the GPU."
+    )
+    memoryUsed: float = Field(description="The amount of GPU memory used in bytes.")
+    memoryTotal: float = Field(description="The total amount of GPU memory in bytes.")
+    processInfo: ProcessGPUUsage = Field(
+        description="GPU usage statistics for the associated process."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainWorker(TrainWorker):
+    """Detailed metadata for a Ray Train worker, including process and GPU stats."""
+
+    processStats: Optional[ProcessStats] = Field(
+        None, description="CPU and memory statistics for the worker process."
+    )
+    gpus: List[GPUStats] = Field(
+        default_factory=list,
+        description="A list of GPUs used by the worker process,"
+        " with detailed statistics.",
+    )
+
+
+@DeveloperAPI
+class TrainRunAttempt(BaseModel):
+    """Metadata for an individual attempt to execute a Train run."""
+
+    run_id: str = Field(description="Unique identifier for the parent Train run.")
+    attempt_id: str = Field(
+        description="Unique identifier for this specific Train run attempt."
+    )
+    status: RunAttemptStatus = Field(
+        description="The current execution status of the Train run attempt."
+    )
+    status_detail: Optional[str] = Field(
+        None,
+        description="Additional details about the status,"
+        " including error messages if applicable.",
+    )
+    start_time_ns: int = Field(
+        description="The UNIX timestamp (in nanoseconds)"
+        " when the Train run attempt started."
+    )
+    end_time_ns: Optional[int] = Field(
+        None,
+        description="The UNIX timestamp (in nanoseconds)"
+        " when the Train run attempt ended. "
+        "If null, the attempt is still ongoing.",
+    )
+    resources: List[TrainResources] = Field(
+        description="The resources (e.g., CPU, GPU) allocated to the Train run attempt."
+    )
+    workers: List[TrainWorker] = Field(
+        description="List of Train workers participating in this attempt, "
+        "sorted by global ranks."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainRunAttempt(TrainRunAttempt):
+    """Detailed metadata for a Train run attempt, including decorated worker data."""
+
+    workers: List[DecoratedTrainWorker] = Field(
+        description="A list of Train workers with detailed statistics, "
+        "sorted by global ranks."
+    )
+
+
+@DeveloperAPI
+class ExecutionOptions(BaseModel):
+    """ExecutionOptions for a single Ray Data ingest pipeline."""
+
+    resource_limits: Dict[str, Any] = Field(
+        description="The resource limits applied to the Ray Data execution plan."
+    )
+    exclude_resources: Dict[str, Any] = Field(
+        description="The resources excluded from the Ray Data execution plan "
+        "(e.g. resources reserved by Ray Train workers)."
+    )
+
+    @field_validator("resource_limits", "exclude_resources", mode="before")
+    @classmethod
+    def _sanitize_dict(cls, v):
+        return _to_json_serializable_value(v)
+
+    preserve_order: bool = Field(
+        description="Whether to preserve the order of outputs across operators."
+    )
+    actor_locality_enabled: bool = Field(
+        description="Whether actor-based locality optimizations are enabled."
+    )
+    verbose_progress: bool = Field(
+        description="Whether verbose progress reporting is enabled."
+    )
+
+
+@DeveloperAPI
+class DataExecutionOptions(BaseModel):
+    """ExecutionOptions for a Ray Train run, split into defaults and per-dataset overrides."""
+
+    default: ExecutionOptions = Field(
+        description="Execution options applied to any dataset without a per-dataset override."
+    )
+    per_dataset_execution_options: Dict[str, ExecutionOptions] = Field(
+        default_factory=dict,
+        description="Per-dataset execution option overrides, keyed by dataset name.",
+    )
+
+
+@DeveloperAPI
+class DataConfig(BaseModel):
+    """Configuration for dataset splitting and execution options within Ray Train."""
+
+    datasets_to_split: Union[Literal["all"], List[str]] = Field(
+        description="Which datasets to split; either 'all' or a list of dataset names."
+    )
+    execution_options: Optional[Dict] = Field(
+        default=None,
+        deprecated="DEPRECATED: Use data_execution_options instead.",
+    )
+    data_execution_options: DataExecutionOptions = Field(
+        description="Data execution options"
+    )
+    enable_shard_locality: bool = Field(
+        description="Whether to enable shard locality optimization."
+    )
+
+
+@DeveloperAPI
+class ScalingConfig(BaseModel):
+    """Scaling config for a Train run."""
+
+    num_workers: Union[int, Tuple[int, int]] = Field(
+        description="The number of workers for the Train run."
+    )
+    use_gpu: bool = Field(description="Whether to use GPUs for the Train run.")
+    resources_per_worker: Optional[Dict[str, float]] = Field(
+        None, description="The resources per worker for a Train run."
+    )
+    placement_strategy: str = Field(
+        description="The placement strategy for the Train run."
+    )
+    accelerator_type: Optional[str] = Field(
+        None, description="The accelerator type for the Train run."
+    )
+    use_tpu: bool = Field(description="Whether to use TPUs for the Train run.")
+    topology: Optional[str] = Field(None, description="The topology for the Train run.")
+    bundle_label_selector: Optional[
+        Union[Dict[str, str], List[Dict[str, str]]]
+    ] = Field(None, description="The bundle label selector for the Train run.")
+
+
+@DeveloperAPI
+class FailureConfig(BaseModel):
+    """Failure config for a Train run."""
+
+    max_failures: int = Field(
+        description="The maximum number of failures for a Train run."
+    )
+    controller_failure_limit: int = Field(
+        description="The maximum number of controller failures to tolerate."
+    )
+
+
+@DeveloperAPI
+class CheckpointConfig(BaseModel):
+    """Checkpoint config for a Train run."""
+
+    num_to_keep: Optional[int] = Field(
+        None,
+        description="The number of most recent checkpoints to keep. Older checkpoints may be deleted.",
+    )
+    checkpoint_score_attribute: Optional[str] = Field(
+        None,
+        description="Attribute used to score and rank checkpoints; can be a metric key or attribute.",
+    )
+    checkpoint_score_order: Literal["max", "min"] = Field(
+        description="Order to rank checkpoint scores, 'max' for higher-is-better, 'min' for lower-is-better.",
+    )
+
+
+@DeveloperAPI
+class RunConfig(BaseModel):
+    """Run configuration parameters for a Train run, encompassing failure,
+    runtime environment, checkpoint settings, and storage path."""
+
+    name: str = Field(description="The name of the Train run.")
+    failure_config: FailureConfig = Field(
+        description="The failure config for a Train run."
+    )
+    worker_runtime_env: Dict[str, Any] = Field(
+        description="The worker runtime env for a Train run."
+    )
+
+    @field_validator("worker_runtime_env", mode="before")
+    @classmethod
+    def _sanitize_worker_runtime_env(cls, v):
+        return _to_json_serializable_value(v)
+
+    checkpoint_config: CheckpointConfig = Field(
+        description="The checkpoint config for a Train run."
+    )
+    storage_path: str = Field(description="The storage path for a Train run.")
+    storage_filesystem: Optional[str] = Field(
+        None, description="The storage filesystem for a Train run."
+    )
+
+    @field_validator("storage_filesystem", mode="before")
+    @classmethod
+    def _sanitize_storage_filesystem(cls, v):
+        return _to_json_serializable_value(v)
+
+
+@DeveloperAPI
+class BackendConfig(BaseModel):
+    """Backend config for a Train run."""
+
+    framework: Optional[TrainingFramework] = Field(
+        None, description="The training framework for this backend config."
+    )
+    config: Dict[str, Any] = Field(
+        description="Training framework-specific configuration fields."
+    )
+
+    @field_validator("config", mode="before")
+    @classmethod
+    def _sanitize_config(cls, v):
+        return _to_json_serializable_value(v)
+
+
+@DeveloperAPI
+class RunSettings(BaseModel):
+    """Settings for a Train run, primarily consisting of configs set before a train run starts.
+
+    This includes the train loop config, backend config, scaling config, dataset configs,
+    and runtime configuration.
+    """
+
+    train_loop_config: Optional[Dict] = Field(
+        None, description="The user defined train loop config for a Train run."
+    )
+
+    @field_validator("train_loop_config", mode="before")
+    @classmethod
+    def _sanitize_train_loop_config(cls, v):
+        return _to_json_serializable_value(v)
+
+    backend_config: BackendConfig = Field(
+        description="The backend config for a Train run. Can vary with the framework (e.g. TorchConfig)"
+    )
+    scaling_config: ScalingConfig = Field(
+        description="The scaling config for this Train run."
+    )
+    datasets: List[str] = Field(
+        description="A list of dataset names for a Train run.",
+    )
+    data_config: DataConfig = Field(
+        description="The data config for a Train run.",
+    )
+    run_config: RunConfig = Field(
+        description="Run configuration for this Train run, including failure, runtime environment, checkpoint settings, and storage path."
+    )
+
+
+@DeveloperAPI
+class TrainRun(BaseModel):
+    """Metadata for a Ray Train run, including its details and status."""
+
+    id: str = Field(description="Unique identifier for the Train run.")
+    name: str = Field(description="Human-readable name assigned to the Train run.")
+    job_id: str = Field(description="The Ray Job ID associated with this Train run.")
+    controller_actor_id: str = Field(
+        description="Unique ID of the actor managing the Train run."
+    )
+    status: RunStatus = Field(
+        description="The current execution status of the Train run."
+    )
+    status_detail: Optional[str] = Field(
+        None,
+        description="Additional details about the current status, "
+        "including error messages if applicable.",
+    )
+    start_time_ns: int = Field(
+        description="The UNIX timestamp (in nanoseconds) when the Train run started."
+    )
+    end_time_ns: Optional[int] = Field(
+        None,
+        description="The UNIX timestamp (in nanoseconds) when the Train run ended. "
+        "If null, the run is still in progress.",
+    )
+    controller_log_file_path: Optional[str] = Field(
+        None, description="The path to the log file for the Train run controller."
+    )
+    framework_versions: Dict[str, str] = Field(
+        description="The relevant framework versions for this Train run,"
+        "including the Ray version and training framework version."
+    )
+    run_settings: RunSettings = Field(
+        description="The run settings for this Train run, including train loop config, "
+        "backend config, scaling config, dataset details, and runtime configuration."
+    )
+
+
+@DeveloperAPI
+class DecoratedTrainRun(TrainRun):
+    """Detailed metadata for a Ray Train run, including attempts and job details."""
+
+    attempts: List[DecoratedTrainRunAttempt] = Field(
+        description="A list of attempts made to execute the Train run."
+    )
+    job_details: Optional[JobDetails] = Field(
+        None,
+        description="Detailed information about the job that initiated this Train run.",
+    )
+
+
+@DeveloperAPI
+class TrainRunsResponse(BaseModel):
+    """Response containing a list of decorated Train runs."""
+
+    train_runs: List[DecoratedTrainRun] = Field(
+        description="A list of Train runs with detailed metadata."
+    )

python/ray/train/v2/_internal/state/state_actor.py

@@ -0,0 +1,301 @@
+import copy
+import logging
+import os
+import threading
+import time
+from collections import OrderedDict, defaultdict
+from typing import Dict, Optional
+
+import ray
+from ray._private import ray_constants
+from ray._private.event.export_event_logger import (
+    EventLogType,
+    check_export_api_enabled,
+    get_export_event_logger,
+)
+from ray.actor import ActorHandle
+from ray.train.v2._internal.constants import (
+    CONTROLLERS_TO_POLL_PER_ITERATION,
+    DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
+    DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
+    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+    GET_ACTOR_TIMEOUT_S,
+    STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
+)
+from ray.train.v2._internal.state.schema import (
+    TrainRun,
+    TrainRunAttempt,
+)
+from ray.train.v2._internal.state.util import (
+    is_actor_alive,
+    update_train_run_aborted,
+    update_train_run_attempt_aborted,
+)
+from ray.train.v2._internal.util import time_monotonic
+
+logger = logging.getLogger(__name__)
+
+
+class TrainStateActor:
+    def __init__(
+        self,
+        # TODO: group into single config if we need to do similar polling elsewhere
+        enable_state_actor_reconciliation: bool = False,
+        reconciliation_interval_s: float = 30,
+        get_actor_timeout_s: int = GET_ACTOR_TIMEOUT_S,
+        controllers_to_poll_per_iteration: int = CONTROLLERS_TO_POLL_PER_ITERATION,
+    ):
+        # NOTE: All runs and attempts are stored in memory.
+        # This may be a memory issue for large runs.
+        # TODO: consider cleaning up runs over time.
+        self._runs: Dict[str, TrainRun] = OrderedDict()
+        # {run_id: {attempt_id: TrainRunAttempt}}
+        self._run_attempts: Dict[str, OrderedDict[str, TrainRunAttempt]] = defaultdict(
+            OrderedDict
+        )
+        (
+            self._export_logger,
+            self._is_train_run_export_api_enabled,
+            self._is_train_run_attempt_export_api_enabled,
+        ) = self._init_export_logger()
+
+        # TODO: consider row level locking if loop takes too long.
+        self._runs_lock = threading.RLock()
+        self._run_attempts_lock = threading.RLock()
+
+        # Set env vars related to reconciling train run/attempt state.
+        if enable_state_actor_reconciliation:
+            self._reconciliation_interval_s = reconciliation_interval_s
+            self._controllers_to_poll_per_iteration = controllers_to_poll_per_iteration
+            self._get_actor_timeout_s = get_actor_timeout_s
+            self._start_run_state_reconciliation_thread()
+
+    def _abort_live_runs_with_dead_controllers(
+        self, last_poll_run_id: Optional[str]
+    ) -> str:
+        aborted_run_ids = []
+        with self._runs_lock:
+            runs = list(self._runs.values())
+
+            # Start iterating from poll index.
+            starting_poll_index = 0
+            if last_poll_run_id is not None:
+                for poll_index, run in enumerate(runs):
+                    if run.id == last_poll_run_id:
+                        starting_poll_index = (poll_index + 1) % len(runs)
+                        break
+
+            # Abort runs.
+            num_polled_runs = 0
+            poll_index = starting_poll_index
+            while (
+                poll_index < starting_poll_index + len(runs)
+                and num_polled_runs < self._controllers_to_poll_per_iteration
+            ):
+                run = runs[poll_index % len(runs)]
+                poll_index += 1
+                last_poll_run_id = run.id
+                if run.status.is_terminal():
+                    continue
+                try:
+                    if not is_actor_alive(
+                        run.controller_actor_id, self._get_actor_timeout_s
+                    ):
+                        update_train_run_aborted(run, False)
+                        self.create_or_update_train_run(run)
+                        aborted_run_ids.append(run.id)
+                except ray.util.state.exception.RayStateApiException:
+                    logger.exception(
+                        "State API unavailable when checking if actor is alive. "
+                        "Will check again on next poll."
+                    )
+                num_polled_runs += 1
+
+        # Abort run attempts.
+        with self._run_attempts_lock:
+            for run_id in aborted_run_ids:
+                latest_run_attempt = self._get_latest_run_attempt(run_id)
+                if latest_run_attempt and not latest_run_attempt.status.is_terminal():
+                    update_train_run_attempt_aborted(latest_run_attempt, False)
+                    self.create_or_update_train_run_attempt(latest_run_attempt)
+
+        return last_poll_run_id
+
+    def _start_run_state_reconciliation_thread(self) -> None:
+        def _reconciliation_loop():
+            last_poll_run_id = None
+            latest_poll_time = float("-inf")
+            while True:
+                # Wait for the poll interval to elapse.
+                time_since_last_poll = time_monotonic() - latest_poll_time
+                if time_since_last_poll < self._reconciliation_interval_s:
+                    remaining_time = (
+                        self._reconciliation_interval_s - time_since_last_poll
+                    )
+                    time.sleep(remaining_time)
+
+                last_poll_run_id = self._abort_live_runs_with_dead_controllers(
+                    last_poll_run_id
+                )
+                latest_poll_time = time_monotonic()
+
+        threading.Thread(target=_reconciliation_loop, daemon=True).start()
+
+    def _get_latest_run_attempt(self, run_id: str) -> Optional[TrainRunAttempt]:
+        with self._run_attempts_lock:
+            # NOTE: run_attempts is OrderedDict from attempt_id to TrainRunAttempt.
+            run_attempts = self._run_attempts.get(run_id, {})
+            if not run_attempts:
+                return None
+            return next(reversed(run_attempts.values()))
+
+    def create_or_update_train_run(self, run: TrainRun) -> None:
+        with self._runs_lock:
+            self._runs[run.id] = run
+            run_copy = copy.deepcopy(run)
+        self._maybe_export_train_run(run_copy)
+
+    def create_or_update_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
+        with self._run_attempts_lock:
+            self._run_attempts[run_attempt.run_id][run_attempt.attempt_id] = run_attempt
+            run_attempt_copy = copy.deepcopy(run_attempt)
+        self._maybe_export_train_run_attempt(run_attempt_copy)
+
+    def get_train_runs(self) -> Dict[str, TrainRun]:
+        with self._runs_lock:
+            return self._runs
+
+    def get_train_run_attempts(self) -> Dict[str, Dict[str, TrainRunAttempt]]:
+        with self._run_attempts_lock:
+            return self._run_attempts
+
+    # ============================
+    # Export API
+    # ============================
+
+    def is_export_api_enabled(self) -> bool:
+        return self._export_logger is not None
+
+    def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]:
+        """Initialize the export logger and check if the export API is enabled.
+
+        Returns:
+            A tuple containing:
+                - The export logger (or None if export API is not enabled).
+                - A boolean indicating if the export API is enabled for train runs.
+                - A boolean indicating if the export API is enabled for train run attempts.
+        """
+        # Proto schemas should be imported within the scope of TrainStateActor to
+        # prevent serialization errors.
+        from ray.core.generated.export_event_pb2 import ExportEvent
+
+        is_train_run_export_api_enabled = check_export_api_enabled(
+            ExportEvent.SourceType.EXPORT_TRAIN_RUN
+        )
+        is_train_run_attempt_export_api_enabled = check_export_api_enabled(
+            ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
+        )
+        export_api_enabled = (
+            is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled
+        )
+
+        if not export_api_enabled:
+            return None, False, False
+
+        log_directory = os.path.join(
+            ray._private.worker._global_node.get_session_dir_path(), "logs"
+        )
+        logger = None
+        try:
+            logger = get_export_event_logger(
+                EventLogType.TRAIN_STATE,
+                log_directory,
+            )
+        except Exception:
+            logger.exception(
+                "Unable to initialize the export event logger, so no Train export "
+                "events will be written."
+            )
+
+        if logger is None:
+            return None, False, False
+
+        return (
+            logger,
+            is_train_run_export_api_enabled,
+            is_train_run_attempt_export_api_enabled,
+        )
+
+    def _maybe_export_train_run(self, run: TrainRun) -> None:
+        if not self._is_train_run_export_api_enabled:
+            return
+
+        from ray.train.v2._internal.state.export import train_run_to_proto
+
+        run_proto = train_run_to_proto(run)
+        self._export_logger.send_event(run_proto)
+
+    def _maybe_export_train_run_attempt(self, run_attempt: TrainRunAttempt) -> None:
+        if not self._is_train_run_attempt_export_api_enabled:
+            return
+
+        from ray.train.v2._internal.state.export import train_run_attempt_to_proto
+
+        run_attempt_proto = train_run_attempt_to_proto(run_attempt)
+        self._export_logger.send_event(run_attempt_proto)
+
+
+TRAIN_STATE_ACTOR_NAME = "train_v2_state_actor"
+TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor"
+
+_state_actor_lock: threading.RLock = threading.RLock()
+
+
+def get_or_create_state_actor() -> ActorHandle:
+    """Get or create the Ray Train state actor singleton.
+
+    This is a long-living, detached actor living on the head node
+    that gets initialized when the first Train run happens on the
+    Ray cluster.
+    """
+    with _state_actor_lock:
+        state_actor = (
+            ray.remote(TrainStateActor)
+            .options(
+                num_cpus=0,
+                name=TRAIN_STATE_ACTOR_NAME,
+                namespace=TRAIN_STATE_ACTOR_NAMESPACE,
+                get_if_exists=True,
+                lifetime="detached",
+                resources={"node:__internal_head__": 0.001},
+                # Escape from the parent's placement group
+                scheduling_strategy="DEFAULT",
+                max_restarts=-1,
+                max_task_retries=-1,
+            )
+            .remote(
+                enable_state_actor_reconciliation=ray_constants.env_bool(
+                    ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
+                    DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION,
+                ),
+                reconciliation_interval_s=float(
+                    os.getenv(
+                        STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
+                        DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S,
+                    )
+                ),
+            )
+        )
+
+    return state_actor
+
+
+def get_state_actor() -> Optional[ActorHandle]:
+    """Get the `TrainStateActor` if exists, otherwise return None."""
+    try:
+        return ray.get_actor(
+            name=TRAIN_STATE_ACTOR_NAME,
+            namespace=TRAIN_STATE_ACTOR_NAMESPACE,
+        )
+    except ValueError:
+        return None

python/ray/train/v2/_internal/state/state_manager.py

@@ -0,0 +1,326 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import ray
+from ray.actor import ActorHandle
+from ray.train import BackendConfig
+from ray.train._internal.data_config import DataConfig
+from ray.train.v2._internal.execution.context import DistributedContext
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    ResizeDecision,
+)
+from ray.train.v2._internal.execution.worker_group import ActorMetadata, Worker
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    BackendConfig as BackendConfigSchema,
+    CheckpointConfig as CheckpointConfigSchema,
+    FailureConfig as FailureConfigSchema,
+    RunAttemptStatus,
+    RunConfig as RunConfigSchema,
+    RunSettings,
+    RunStatus,
+    ScalingConfig as ScalingConfigSchema,
+    TrainResources,
+    TrainRun,
+    TrainRunAttempt,
+    TrainWorker,
+)
+from ray.train.v2._internal.state.state_actor import get_or_create_state_actor
+from ray.train.v2._internal.state.util import (
+    construct_data_config,
+    current_time_ns,
+    mark_workers_dead,
+    update_train_run_aborted,
+    update_train_run_attempt_aborted,
+)
+from ray.train.v2._internal.util import TrainingFramework
+from ray.train.v2.api.config import RunConfig, ScalingConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainStateManager:
+    """Manages the state of a train run and run attempts."""
+
+    def __init__(self) -> None:
+        self._state_actor = get_or_create_state_actor()
+        # NOTE: All runs and attempts are stored in memory.
+        # This may be a memory issue for large runs.
+        self._runs: Dict[str, TrainRun] = {}
+        # {run_id: {attempt_id: TrainRunAttempt}}
+        self._run_attempts: Dict[str, Dict[str, TrainRunAttempt]] = defaultdict(dict)
+
+    def create_train_run(
+        self,
+        id: str,
+        name: str,
+        job_id: str,
+        controller_actor_id: str,
+        controller_log_file_path: str,
+        run_config: RunConfig,
+        train_loop_config: Optional[Dict],
+        scaling_config: ScalingConfig,
+        backend_config: BackendConfig,
+        datasets: Dict[str, ray.data.Dataset],
+        dataset_config: DataConfig,
+    ) -> None:
+        run_config_schema = RunConfigSchema(
+            name=run_config.name,
+            failure_config=FailureConfigSchema(
+                max_failures=run_config.failure_config.max_failures,
+                controller_failure_limit=run_config.failure_config.controller_failure_limit,
+            ),
+            worker_runtime_env=run_config.worker_runtime_env,
+            checkpoint_config=CheckpointConfigSchema(
+                num_to_keep=run_config.checkpoint_config.num_to_keep,
+                checkpoint_score_attribute=run_config.checkpoint_config.checkpoint_score_attribute,
+                checkpoint_score_order=run_config.checkpoint_config.checkpoint_score_order,
+            ),
+            storage_path=run_config.storage_path,
+            storage_filesystem=run_config.storage_filesystem,
+        )
+
+        scaling_config_schema = ScalingConfigSchema(
+            num_workers=scaling_config.num_workers,
+            use_gpu=scaling_config.use_gpu,
+            resources_per_worker=scaling_config.resources_per_worker,
+            placement_strategy=scaling_config.placement_strategy,
+            accelerator_type=scaling_config.accelerator_type,
+            use_tpu=scaling_config.use_tpu,
+            topology=scaling_config.topology,
+            bundle_label_selector=scaling_config.label_selector,
+        )
+
+        backend_config_schema = BackendConfigSchema(
+            framework=backend_config.framework,
+            config=backend_config.to_dict(),
+        )
+
+        run_settings = RunSettings(
+            train_loop_config=train_loop_config,
+            backend_config=backend_config_schema,
+            scaling_config=scaling_config_schema,
+            datasets=list(datasets.keys()),
+            data_config=construct_data_config(dataset_config),
+            run_config=run_config_schema,
+        )
+
+        run = TrainRun(
+            id=id,
+            name=name,
+            job_id=job_id,
+            status=RunStatus.INITIALIZING,
+            status_detail=None,
+            controller_actor_id=controller_actor_id,
+            start_time_ns=current_time_ns(),
+            end_time_ns=None,
+            controller_log_file_path=controller_log_file_path,
+            framework_versions={"ray": ray.__version__},
+            run_settings=run_settings,
+        )
+        self._runs[run.id] = run
+        # Block so the initial run state isn't lost if the controller exits
+        # right after. Without this, the .remote() task could still be in the
+        # caller's outbound queue when the controller dies, leaving the state
+        # actor with no record of the run.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_scheduling(
+        self,
+        run_id: str,
+        resize_decision: Optional[ResizeDecision] = None,
+    ) -> None:
+        if resize_decision is not None:
+            status_detail = _get_scheduling_status_detail(
+                resize_decision.num_workers, resize_decision.resources_per_worker
+            )
+        else:
+            status_detail = None
+
+        run = self._runs[run_id]
+        run.status = RunStatus.SCHEDULING
+        run.status_detail = status_detail
+        self._create_or_update_train_run(run)
+
+    def update_train_run_running(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RUNNING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_restarting(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RESTARTING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_resizing(
+        self,
+        run_id: str,
+    ) -> None:
+        run = self._runs[run_id]
+        run.status = RunStatus.RESIZING
+        run.status_detail = None
+        self._create_or_update_train_run(run)
+
+    def update_train_run_finished(
+        self,
+        run_id: str,
+    ):
+        run = self._runs[run_id]
+        run.status = RunStatus.FINISHED
+        run.status_detail = None
+        run.end_time_ns = current_time_ns()
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_errored(
+        self,
+        run_id: str,
+        status_detail: str,
+    ):
+        run = self._runs[run_id]
+        run.status = RunStatus.ERRORED
+        run.status_detail = status_detail
+        run.end_time_ns = current_time_ns()
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_aborted(
+        self,
+        run_id: str,
+    ):
+        run = self._runs[run_id]
+        update_train_run_aborted(run=run, graceful=True)
+        # Block on terminal status so the final state isn't lost if the controller exits right after.
+        self._create_or_update_train_run(run, block=True)
+
+    def update_train_run_framework_versions(
+        self, run_id: str, framework_versions: Dict[str, str]
+    ):
+        run = self._runs[run_id]
+        run.framework_versions = framework_versions
+        self._create_or_update_train_run(run)
+
+    def create_train_run_attempt(
+        self,
+        run_id: str,
+        attempt_id: str,
+        num_workers: int,
+        resources_per_worker: Dict[str, float],
+    ) -> None:
+        status_detail = _get_scheduling_status_detail(num_workers, resources_per_worker)
+        resources = [
+            TrainResources(resources=resources_per_worker) for _ in range(num_workers)
+        ]
+        run_attempt = TrainRunAttempt(
+            run_id=run_id,
+            attempt_id=attempt_id,
+            start_time_ns=current_time_ns(),
+            status=RunAttemptStatus.PENDING,
+            status_detail=status_detail,
+            resources=resources,
+            workers=[],  # Not started yet.
+        )
+
+        self._run_attempts[run_id][attempt_id] = run_attempt
+        self._create_or_update_train_run_attempt(run_attempt)
+
+    def update_train_run_attempt_running(
+        self, run_id: str, attempt_id: str, workers: List[Worker]
+    ) -> None:
+        def _convert_worker(worker: Worker) -> TrainWorker:
+
+            actor: ActorHandle = worker.actor
+            distributed_context: DistributedContext = worker.distributed_context
+            actor_metadata: ActorMetadata = worker.metadata
+
+            return TrainWorker(
+                world_rank=distributed_context.world_rank,
+                local_rank=distributed_context.local_rank,
+                node_rank=distributed_context.node_rank,
+                actor_id=actor._actor_id.hex(),
+                node_id=actor_metadata.node_id,
+                node_ip=actor_metadata.node_ip,
+                pid=actor_metadata.pid,
+                gpu_ids=actor_metadata.gpu_ids,
+                status=ActorStatus.ALIVE,
+                resources=TrainResources(resources=worker.resources),
+                log_file_path=worker.log_file_path,
+            )
+
+        workers: List[TrainWorker] = [_convert_worker(worker) for worker in workers]
+
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.RUNNING
+        run_attempt.status_detail = None
+        run_attempt.workers = workers
+        self._create_or_update_train_run_attempt(run_attempt)
+
+    def update_train_run_attempt_finished(
+        self,
+        run_id: str,
+        attempt_id: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.FINISHED
+        run_attempt.status_detail = None
+        run_attempt.end_time_ns = current_time_ns()
+        mark_workers_dead(run_attempt)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def update_train_run_attempt_errored(
+        self,
+        run_id: str,
+        attempt_id: str,
+        status_detail: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        run_attempt.status = RunAttemptStatus.ERRORED
+        run_attempt.status_detail = status_detail
+        run_attempt.end_time_ns = current_time_ns()
+        mark_workers_dead(run_attempt)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def update_train_run_attempt_aborted(
+        self,
+        run_id: str,
+        attempt_id: str,
+    ):
+        run_attempt = self._run_attempts[run_id][attempt_id]
+        update_train_run_attempt_aborted(run_attempt=run_attempt, graceful=True)
+        # Block to avoid case where controller is dead but attempt is not terminal.
+        self._create_or_update_train_run_attempt(run_attempt, block=True)
+
+    def get_train_run_framework(self, run_id: str) -> Optional[TrainingFramework]:
+        run = self._runs[run_id]
+        return run.run_settings.backend_config.framework
+
+    def _create_or_update_train_run(
+        self, run: TrainRun, *, block: bool = False
+    ) -> None:
+        ref = self._state_actor.create_or_update_train_run.remote(run)
+        if block:
+            ray.get(ref)
+
+    def _create_or_update_train_run_attempt(
+        self, run_attempt: TrainRunAttempt, *, block: bool = False
+    ) -> None:
+        ref = self._state_actor.create_or_update_train_run_attempt.remote(run_attempt)
+        if block:
+            ray.get(ref)
+
+
+def _get_scheduling_status_detail(
+    num_workers: int, resources_per_worker: Dict[str, float]
+) -> str:
+    return f"Scheduling {num_workers} workers, each requiring: {resources_per_worker}."

python/ray/train/v2/_internal/state/util.py

@@ -0,0 +1,97 @@
+import time
+
+from ray.data._internal.execution.interfaces.execution_options import ExecutionOptions
+from ray.train._internal.data_config import DataConfig
+from ray.train.v2._internal.state.schema import (
+    ActorStatus,
+    DataConfig as DataConfigSchema,
+    DataExecutionOptions,
+    ExecutionOptions as ExecutionOptionsSchema,
+    RunAttemptStatus,
+    RunStatus,
+    TrainRun,
+    TrainRunAttempt,
+)
+from ray.util.state import get_actor
+
+_GRACEFUL_ABORT_STATUS_DETAIL = "Run aborted due to user interrupt (SIGINT)."
+_DEAD_CONTROLLER_ABORT_STATUS_DETAIL = (
+    "Run aborted because the driver process exited unexpectedly."
+)
+
+
+def update_train_run_aborted(run: TrainRun, graceful: bool) -> None:
+    run.status = RunStatus.ABORTED
+    if graceful:
+        run.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
+    else:
+        run.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
+    run.end_time_ns = current_time_ns()
+
+
+def update_train_run_attempt_aborted(
+    run_attempt: TrainRunAttempt, graceful: bool
+) -> None:
+    if graceful:
+        run_attempt.status_detail = _GRACEFUL_ABORT_STATUS_DETAIL
+    else:
+        run_attempt.status_detail = _DEAD_CONTROLLER_ABORT_STATUS_DETAIL
+    run_attempt.status = RunAttemptStatus.ABORTED
+    run_attempt.end_time_ns = current_time_ns()
+    mark_workers_dead(run_attempt)
+
+
+def mark_workers_dead(run_attempt: TrainRunAttempt) -> None:
+    for worker in run_attempt.workers:
+        worker.status = ActorStatus.DEAD
+
+
+def current_time_ns() -> int:
+    return time.time_ns()
+
+
+def is_actor_alive(actor_id: str, timeout: int) -> bool:
+    """Returns whether actor is alive."""
+    actor_state = get_actor(actor_id, timeout=timeout)
+    return actor_state and actor_state.state != "DEAD"
+
+
+def construct_data_config(data_config: DataConfig) -> DataConfigSchema:
+    """Serialize a user-facing DataConfig into the exportable schema.
+
+    Note: This function assumes data_config._execution_options (a defaultdict)
+    hasn't been read between initialization of the field and this function call.
+    Any read materializes a dataset key and affects the data config shape,
+    wrongly capturing a per dataset execution options even if the user only
+    provided a default.
+    """
+    exec_options = data_config._execution_options
+
+    per_dataset_execution_options = {}
+    if exec_options:
+        per_dataset_execution_options = {
+            ds_name: execution_options_to_model(opts)
+            for ds_name, opts in exec_options.items()
+        }
+
+    return DataConfigSchema(
+        datasets_to_split=data_config._datasets_to_split,
+        data_execution_options=DataExecutionOptions(
+            default=execution_options_to_model(exec_options.default_factory()),
+            per_dataset_execution_options=per_dataset_execution_options,
+        ),
+        enable_shard_locality=data_config._enable_shard_locality,
+    )
+
+
+def execution_options_to_model(
+    execution_options: ExecutionOptions,
+) -> ExecutionOptionsSchema:
+    """Convert a ray.data ExecutionOptions object into the export schema model."""
+    return ExecutionOptionsSchema(
+        resource_limits=execution_options.resource_limits.to_resource_dict(),
+        exclude_resources=execution_options.exclude_resources.to_resource_dict(),
+        preserve_order=execution_options.preserve_order,
+        actor_locality_enabled=execution_options.actor_locality_enabled,
+        verbose_progress=execution_options.verbose_progress,
+    )