wehub-resource-sync
153 lines
4.9 KiB
Python
153 lines
4.9 KiB
Python
import logging
|
|
import os
|
|
import threading
|
|
from typing import Dict, Optional
|
|
|
|
import ray
|
|
from ray._private.event.export_event_logger import (
|
|
EventLogType,
|
|
check_export_api_enabled,
|
|
get_export_event_logger,
|
|
)
|
|
from ray.actor import ActorHandle
|
|
from ray.train._internal.state.schema import TrainRunInfo
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class TrainStateActor:
|
|
def __init__(self):
|
|
self._run_infos: Dict[str, TrainRunInfo] = {}
|
|
|
|
(
|
|
self._export_logger,
|
|
self._is_train_run_export_api_enabled,
|
|
self._is_train_run_attempt_export_api_enabled,
|
|
) = self._init_export_logger()
|
|
|
|
def register_train_run(self, run_info: TrainRunInfo) -> None:
|
|
# Register a new train run.
|
|
self._run_infos[run_info.id] = run_info
|
|
|
|
self._maybe_export_train_run(run_info)
|
|
self._maybe_export_train_run_attempt(run_info)
|
|
|
|
def get_train_run(self, run_id: str) -> Optional[TrainRunInfo]:
|
|
# Retrieve a registered run with its id
|
|
return self._run_infos.get(run_id, None)
|
|
|
|
def get_all_train_runs(self) -> Dict[str, TrainRunInfo]:
|
|
# Retrieve all registered train runs
|
|
return self._run_infos
|
|
|
|
# ============================
|
|
# Export API
|
|
# ============================
|
|
|
|
def is_export_api_enabled(self) -> bool:
|
|
return self._export_logger is not None
|
|
|
|
def _init_export_logger(self) -> tuple[Optional[logging.Logger], bool, bool]:
|
|
"""Initialize the export logger and check if the export API is enabled.
|
|
|
|
Returns:
|
|
A tuple containing:
|
|
- The export logger (or None if export API is not enabled).
|
|
- A boolean indicating if the export API is enabled for train runs.
|
|
- A boolean indicating if the export API is enabled for train run attempts.
|
|
"""
|
|
# Proto schemas should be imported within the scope of TrainStateActor to
|
|
# prevent serialization errors.
|
|
from ray.core.generated.export_event_pb2 import ExportEvent
|
|
|
|
is_train_run_export_api_enabled = check_export_api_enabled(
|
|
ExportEvent.SourceType.EXPORT_TRAIN_RUN
|
|
)
|
|
is_train_run_attempt_export_api_enabled = check_export_api_enabled(
|
|
ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
|
|
)
|
|
export_api_enabled = (
|
|
is_train_run_export_api_enabled or is_train_run_attempt_export_api_enabled
|
|
)
|
|
|
|
if not export_api_enabled:
|
|
return None, False, False
|
|
|
|
log_directory = os.path.join(
|
|
ray._private.worker._global_node.get_session_dir_path(), "logs"
|
|
)
|
|
logger = None
|
|
try:
|
|
logger = get_export_event_logger(
|
|
EventLogType.TRAIN_STATE,
|
|
log_directory,
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"Unable to initialize the export event logger, so no Train export "
|
|
"events will be written."
|
|
)
|
|
|
|
if logger is None:
|
|
return None, False, False
|
|
|
|
return (
|
|
logger,
|
|
is_train_run_export_api_enabled,
|
|
is_train_run_attempt_export_api_enabled,
|
|
)
|
|
|
|
def _maybe_export_train_run(self, run_info: TrainRunInfo) -> None:
|
|
if not self._is_train_run_export_api_enabled:
|
|
return
|
|
|
|
from ray.train._internal.state.export import train_run_info_to_proto_run
|
|
|
|
run_proto = train_run_info_to_proto_run(run_info)
|
|
self._export_logger.send_event(run_proto)
|
|
|
|
def _maybe_export_train_run_attempt(self, run_info: TrainRunInfo) -> None:
|
|
if not self._is_train_run_attempt_export_api_enabled:
|
|
return
|
|
|
|
from ray.train._internal.state.export import train_run_info_to_proto_attempt
|
|
|
|
run_attempt_proto = train_run_info_to_proto_attempt(run_info)
|
|
self._export_logger.send_event(run_attempt_proto)
|
|
|
|
|
|
TRAIN_STATE_ACTOR_NAME = "train_state_actor"
|
|
TRAIN_STATE_ACTOR_NAMESPACE = "_train_state_actor"
|
|
|
|
_state_actor_lock: threading.RLock = threading.RLock()
|
|
|
|
|
|
def get_or_create_state_actor() -> ActorHandle:
|
|
"""Get or create a `TrainStateActor` on the head node."""
|
|
with _state_actor_lock:
|
|
state_actor = TrainStateActor.options(
|
|
name=TRAIN_STATE_ACTOR_NAME,
|
|
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
|
|
get_if_exists=True,
|
|
lifetime="detached",
|
|
resources={"node:__internal_head__": 0.001},
|
|
# Escape from the parent's placement group
|
|
scheduling_strategy="DEFAULT",
|
|
).remote()
|
|
|
|
# Ensure the state actor is ready
|
|
ray.get(state_actor.__ray_ready__.remote())
|
|
return state_actor
|
|
|
|
|
|
def get_state_actor() -> Optional[ActorHandle]:
|
|
"""Get the `TrainStateActor` if exists, otherwise return None."""
|
|
try:
|
|
return ray.get_actor(
|
|
name=TRAIN_STATE_ACTOR_NAME,
|
|
namespace=TRAIN_STATE_ACTOR_NAMESPACE,
|
|
)
|
|
except ValueError:
|
|
return None
|