[Train][Observability] Track Train Run Info with TrainStateActor

python/ray/train/_internal/backend_executor.py

@@ -4,6 +4,8 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar
 
+from python.ray.train._internal.stats import TrainRunStatsManager
+
 import ray
 import ray._private.ray_constants as ray_constants
 from ray._private.ray_constants import env_integer
@@ -23,6 +25,7 @@
 from ray.train.backend import BackendConfig
 from ray.train.constants import (
     ENABLE_DETAILED_AUTOFILLED_METRICS_ENV,
+    ENABLE_RAY_TRAIN_DASHBOARD_ENV,
     ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
     ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
     TRAIN_ENABLE_WORKER_SPREAD_ENV,
@@ -118,6 +121,8 @@ def __init__(
             )
         ]
 
+        self.dashboard_enabled = env_integer(ENABLE_RAY_TRAIN_DASHBOARD_ENV, 0)
+
     def start(
         self,
         initialization_hook: Optional[Callable[[], None]] = None,
@@ -194,6 +199,10 @@ def _set_driver_dataset_context(ctx: DataContext):
             self._increment_failures()
             self._restart()
 
+        # Setup StatsActorManager for Ray Train Dashboard
+        if self.dashboard_enabled:
+            self.stats_manager = TrainRunStatsManager()
+
     def _create_placement_group(self):
         """Creates a placement group if it does not exist.
 
@@ -432,7 +441,6 @@ def start_training(
         data_config: DataConfig,
         storage: StorageContext,
         checkpoint: Optional[Checkpoint] = None,
-        on_session_init: Callable[[], None] = None,
     ) -> None:
         """Executes a training function on all workers in a separate thread.
 
@@ -528,8 +536,19 @@ def initialize_session(
 
         self.get_with_failure_handling(futures)
 
-        if on_session_init:
-            on_session_init()
+        # Register Train Run before training starts
+        if self.dashboard_enabled:
+            session = get_session()
+            trainer_actor_id = ray.runtime_context.get_runtime_context().get_actor_id()
+
+            self.stats_manager.register_train_run(
+                run_id=session.run_id,
+                run_name=session.experiment_name,
+                trial_name=session.trial_name,
+                trainer_actor_id=trainer_actor_id,
+                datasets=datasets,
+                worker_group=self.worker_group,
+            )
 
         # Run the training function asynchronously in its own thread.
         def train_async():

python/ray/train/_internal/schema.py

@@ -0,0 +1,63 @@
+from typing import List, Optional
+
+from ray._private.pydantic_compat import BaseModel, Field
+from ray.util.annotations import DeveloperAPI
+
+try:
+    import pydantic  # noqa: F401
+except ImportError:
+    raise ModuleNotFoundError(
+        "pydantic isn't installed. "
+        "To install pydantic, please run 'pip install pydantic'"
+    )
+
+
+@DeveloperAPI
+class TrainWorkerInfo(BaseModel):
+    """Metadata of a Ray Train worker."""
+
+    actor_id: str = Field(description="Actor ID of the worker.")
+    world_rank: int = Field(description="World rank.")
+    local_rank: int = Field(description="Local rank.")
+    node_rank: int = Field(description="Node rank.")
+    gpu_ids: Optional[List[str]] = Field(
+        description="A list of GPU ids allocated to that worker."
+    )
+    node_id: Optional[str] = Field(
+        description="ID of the node that the worker is running on."
+    )
+    node_ip: Optional[str] = Field(
+        description="IP address of the node that the worker is running on."
+    )
+    pid: Optional[str] = Field(description="PID of the worker.")
+
+
+@DeveloperAPI
+class TrainDatasetInfo(BaseModel):
+    name: str = Field(
+        description="The key of the dataset dict specified in Ray Train Trainer."
+    )
+    plan_name: str = Field(description="The name of the internal dataset plan.")
+    plan_uuid: str = Field(description="The uuid of the internal dataset plan.")
+
+
+@DeveloperAPI
+class TrainRunInfo(BaseModel):
+    """Metadata for a Ray Train run and information about its workers."""
+
+    name: str = Field(description="The name of the Train run.")
+    id: str = Field(description="The unique identifier for each Train run.")
+    job_id: str = Field(description="Ray Job ID.")
+    trial_name: str = Field(
+        description=(
+            "Trial name. It should be different among different Train runs, "
+            "except for those that are restored from checkpoints."
+        )
+    )
+    trainer_actor_id: str = Field(description="Actor Id of the Trainer.")
+    workers: List[TrainWorkerInfo] = Field(
+        description="A List of Train workers sorted by global ranks."
+    )
+    datasets: List[TrainDatasetInfo] = Field(
+        description="A List of dataset info for this Train run."
+    )

python/ray/train/_internal/session.py

@@ -9,7 +9,7 @@
 import warnings
 from dataclasses import dataclass
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Type
 
 import ray
 from ray.air._internal.session import _get_session
@@ -60,7 +60,9 @@ class TrialInfo:
     resources: Dict[str, float]
     logdir: str
     driver_ip: str
+    run_id: str
     experiment_name: Optional[str] = None
+    datasets_info: Optional[List[Dict[str, Any]]] = None
 
 
 class _FutureTrainingResult:
@@ -461,6 +463,10 @@ def trial_name(self) -> str:
     def trial_id(self) -> str:
         return self.trial_info.id
 
+    @property
+    def run_id(self) -> str:
+        return self.trial_info.run_id
+
     @property
     def trial_resources(self) -> "PlacementGroupFactory":
         return self.trial_info.resources
@@ -478,7 +484,8 @@ def get_dataset_shard(
             warnings.warn(
                 "No dataset passed in. Returning None. Make sure to "
                 "pass in a Dataset to Trainer.run to use this "
-                "function."
+                "function.",
+                stacklevel=2,
             )
         elif isinstance(shard, dict):
             if not dataset_name:
@@ -646,7 +653,8 @@ def wrapper(*args, **kwargs):
                     warnings.warn(
                         f"`{fn_name}` is meant to only be "
                         "called inside a function that is executed by a Tuner"
-                        f" or Trainer. Returning `{default_value}`."
+                        f" or Trainer. Returning `{default_value}`.",
+                        stacklevel=2,
                     )
                 return default_value
             return fn(*args, **kwargs)
@@ -818,6 +826,13 @@ def get_trial_id() -> str:
     return _get_session().trial_id
 
 
+@PublicAPI(stability="alpha")
+@_warn_session_misuse()
+def get_run_id() -> str:
+    """Unique Train Run id for the corresponding trial."""
+    return _get_session().run_id
+
+
 @PublicAPI(stability="beta")
 @_warn_session_misuse()
 def get_trial_resources() -> "PlacementGroupFactory":

python/ray/train/_internal/stats.py

@@ -0,0 +1,117 @@
+import logging
+import os
+import threading
+from typing import Dict, Optional
+
+import ray
+from ray.data import Dataset
+from ray.train._internal.schema import TrainDatasetInfo, TrainRunInfo, TrainWorkerInfo
+from ray.train._internal.utils import check_for_failure
+from ray.train._internal.worker_group import WorkerGroup
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_cpus=0)
+class TrainStatsActor:
+    def __init__(self):
+        self.train_runs = dict()
+
+    def register_train_run(self, run_info: TrainRunInfo):
+        # Register a new train run.
+        self.train_runs[run_info.id] = run_info
+
+    def get_train_run(self, run_id: str) -> Optional[TrainRunInfo]:
+        # Retrieve a registered run with its id
+        return self.train_runs.get(run_id, None)
+
+    def get_all_train_runs(self) -> Dict[str, TrainRunInfo]:
+        # Retrieve all registered train runs
+        return self.train_runs
+
+
+TRAIN_STATS_ACTOR_NAME = "train_stats_actor"
+TRAIN_STATS_ACTOR_NAMESPACE = "_train_stats_actor"
+
+_stats_actor_lock: threading.RLock = threading.RLock()
+
+
+def get_or_launch_stats_actor():
+    """Create or launch a `TrainStatsActor` on the head node."""
+    with _stats_actor_lock:
+        return TrainStatsActor.options(
+            name=TRAIN_STATS_ACTOR_NAME,
+            namespace=TRAIN_STATS_ACTOR_NAMESPACE,
+            get_if_exists=True,
+            lifetime="detached",
+            resources={"node:__internal_head__": 0.001},
+        ).remote()
+
+
+class TrainRunStatsManager:
+    """A class that aggregates and reports train run info to TrainStatsActor.
+
+    This manager class is created on the train controller layer for each run.
+    """
+
+    def __init__(self) -> None:
+        self.stats_actor = get_or_launch_stats_actor()
+
+    def register_train_run(
+        self,
+        run_id: str,
+        run_name: str,
+        trial_name: str,
+        trainer_actor_id: str,
+        datasets: Dict[str, Dataset],
+        worker_group: WorkerGroup,
+    ) -> None:
+        """Collect Train Run Info and report to StatsActor."""
+
+        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=core_context.get_node_ip_address(),
+                gpu_ids=core_context.get_accelerator_ids().get("GPU", []),
+                pid=os.getpid(),
+            )
+
+        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.warning("Failed to collect infomation for Ray Train Worker.")
+            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,
+                plan_name=ds._plan._dataset_name,
+                plan_uuid=ds._plan._dataset_uuid,
+            )
+            for ds_name, ds in datasets.items()
+        ]
+
+        train_run_info = TrainRunInfo(
+            id=run_id,
+            name=run_name,
+            trial_name=trial_name,
+            trainer_actor_id=trainer_actor_id,
+            workers=worker_info_list,
+            datasets=dataset_info_list,
+        )
+
+        self.stats_actor.register_train_run.remote(train_run_info)

