chore: import upstream snapshot with attribution

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

@@ -0,0 +1,132 @@
+import logging
+import os
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, List
+
+import ray
+from ray.train._internal.state.schema import (
+    ActorStatusEnum,
+    RunStatusEnum,
+    TrainDatasetInfo,
+    TrainRunInfo,
+    TrainWorkerInfo,
+)
+from ray.train._internal.utils import check_for_failure
+from ray.train._internal.worker_group import WorkerGroup
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+class TrainRunStateManager:
+    """A class that aggregates and reports train run info to TrainStateActor.
+
+    This manager class is created on the train controller layer for each run.
+    """
+
+    def __init__(self, state_actor) -> None:
+        self.state_actor = state_actor
+        self.train_run_info_dict = defaultdict(dict)
+
+    def register_train_run(
+        self,
+        run_id: str,
+        job_id: str,
+        run_name: str,
+        run_status: str,
+        controller_actor_id: str,
+        datasets: Dict[str, "Dataset"],
+        worker_group: WorkerGroup,
+        start_time_ms: float,
+        resources: List[Dict[str, float]],
+        status_detail: str = "",
+    ) -> None:
+        """Collect Train Run Info and report to StateActor."""
+
+        if not self.state_actor:
+            logger.warning(
+                "Unable to register train run since `TrainStateActor` is not started."
+            )
+            return
+
+        def collect_train_worker_info():
+            train_context = ray.train.get_context()
+            core_context = ray.runtime_context.get_runtime_context()
+            return TrainWorkerInfo(
+                world_rank=train_context.get_world_rank(),
+                local_rank=train_context.get_local_rank(),
+                node_rank=train_context.get_node_rank(),
+                actor_id=core_context.get_actor_id(),
+                node_id=core_context.get_node_id(),
+                node_ip=ray.util.get_node_ip_address(),
+                gpu_ids=ray.get_gpu_ids(),
+                pid=os.getpid(),
+                resources=resources[0],
+                status=ActorStatusEnum.ALIVE,
+            )
+
+        futures = [
+            worker_group.execute_single_async(index, collect_train_worker_info)
+            for index in range(len(worker_group))
+        ]
+        success, exception = check_for_failure(futures)
+
+        if not success:
+            logger.error(
+                "Failed to collect run information from the Ray Train "
+                f"workers:\n{exception}"
+            )
+            return
+
+        worker_info_list = ray.get(futures)
+        worker_info_list = sorted(worker_info_list, key=lambda info: info.world_rank)
+
+        dataset_info_list = [
+            TrainDatasetInfo(
+                name=ds_name,
+                dataset_name=ds._dataset_name,
+                dataset_uuid=ds._uuid,
+            )
+            for ds_name, ds in datasets.items()
+        ]
+
+        updates = dict(
+            id=run_id,
+            job_id=job_id,
+            name=run_name,
+            controller_actor_id=controller_actor_id,
+            workers=worker_info_list,
+            datasets=dataset_info_list,
+            start_time_ms=start_time_ms,
+            run_status=run_status,
+            status_detail=status_detail,
+            resources=resources,
+        )
+
+        # Clear the cached info to avoid registering the same run twice
+        self.train_run_info_dict[run_id] = {}
+        self._update_train_run_info(run_id, updates)
+
+    def end_train_run(
+        self,
+        run_id: str,
+        run_status: RunStatusEnum,
+        status_detail: str,
+        end_time_ms: int,
+    ):
+        """Update the train run status when the training is finished."""
+        updates = dict(
+            run_status=run_status,
+            status_detail=status_detail,
+            end_time_ms=end_time_ms,
+        )
+        self._update_train_run_info(run_id, updates)
+
+    def _update_train_run_info(self, run_id: str, updates: Dict[str, Any]) -> None:
+        """Update specific fields of a registered TrainRunInfo instance."""
+        if run_id in self.train_run_info_dict:
+            self.train_run_info_dict[run_id].update(updates)
+            train_run_info = TrainRunInfo(**self.train_run_info_dict[run_id])
+            ray.get(self.state_actor.register_train_run.remote(train_run_info))