chore: import upstream snapshot with attribution

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

@@ -0,0 +1,221 @@
+import importlib
+import logging
+from typing import TYPE_CHECKING, Dict, Optional
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+import ray
+from ray.train.v2._internal.execution.callback import (
+    ControllerCallback,
+    WorkerGroupCallback,
+)
+from ray.train.v2._internal.execution.context import TrainRunContext
+from ray.train.v2._internal.execution.controller.state import (
+    AbortedState,
+    ErroredState,
+    FinishedState,
+    ReschedulingState,
+    ResizingState,
+    RestartingState,
+    RunningState,
+    SchedulingState,
+    ShuttingDownState,
+    TrainControllerState,
+)
+from ray.train.v2._internal.execution.scaling_policy.scaling_policy import (
+    ResizeDecision,
+)
+from ray.train.v2._internal.execution.worker_group import (
+    WorkerGroup,
+    WorkerGroupContext,
+    WorkerGroupState,
+)
+from ray.train.v2._internal.execution.worker_group.poll import WorkerGroupPollStatus
+from ray.train.v2._internal.logging.logging import (
+    get_train_application_controller_log_path,
+)
+from ray.train.v2._internal.state.state_manager import TrainStateManager
+from ray.train.v2._internal.util import TrainingFramework
+
+logger = logging.getLogger(__name__)
+
+
+def _get_framework_version(framework: Optional[TrainingFramework]):
+    versions = {}
+
+    try:
+        import ray
+
+        versions["ray"] = ray.__version__
+    except ImportError:
+        logger.warning("Failed to collect ray version on worker.")
+
+    if framework is None:
+        return versions
+
+    for module_name in framework.module_names():
+        try:
+            module = importlib.import_module(module_name)
+            versions[module_name] = module.__version__
+        except ModuleNotFoundError:
+            # Module is not installed, skip without recording a version.
+            continue
+        except Exception:
+            logger.warning(f"Failed to collect {module_name} version on worker.")
+            continue
+
+    return versions
+
+
+class StateManagerCallback(ControllerCallback, WorkerGroupCallback):
+    def __init__(self, datasets: Dict[str, "Dataset"]):
+        self._datasets = datasets
+
+    def after_controller_start(self, train_run_context: TrainRunContext):
+        self._state_manager = TrainStateManager()
+        self._run_name = train_run_context.get_run_config().name
+        self._run_id = train_run_context.run_id
+
+        # TODO: Should this be generated by the caller?
+        # NOTE: These must be called on the Controller.
+        #       The Callback is first initialized on the Driver.
+        core_context = ray.runtime_context.get_runtime_context()
+        self._job_id = core_context.get_job_id()
+        self._controller_actor_id = core_context.get_actor_id()
+        controller_log_file_path = get_train_application_controller_log_path()
+        self._state_manager.create_train_run(
+            id=self._run_id,
+            name=self._run_name,
+            job_id=self._job_id,
+            controller_actor_id=self._controller_actor_id,
+            controller_log_file_path=controller_log_file_path,
+            run_config=train_run_context.run_config,
+            train_loop_config=train_run_context.train_loop_config,
+            scaling_config=train_run_context.scaling_config,
+            backend_config=train_run_context.backend_config,
+            datasets=self._datasets,
+            dataset_config=train_run_context.dataset_config,
+        )
+
+    def after_controller_state_update(
+        self,
+        previous_state: TrainControllerState,
+        current_state: TrainControllerState,
+    ):
+        if previous_state._state_type == current_state._state_type:
+            return
+
+        logger.info(
+            f"[State Transition] {previous_state._state_type.state_name} -> "
+            f"{current_state._state_type.state_name}."
+        )
+
+        if isinstance(current_state, SchedulingState):
+            # TODO: This should probably always be ResizeDecision.
+            if isinstance(current_state.scaling_decision, ResizeDecision):
+                resize_decision = current_state.scaling_decision
+            else:
+                resize_decision = None
+
+            self._state_manager.update_train_run_scheduling(
+                run_id=self._run_id,
+                resize_decision=resize_decision,
+            )
+
+        elif isinstance(current_state, RunningState):
+            self._state_manager.update_train_run_running(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, RestartingState):
+            self._state_manager.update_train_run_restarting(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ResizingState):
+            self._state_manager.update_train_run_resizing(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ErroredState):
+            self._state_manager.update_train_run_errored(
+                run_id=self._run_id,
+                status_detail=str(current_state.training_failed_error),
+            )
+
+        elif isinstance(current_state, FinishedState):
+            self._state_manager.update_train_run_finished(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, AbortedState):
+            self._state_manager.update_train_run_aborted(
+                run_id=self._run_id,
+            )
+
+        elif isinstance(current_state, ReschedulingState):
+            # substate of SchedulingState
+            pass
+
+        elif isinstance(current_state, ShuttingDownState):
+            # substate of RunningState
+            pass
+
+    def before_worker_group_start(self, worker_group_context: WorkerGroupContext):
+        self._state_manager.create_train_run_attempt(
+            run_id=self._run_id,
+            attempt_id=worker_group_context.run_attempt_id,
+            num_workers=worker_group_context.num_workers,
+            resources_per_worker=worker_group_context.resources_per_worker,
+        )
+
+    def after_worker_group_start(self, worker_group: WorkerGroup):
+        worker_group_context: WorkerGroupContext = (
+            worker_group.get_worker_group_context()
+        )
+        worker_group_state: WorkerGroupState = worker_group.get_worker_group_state()
+        self._state_manager.update_train_run_attempt_running(
+            run_id=self._run_id,
+            attempt_id=worker_group_context.run_attempt_id,
+            workers=worker_group_state.workers,
+        )
+
+        # Update train run framework version
+        framework = self._state_manager.get_train_run_framework(self._run_id)
+        framework_versions = worker_group.execute_single(
+            0, _get_framework_version, framework
+        )
+        self._state_manager.update_train_run_framework_versions(
+            run_id=self._run_id,
+            framework_versions=framework_versions,
+        )
+
+    def before_worker_group_shutdown(self, worker_group: WorkerGroup):
+        worker_group_context: WorkerGroupContext = (
+            worker_group.get_worker_group_context()
+        )
+        # TODO: Consider passing error reason directly to the callback.
+        # Something along the lines of:
+        #    WorkerGroup.shutdown(reason)
+        #    -> WorkerGroupCallback.before_worker_group_shutdown(reason)
+        worker_group_poll_status: Optional[
+            WorkerGroupPollStatus
+        ] = worker_group.get_latest_poll_status()
+        if worker_group_poll_status and worker_group_poll_status.errors:
+            self._state_manager.update_train_run_attempt_errored(
+                run_id=self._run_id,
+                attempt_id=worker_group_context.run_attempt_id,
+                status_detail=worker_group_poll_status.get_error_string(),
+            )
+        else:
+            self._state_manager.update_train_run_attempt_finished(
+                run_id=self._run_id,
+                attempt_id=worker_group_context.run_attempt_id,
+            )
+
+    def before_worker_group_abort(self, worker_group_context: WorkerGroupContext):
+        self._state_manager.update_train_run_attempt_aborted(
+            self._run_id,
+            worker_group_context.run_attempt_id,
+        )