wehub-resource-sync
133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
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))
|