Skip to content

feat(monitor): unify explorer/trainer wandb logging into one run #590

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MengsD wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(monitor): unify explorer/trainer wandb logging into one run #590

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 84 additions & 77 deletions trinity/cli/launcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Launch the trainer"""
import contextlib
import os
import time
import traceback
Expand Down Expand Up @@ -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 = {
Expand Down
5 changes: 5 additions & 0 deletions trinity/common/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ""

Expand Down
8 changes: 8 additions & 0 deletions trinity/common/config_validator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <checkpoint_root_dir>/<project>/<group>/<name>/monitor
config.monitor.cache_dir = os.path.join(config.checkpoint_job_dir, "monitor")
Expand Down
Loading
Loading