launcher.py

"""Launch the trainer"""
import asyncio
import os
import sys
import traceback
from pprint import pprint
from typing import Optional

import ray
import typer
from typing_extensions import Annotated

from trinity.common.config import Config, load_config
from trinity.common.constants import DEBUG_NAMESPACE, PLUGIN_DIRS_ENV_VAR
from trinity.manager.checkpoint_converter import Converter
from trinity.manager.state_manager import StateManager
from trinity.utils.dlc_utils import is_running, setup_ray_cluster, stop_ray_cluster
from trinity.utils.log import get_logger
from trinity.utils.plugin_loader import load_plugins

logger = get_logger(__name__)

app = typer.Typer(
    help="Trinity CLI - Launch and manage Trinity-RFT processes.",
    pretty_exceptions_show_locals=False,
    pretty_exceptions_short=True,
)


def bench(config: Config) -> None:
    """Evaluate model."""
    from trinity.explorer.explorer import Explorer

    config.explorer.name = "benchmark"
    try:
        explorer = Explorer.get_actor(config)
        ray.get(explorer.prepare.remote())
        ray.get(explorer.benchmark.remote())
        logger.info("Benchmark finished.")
        ray.get(explorer.shutdown.remote())
    except Exception:
        logger.error(f"Benchmark failed:\n{traceback.format_exc()}")


def explore(config: Config) -> None:
    """Run explorer."""
    from trinity.explorer.explorer import Explorer

    try:
        explorer = Explorer.get_actor(config)
        ray.get(explorer.prepare.remote())
        ray.get(explorer.sync_weight.remote())
        ray.get(explorer.explore.remote())
        ray.get(explorer.shutdown.remote())
    except Exception:
        logger.error(f"Explorer failed:\n{traceback.format_exc()}")


def train(config: Config) -> None:
    """Run trainer."""
    from trinity.trainer.trainer import Trainer

    try:
        trainer = Trainer.get_actor(config)
        ray.get(trainer.prepare.remote())
        ray.get(trainer.sync_weight.remote())
        ray.get(trainer.train.remote())
        ray.get(trainer.shutdown.remote())
    except Exception:
        logger.error(f"Trainer failed:\n{traceback.format_exc()}")


def serve(config: Config) -> None:
    """Run explorer in server mode."""
    from trinity.explorer.explorer import Explorer

    try:
        explorer = Explorer.get_actor(config)
        ray.get(explorer.prepare.remote())
        ray.get(explorer.sync_weight.remote())
        ray.get(explorer.serve.remote())
        ray.get(explorer.shutdown.remote())
    except Exception:
        logger.error(f"Explorer failed:\n{traceback.format_exc()}")


def both(config: Config) -> None:
    """Setup both explorer and trainer.

    For the explorer, a step contains `batch_size * sync_interval` number
    of rollout tasks.

    For the trainer, it has to consume all experiences generated by the explorer in
    the latest step. The specific number of experiences may vary for different
    algorithms and tasks.
    """
    from trinity.explorer.explorer import Explorer
    from trinity.trainer.trainer import Trainer

    try:
        explorer = Explorer.get_actor(config)
        trainer = Trainer.get_actor(config)
        ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()])
        ray.get(
            [
                explorer.prepare.remote(),
                trainer.prepare.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"
                "==========================================================="
            )
            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)
        ray.wait(
            [explorer.shutdown.remote(), trainer.shutdown.remote()],
            timeout=config.synchronizer.sync_timeout,
            num_returns=2,
        )
    except Exception:
        logger.error(f"Explorer or Trainer failed:\n{traceback.format_exc()}")


MODE_MAP = {
    "explore": explore,
    "train": train,
    "both": both,
    "bench": bench,
    "serve": serve,
    "colocate": both,
}


def run_stage(config: Config) -> None:
    ray.init(
        address=config.cluster.ray_address,
        ignore_reinit_error=True,
        namespace=config.ray_namespace,
        runtime_env={"env_vars": config.get_envs()},
    )
    pprint(config)
    try:
        from trinity.buffer.pipelines.task_pipeline import check_and_run_task_pipeline

        check_and_run_task_pipeline(config)
        MODE_MAP[config.mode](config)
    finally:
        if config.monitor.enable_ray_timeline:
            timeline_file = os.path.join(config.monitor.cache_dir, "timeline.json")
            logger.info(f"Exporting Ray timeline to {timeline_file}...")
            ray.timeline(filename=timeline_file)
            logger.info("Done. You can open the timeline file in `chrome://tracing`")
        ray.shutdown()