python/ray/train/base_trainer.py

@@ -361,7 +361,8 @@ def training_loop(self):
                 f"Invalid trainer type. You are attempting to restore a trainer of type"
                 f" {trainer_cls} with `{cls.__name__}.restore`, "
                 "which will most likely fail. "
-                f"Use `{trainer_cls.__name__}.restore` instead."
+                f"Use `{trainer_cls.__name__}.restore` instead.",
+                stacklevel=2,
             )
 
         original_datasets = param_dict.pop("datasets", {})

python/ray/train/constants.py

@@ -78,6 +78,10 @@ def _get_ray_train_session_dir() -> str:
 # Defaults to 0, which always retries on node preemption failures.
 RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE = "RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE"
 
+# Set this to 1 to start a StatsActor and collect infomation for Ray Train Dashboard
+# Defaults to 0
+ENABLE_RAY_TRAIN_DASHBOARD_ENV = "ENABLE_RAY_TRAIN_DASHBOARD_ENV"
+
 # NOTE: When adding a new environment variable, please track it in this list.
 TRAIN_ENV_VARS = {
     ENABLE_DETAILED_AUTOFILLED_METRICS_ENV,
@@ -87,6 +91,7 @@ def _get_ray_train_session_dir() -> str:
     TRAIN_ENABLE_WORKER_SPREAD_ENV,
     RAY_CHDIR_TO_TRIAL_DIR,
     RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE,
+    ENABLE_RAY_TRAIN_DASHBOARD_ENV,
 }
 
 # Key for AIR Checkpoint metadata in TrainingResult metadata

python/ray/train/data_parallel_trainer.py

@@ -431,10 +431,12 @@ def training_loop(self) -> None:
         trial_info = TrialInfo(
             name=session.get_trial_name(),
             id=session.get_trial_id(),
+            run_id=session.get_run_id(),
             resources=session.get_trial_resources(),
             logdir=session.get_trial_dir(),
             driver_ip=ray.util.get_node_ip_address(),
             experiment_name=session.get_experiment_name(),
+            datasets_info=self.datasets_info,
         )
 
         backend_executor = self._backend_executor_cls(

python/ray/tune/trainable/function_trainable.py

@@ -1,6 +1,7 @@
 import inspect
 import logging
 import os
+import uuid
 from functools import partial
 from numbers import Number
 from typing import Any, Callable, Dict, Optional, Type
@@ -58,6 +59,7 @@ def setup(self, config):
                 logdir=self._storage.trial_driver_staging_path,
                 driver_ip=None,
                 experiment_name=self._storage.experiment_dir_name,
+                run_id=uuid.uuid4().hex,
             ),
             storage=self._storage,
             synchronous_result_reporting=True,