chore: import upstream snapshot with attribution

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,
+    )