# ---------------------------------------------------------------------------
# CLI commands
# ---------------------------------------------------------------------------


@app.command()
def run(
    config: Annotated[
        str,
        typer.Option("--config", help="Path to the config file."),
    ],
    dlc: Annotated[
        bool,
        typer.Option("--dlc", help="Specify when running in Aliyun PAI DLC."),
    ] = False,
    plugin_dir: Annotated[
        Optional[str],
        typer.Option("--plugin-dir", help="Path to the directory containing plugin modules."),
    ] = None,
) -> None:
    """Run RFT process."""
    if plugin_dir:
        os.environ[PLUGIN_DIRS_ENV_VAR] = plugin_dir
    load_plugins()
    cfg = load_config(config)

    if dlc:
        cluster_namespace = f"{cfg.project}-{cfg.name}"
        cfg.cluster.ray_address = setup_ray_cluster(namespace=cluster_namespace)

    if not is_running():
        raise RuntimeError("Ray is not running, please start it by `ray start --head`.")

    try:
        if cfg.stages:
            from trinity.trainer.verl.utils import get_latest_hf_checkpoint_path

            state_manager = StateManager(
                path=os.path.join(cfg.checkpoint_root_dir, cfg.project, cfg.name)
            )
            latest_stage = state_manager.load_stage().get("latest_stage", 0)
            prev_stage_checkpoint = None
            for i, stage_config in enumerate(cfg):
                if i < latest_stage:
                    logger.info(
                        "===========================================================\n"
                        f"> Skipping completed stage {i + 1}/{len(cfg.stages)}...\n"
                        "==========================================================="
                    )
                    stage_config.check_and_update()
                else:
                    logger.info(
                        "===========================================================\n"
                        f"> Starting stage {i + 1}/{len(cfg.stages)}...\n"
                        "==========================================================="
                    )
                    state_manager.save_stage(i)
                    if prev_stage_checkpoint is not None:
                        stage_config.model.model_path = prev_stage_checkpoint
                    stage_config.check_and_update()
                    run_stage(stage_config)
                    logger.info(
                        "===========================================================\n"
                        f"> Stage {i + 1}/{len(cfg.stages)} finished.\n"
                        "==========================================================="
                    )
                prev_stage_checkpoint = get_latest_hf_checkpoint_path(stage_config)
        else:
            cfg.check_and_update()
            run_stage(cfg)

    finally:
        if dlc:
            stop_ray_cluster(namespace=cluster_namespace)


@app.command()
def studio(
    port: Annotated[
        int,
        typer.Option("--port", help="The port for Trinity-Studio."),
    ] = 8501,
) -> None:
    """Run studio to manage configurations."""
    from trinity.manager.config_manager import ConfigManager

    ConfigManager.run(port)


