wehub-resource-sync
222 lines
7.9 KiB
Python
222 lines
7.9 KiB
Python
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,
|
|
)
|