diff --git a/trinity/cli/launcher.py b/trinity/cli/launcher.py index 7b0977803ba..aaed86a0385 100644 --- a/trinity/cli/launcher.py +++ b/trinity/cli/launcher.py @@ -1,4 +1,5 @@ """Launch the trainer""" +import contextlib import os import time import traceback @@ -234,86 +235,92 @@ def both(config: Config) -> StageStatus: from trinity.explorer.explorer import Explorer from trinity.manager.synchronizer import Synchronizer from trinity.trainer.trainer import Trainer + from trinity.utils.monitor import SharedRun + + # A shared run (one backend run shared by explorer + trainer) is opt-in. When enabled, + # the orchestration runs inside SharedRun, which creates the shared run on entry and + # tears it down on exit; when disabled, it runs unwrapped — same behavior as before. + shared_run_ctx = SharedRun(config) if config.monitor.shared_run else contextlib.nullcontext() + with shared_run_ctx: + explorer = Explorer.get_actor(config) + trainer = Trainer.get_actor(config) + started_at = time.perf_counter() + try: + ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()]) + ray.get( + [ + explorer.prepare.remote(), + trainer.prepare.remote(), + ] + ) + # Set up NCCL weight sync group between Trainer and Explorer. + # This must happen after both sides are prepared (Trainer has model + # meta cached, Explorer has rollout models created) and before the + # first weight sync. + if config.synchronizer.sync_method == SyncMethod.NCCL: + synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace) + ray.get(synchronizer.coordinate_weight_sync_setup.remote()) + ray.get( + [ + explorer.sync_weight.remote(), + trainer.sync_weight.remote(), + ] + ) + ready_ref, wait_ref = ray.wait( + [ + explorer.explore.remote(), + trainer.train.remote(), + ], + num_returns=1, + ) - explorer = Explorer.get_actor(config) - trainer = Trainer.get_actor(config) - started_at = time.perf_counter() - try: - ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()]) - ray.get( - [ - explorer.prepare.remote(), - trainer.prepare.remote(), - ] - ) - # Set up NCCL weight sync group between Trainer and Explorer. - # This must happen after both sides are prepared (Trainer has model - # meta cached, Explorer has rollout models created) and before the - # first weight sync. - if config.synchronizer.sync_method == SyncMethod.NCCL: - synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace) - ray.get(synchronizer.coordinate_weight_sync_setup.remote()) - ray.get( - [ - explorer.sync_weight.remote(), - trainer.sync_weight.remote(), - ] - ) - ready_ref, wait_ref = ray.wait( - [ - explorer.explore.remote(), - trainer.train.remote(), - ], - num_returns=1, - ) - - ready = ray.get(ready_ref[0]) - if ready == config.trainer.name: - logger.info( - "===========================================================\n" - "> Launcher detected that the `Trainer` process has finished.\n" - "> Stopping the explorer process immediately.\n" - "===========================================================" + ready = ray.get(ready_ref[0]) + if ready == config.trainer.name: + logger.info( + "===========================================================\n" + "> Launcher detected that the `Trainer` process has finished.\n" + "> Stopping the explorer process immediately.\n" + "===========================================================" + ) + ray.wait(wait_ref, timeout=5) + elif ready == config.explorer.name: + logger.info( + "===============================================================\n" + "> Launcher detected that the `Explorer` process has finished.\n" + "> `Trainer` process may need to save the model checkpoint.\n" + f"> Waiting {config.synchronizer.sync_timeout} s for the trainer process...\n" + "> You can force stop the `Trainer` process by pressing Ctrl+C.\n" + "===============================================================" + ) + ray.wait(wait_ref, timeout=config.synchronizer.sync_timeout) + return StageStatus( + stage="both", + success=True, + total_time_sec=time.perf_counter() - started_at, ) - ray.wait(wait_ref, timeout=5) - elif ready == config.explorer.name: - logger.info( - "===============================================================\n" - "> Launcher detected that the `Explorer` process has finished.\n" - "> `Trainer` process may need to save the model checkpoint.\n" - f"> Waiting {config.synchronizer.sync_timeout} s for the trainer process...\n" - "> You can force stop the `Trainer` process by pressing Ctrl+C.\n" - "===============================================================" + except Exception as exc: + error = _build_stage_error(exc) + logger.error(f"Explorer or Trainer failed:\n{error.traceback_text}") + return StageStatus( + stage="both", + success=False, + total_time_sec=time.perf_counter() - started_at, + error=error, + ) + finally: + # Tear down the NCCL weight sync group before shutting down actors. + # Best-effort: if actors or Synchronizer are already dead, skip. + if config.synchronizer.sync_method == SyncMethod.NCCL: + try: + synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace) + ray.get(synchronizer.coordinate_weight_sync_teardown.remote(), timeout=30) + except Exception: + logger.warning("Weight sync teardown skipped (actors may have already exited).") + ray.wait( + [explorer.shutdown.remote(), trainer.shutdown.remote()], + timeout=config.synchronizer.sync_timeout, + num_returns=2, ) - ray.wait(wait_ref, timeout=config.synchronizer.sync_timeout) - return StageStatus( - stage="both", - success=True, - total_time_sec=time.perf_counter() - started_at, - ) - except Exception as exc: - error = _build_stage_error(exc) - logger.error(f"Explorer or Trainer failed:\n{error.traceback_text}") - return StageStatus( - stage="both", - success=False, - total_time_sec=time.perf_counter() - started_at, - error=error, - ) - finally: - # Tear down the NCCL weight sync group before shutting down actors. - # Best-effort: if actors or Synchronizer are already dead, skip. - if config.synchronizer.sync_method == SyncMethod.NCCL: - try: - synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace) - ray.get(synchronizer.coordinate_weight_sync_teardown.remote(), timeout=30) - except Exception: - logger.warning("Weight sync teardown skipped (actors may have already exited).") - ray.wait( - [explorer.shutdown.remote(), trainer.shutdown.remote()], - timeout=config.synchronizer.sync_timeout, - num_returns=2, - ) MODE_MAP = { diff --git a/trinity/common/config.py b/trinity/common/config.py index 2d231490c55..4e668f11206 100644 --- a/trinity/common/config.py +++ b/trinity/common/config.py @@ -857,6 +857,11 @@ class MonitorConfig: # whether to enable ray timeline profile # the output file will be saved to `cache_dir/timeline.json` enable_ray_timeline: bool = False + # whether to merge explorer & trainer into one shared run (both mode, wandb only; + # validator downgrades to False for backends that don't support it) + shared_run: bool = False + # ! DO NOT SET, the primary run id the launcher injects for actors to join + run_id: Optional[str] = None # ! DO NOT SET, automatically generated as checkpoint_job_dir/monitor cache_dir: str = "" diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index 07f0374ac37..81649a6971e 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -926,6 +926,14 @@ def validate(self, config: Config) -> None: monitor_cls = MONITOR.get(config.monitor.monitor_type) if monitor_cls is None: raise ValueError(f"Invalid monitor type: {config.monitor.monitor_type}") + # shared_run requires a backend that can merge explorer+trainer into one run; + # downgrade to separate runs (instead of erroring) for backends that can't. + if config.monitor.shared_run and not monitor_cls.supports_shared_run: + self.logger.warning( + f"monitor `{config.monitor.monitor_type}` does not support shared_run; " + "disabling it (explorer and trainer will use separate runs)." + ) + config.monitor.shared_run = False set_if_none(config.monitor, "monitor_args", monitor_cls.default_args()) # create a job dir in ////monitor config.monitor.cache_dir = os.path.join(config.checkpoint_job_dir, "monitor") diff --git a/trinity/utils/monitor.py b/trinity/utils/monitor.py index 7b693716b84..0313f3de388 100644 --- a/trinity/utils/monitor.py +++ b/trinity/utils/monitor.py @@ -1,8 +1,9 @@ """Monitor""" +import json import os from abc import ABC, abstractmethod -from typing import Dict +from typing import Dict, Literal, Optional import numpy as np import pandas as pd @@ -38,10 +39,15 @@ }, ) +_logger = get_logger(__name__) + class Monitor(ABC): """Monitor""" + # whether this backend can merge explorer+trainer into one shared run + supports_shared_run: bool = False + def __init__( self, project: str, @@ -74,6 +80,92 @@ def default_args(cls) -> Dict: """Return default arguments for the monitor.""" return {} + # ---- shared-run hooks (launcher/primary side) ------------------------------- + # Default implementations make a backend "shared-run unaware": init_shared declines + # (-> separate runs), poll_shared is a no-op. Backends opt in by overriding. + + @classmethod + def init_shared(cls, config: Config) -> Optional[object]: + """Create the primary shared run; default declines (unsupported backend). + + Args: + config: Global config; the run id is written back to ``config.monitor.run_id``. + + Returns: + The primary run handle, or ``None`` if the backend can't share runs. + """ + _logger.warning( + "monitor `%s` does not support shared_run; falling back to separate runs.", + config.monitor.monitor_type, + ) + return None + + @classmethod + def poll_shared(cls, handle, config: Config, declared: set) -> None: + """Re-declare newly-seen bindings on the primary; default no-op. + + Args: + handle: The primary run. + config: Global config. + declared: Keys already declared on the primary (mutated to add new ones). + """ + + @staticmethod + def finish_shared(handle) -> None: + """Finish the primary shared run, best effort. + + Args: + handle: The primary run to finish, or ``None`` to skip. + """ + if handle is None: + return + try: + handle.finish() + except Exception: + _logger.warning("Failed to finish shared run", exc_info=True) + + @staticmethod + def load_shared_step_bindings(cache_dir: str) -> Dict[str, str]: + """Merge all actors' binding files into one mapping. + + Args: + cache_dir: Monitor cache dir holding ``shared_step_bindings/*.json``. + + Returns: + Merged ``metric key -> step axis`` mapping. + """ + bindings: Dict[str, str] = {} + if not cache_dir: + return bindings + d = os.path.join(cache_dir, "shared_step_bindings") + if not os.path.isdir(d): + return bindings + for fn in os.listdir(d): + if not fn.endswith(".json"): + continue + try: + with open(os.path.join(d, fn)) as f: + bindings.update(json.load(f)) + except Exception: + _logger.warning("Failed to read bindings file %s", fn, exc_info=True) + return bindings + + @staticmethod + def apply_shared_step_bindings(handle, bindings: Dict[str, str]) -> None: + """Declare each ``metric -> step axis`` binding on the primary run. + + Args: + handle: The primary run. + bindings: ``metric key -> step axis`` mapping to declare. + """ + if handle is None or not bindings: + return + for key, step_metric in bindings.items(): + try: + handle.define_metric(key, step_metric=step_metric) + except Exception: + _logger.warning("Failed to bind %s -> %s", key, step_metric, exc_info=True) + def format_data_str(self, data: dict, step: int) -> str: cleaned_data = { k: ( @@ -124,26 +216,121 @@ class WandbMonitor(Monitor): api_key (`Optional[str]`): The API key for W&B. If not provided, use the environment variable `WANDB_API_KEY`. """ - def __init__( - self, project: str, group: str, name: str, role: str, config: Config = None - ) -> None: - assert wandb is not None, "wandb is not installed. Please install it to use WandbMonitor." - if not group: - group = name - monitor_args = config.monitor.monitor_args or {} + supports_shared_run: bool = True + + @staticmethod + def _apply_credentials(monitor_args: dict) -> None: + """Export ``base_url``/``api_key`` from monitor_args into wandb env vars. + + Args: + monitor_args: Monitor args optionally holding ``base_url``/``api_key``. + """ if base_url := monitor_args.get("base_url"): os.environ["WANDB_BASE_URL"] = base_url if api_key := monitor_args.get("api_key"): os.environ["WANDB_API_KEY"] = api_key - self.logger = wandb.init( - project=project, - group=group, - name=f"{name}_{role}", - tags=[role], + + @classmethod + def init_shared(cls, config: Config) -> Optional[object]: + """Create the primary run (x_primary) and inject ``config.monitor.run_id``. + + Args: + config: Global config; the run id is written back to ``config.monitor.run_id``. + + Returns: + The primary run, or ``None`` in offline mode (no server-side merge). + """ + assert wandb is not None, "wandb is not installed. Please install it to use WandbMonitor." + cls._apply_credentials(config.monitor.monitor_args or {}) + if os.environ.get("WANDB_MODE") == "offline": + _logger.warning("wandb offline mode does not support shared runs; using separate runs.") + return None + run = wandb.init( + project=config.project, + group=config.group or config.name, + name=config.name, config=config, save_code=False, + settings=wandb.Settings(mode="shared", x_primary=True), ) + config.monitor.run_id = run.id + cls._define_shared_step_axes(run, config) + return run + + @staticmethod + def _define_shared_step_axes(run, config: Config) -> None: + """Declare each role's hidden ``{role}/step`` axis on the primary run. + + Args: + run: The primary run. + config: Global config providing the explorer/trainer role names. + """ + # hidden: each axis is an x-axis only, with no plot of its own + for role in (config.explorer.name, config.trainer.name): + run.define_metric(f"{role}/step", hidden=True) + + @classmethod + def poll_shared(cls, handle, config: Config, declared: set) -> None: + """Declare any binding keys not yet in ``declared`` on the primary. + + Args: + handle: The primary run. + config: Global config (provides ``cache_dir``). + declared: Keys already declared; newly-seen keys are added in place. + """ + bindings = cls.load_shared_step_bindings(config.monitor.cache_dir) + new = {k: v for k, v in bindings.items() if k not in declared} + if not new: + return + # define_metric is retroactive: already-written history rows realign once declared + cls.apply_shared_step_bindings(handle, new) + declared.update(new) + + def __init__( + self, project: str, group: str, name: str, role: str, config: Config = None + ) -> None: + assert wandb is not None, "wandb is not installed. Please install it to use WandbMonitor." + if not group: + group = name self.console_logger = get_logger(__name__, in_ray_actor=True) + # init these before wandb.init so a failed init still leaves close()/__del__ safe + # (Option A: shared_writer is created only after a successful shared init below) + self.logger = None + self.shared_writer: Optional[SharedRunWriter] = None + self._apply_credentials(config.monitor.monitor_args or {}) + # join the launcher-created primary run as a secondary writer when enabled + self.shared = bool(config.monitor.shared_run and config.monitor.run_id) + if self.shared: + self.logger = wandb.init( + project=project, + group=group, + id=config.monitor.run_id, + config=config, + resume="allow", + reinit=True, + save_code=False, + settings=wandb.Settings( + mode="shared", + x_primary=False, + x_update_finish_state=False, # this writer's close must not finish the run + ), + ) + self.shared_writer = SharedRunWriter( + self.logger, + role, + getattr(config.monitor, "cache_dir", "") or "", + self.console_logger, + config, + ) + else: + self.logger = wandb.init( + project=project, + group=group, + name=f"{name}_{role}", + tags=[role], + config=config, + save_code=False, + ) def log_table(self, table_name: str, experiences_table: pd.DataFrame, step: int): experiences_table = wandb.Table(dataframe=experiences_table) @@ -151,11 +338,17 @@ def log_table(self, table_name: str, experiences_table: pd.DataFrame, step: int) def log(self, data: dict, step: int, commit: bool = False) -> None: """Log metrics.""" - self.logger.log(data, step=step, commit=commit) + if self.shared: + self.shared_writer.record(data, step) + else: + self.logger.log(data, step=step, commit=commit) self.console_logger.info(f"{self.format_data_str(data, step)}") def close(self) -> None: - self.logger.finish() + if self.shared_writer is not None: + self.shared_writer.flush() + if self.logger is not None: + self.logger.finish() @classmethod def default_args(cls) -> Dict: @@ -307,3 +500,155 @@ def close(self) -> None: def default_args(cls) -> Dict: """Return default arguments for the monitor.""" return {} + + +class SharedRun: + """Launcher-side context manager that creates and finishes a shared run. + + A no-op unless ``config.monitor.shared_run`` is on and the backend supports it. + The secondaries keep the live step-axis bindings aligned (each carries the full set), + so the launcher only declares them once on the primary at exit -- which is what makes + the x-axis correct after finalize (the primary's config wins when the run finalizes). + """ + + def __init__(self, config: Config) -> None: + """Resolve the backend (no run created yet). + + Args: + config: Global config selecting the monitor backend. + """ + self.config = config + self.handle = None + self._cls = MONITOR.get(config.monitor.monitor_type) + self._declared: set = set() + + def __enter__(self) -> "SharedRun": + """Create the primary run. + + Returns: + Self, for use as a context manager. + """ + if not self.config.monitor.shared_run or self._cls is None: + return self + self.handle = self._cls.init_shared(self.config) # create primary, inject run_id + return self + + def __exit__(self, *exc) -> Literal[False]: + """Declare the bindings once on the primary, then finish the run. + + Args: + *exc: Exception triple from the ``with`` block (unused). + + Returns: + ``False``, so exceptions from the with-body are never suppressed. + """ + if self.handle is not None: + try: + # one-shot: declare all bindings on the primary before finalize + self._cls.poll_shared(self.handle, self.config, self._declared) + except Exception: + _logger.warning("shared-run poll failed", exc_info=True) + self._cls.finish_shared(self.handle) + return False # never suppress exceptions from the with-body + + +class SharedRunWriter: + """Secondary-side writer: buffers one row per step and persists step-axis bindings. + + Created only after a successful ``wandb.init`` (Option A), so ``run`` is always valid. + """ + + def __init__(self, run, role: str, cache_dir: str, console_logger, config: Config) -> None: + """Bind the shared run and init the per-step buffer / binding state. + + Args: + run: The joined secondary wandb run. + role: This actor's role, used as the ``{role}/`` metric prefix and step axis. + cache_dir: Monitor cache dir; ``""`` disables binding persistence. + console_logger: Logger for best-effort persistence warnings. + config: Global config, for the explorer/trainer role names. + """ + self._run = run + self.role = role + self._console_logger = console_logger + self._cache_dir = cache_dir + # buffer one step's metrics, emitted as a single row when the step advances + self._buffer: dict = {} + self._buffer_step = None + # full keys seen so far, so the primary binds each to the step axis exactly once + self._logged_keys: set = set() + # keys already define_metric'd on THIS secondary handle (they persist in its config) + self._declared: set = set() + # cache file the launcher reads at finish; None disables persistence + self._bindings_path = ( + os.path.join(cache_dir, "shared_step_bindings", f"{role}.json") if cache_dir else None + ) + self._persisted_count = 0 + # Declare BOTH roles' hidden step axes on this secondary so its config matches the + # primary's and its syncs don't drop the merged step-axis defs. + for r in (config.explorer.name, config.trainer.name): + self._run.define_metric(f"{r}/step", hidden=True) + + def record(self, data: dict, step: int) -> None: + """Buffer a step's metrics, flushing the previous step when ``step`` advances. + + Args: + data: Metric name -> value for this step. + step: This role's step counter. + """ + if self._buffer_step is not None and step != self._buffer_step: + self.flush() + self._buffer_step = step + self._buffer.update({f"{self.role}/{k}": v for k, v in data.items()}) + + def flush(self) -> None: + """Emit the buffered step as one wandb history row.""" + if self._buffer_step is None: + return + self._logged_keys.update(self._buffer.keys()) + # Declare the FULL binding set on this secondary handle, not just this role's keys: + # in wandb shared mode a writer's config sync replaces the merged metric defs, so a + # secondary holding only its own role's bindings drops the other role's on every + # sync -- the two secondaries clobber each other and only one role's x-axis is + # correct at a time. Every role's published bindings + this step's own keys make + # this secondary carry the whole set, so no sync drops anything (verified: both + # roles stay bound throughout the run, no flicker). + full = dict(Monitor.load_shared_step_bindings(self._cache_dir)) + full.update({k: f"{self.role}/step" for k in self._buffer}) + for key, step_metric in full.items(): + if key not in self._declared: + try: + self._run.define_metric(key, step_metric=step_metric) + except Exception: + self._console_logger.warning("Failed to bind %s", key, exc_info=True) + self._declared.add(key) + # inject the axis value only at log time so it never enters the buffer + self._run.log({**self._buffer, f"{self.role}/step": self._buffer_step}) + self._buffer = {} + self._buffer_step = None + self._persist_bindings() + + def step_bindings(self) -> Dict[str, str]: + """Return this role's bindings for the primary to declare. + + Returns: + ``metric key -> {role}/step`` for every key this role has logged. + """ + # include any buffered-but-not-yet-flushed keys so none are missed + self._logged_keys.update(self._buffer.keys()) + return {key: f"{self.role}/step" for key in self._logged_keys} + + def _persist_bindings(self) -> None: + """Atomically write the bindings to the cache file for the launcher.""" + # nothing to do without a cache path, or when no new keys appeared since last write + if not self._bindings_path or len(self._logged_keys) == self._persisted_count: + return + try: + os.makedirs(os.path.dirname(self._bindings_path), exist_ok=True) + tmp = f"{self._bindings_path}.tmp" + with open(tmp, "w") as f: + json.dump(self.step_bindings(), f) + os.replace(tmp, self._bindings_path) + self._persisted_count = len(self._logged_keys) + except Exception: + self._console_logger.warning("Failed to persist step bindings", exc_info=True)