@app.command()
def debug(
    config: Annotated[
        str,
        typer.Option("--config", help="Path to the config file."),
    ],
    module: Annotated[
        str,
        typer.Option(
            "--module",
            help="The module to debug: 'inference_model', 'workflow', or 'viewer'.",
        ),
    ],
    plugin_dir: Annotated[
        Optional[str],
        typer.Option("--plugin-dir", help="Path to the directory containing plugin modules."),
    ] = None,
    output_dir: Annotated[
        str,
        typer.Option("--output-dir", help="The output directory for debug files."),
    ] = "debug_output",
    disable_overwrite: Annotated[
        bool,
        typer.Option("--disable-overwrite", help="Disable overwriting the output directory."),
    ] = False,
    enable_profiling: Annotated[
        bool,
        typer.Option("--enable-profiling", help="Whether to use viztracer for workflow profiling."),
    ] = False,
    port: Annotated[
        int,
        typer.Option("--port", help="The port for Experience Viewer."),
    ] = 8502,
) -> None:
    """Debug a workflow implementation."""
    valid_modules = ("inference_model", "workflow", "viewer")
    if module not in valid_modules:
        raise typer.BadParameter(f"Only support {valid_modules} for debugging, got '{module}'")

    if plugin_dir:
        os.environ[PLUGIN_DIRS_ENV_VAR] = plugin_dir
    load_plugins()
    cfg = load_config(config)
    cfg.mode = "explore"
    cfg.ray_namespace = DEBUG_NAMESPACE
    cfg.check_and_update()
    sys.path.insert(0, os.getcwd())
    ray.init(
        namespace=cfg.ray_namespace,
        runtime_env={"env_vars": cfg.get_envs()},
        ignore_reinit_error=True,
    )

    from trinity.common.models import create_debug_explorer_model

    if module == "inference_model":
        asyncio.run(create_debug_explorer_model(cfg))

    elif module == "workflow":
        from trinity.explorer.workflow_runner import DebugWorkflowRunner

        runner = DebugWorkflowRunner(cfg, output_dir, enable_profiling, disable_overwrite)
        asyncio.run(runner.debug())

    elif module == "viewer":
        from trinity.buffer.viewer import SQLExperienceViewer

        output_dir_abs = os.path.abspath(output_dir).rstrip("/")

        SQLExperienceViewer.run_viewer(
            model_path=cfg.model.model_path,
            db_url=f"sqlite:///{os.path.join(output_dir_abs, 'debug_buffer.db')}",
            table_name="debug_buffer",
            port=port,
        )


@app.command()
def convert(
    checkpoint_dir: Annotated[
        str,
        typer.Option("--checkpoint-dir", help="The path to the checkpoint directory."),
    ],
    base_model_dir: Annotated[
        Optional[str],
        typer.Option("--base-model-dir", help="The path to the base model."),
    ] = None,
) -> None:
    """Convert checkpoints to huggingface format."""
    dir_path = checkpoint_dir
    if "global_step_" in dir_path:
        while not os.path.basename(dir_path).startswith("global_step_"):
            dir_path = os.path.dirname(dir_path)
    converter = Converter(base_model_dir)
    converter.convert(dir_path)


@app.command()
def log(
    log_dir: Annotated[
        str,
        typer.Option(
            "--log-dir",
            "-d",
            help="Path to the log directory. If provided, it will be used directly and ignore --config.",
        ),
    ] = "",
    config: Annotated[
        str,
        typer.Option(
            "--config",
            help="Path to the config file. If provided, it will automatically locate the log directory based on the config.",
        ),
    ] = "",
    keyword: Annotated[
        Optional[str],
        typer.Option("--keyword", "-k", help="The keyword to filter log files."),
    ] = None,
    level: Annotated[
        str,
        typer.Option("--level", "-l", help="The minimum log level to display."),
    ] = "INFO",
    last_n_lines: Annotated[
        int,
        typer.Option("--last-n-lines", "-n", help="Number of last lines to display when starting."),
    ] = 0,
    search_pattern: Annotated[
        Optional[str],
        typer.Option(
            "--search-pattern",
            "-p",
            help="The pattern to search in log files. Only search for history logs and display all lines containing the pattern.",
        ),
    ] = None,
    no_color: Annotated[
        bool,
        typer.Option("--no-color", help="Disable colored output."),
    ] = False,
) -> None:
    """Monitor log files in real-time."""
    from trinity.manager.log_manager import LogManager

    if not config and not log_dir:
        raise typer.BadParameter("Either --config or --log-dir must be provided.")
    if not log_dir:
        cfg = load_config(config)
        checkpoint_job_dir = cfg.get_checkpoint_job_dir()
        # we do not use check_and_update here because user may use this command
        # in another environment
        log_dir = os.path.join(checkpoint_job_dir, "log")
        if not os.path.exists(log_dir):
            raise FileNotFoundError(f"Log directory not found: {log_dir}")
    if not os.path.exists(log_dir):
        raise FileNotFoundError(f"Log directory not found: {log_dir}")

    log_manager = LogManager(
        log_dir,
        keyword=keyword,
        min_level=level,
        color_output=not no_color,
        last_n_lines=last_n_lines,
        search_pattern=search_pattern,
    )
    log_manager.monitor()


def main() -> None:
    """The main entrypoint."""
    app()


if __name__ == "__main__":
    main()