Repository: oumi-ai/oumi
File: src/oumi/train.py (Python)

Code before:
            runtime_env={
                "env_vars": {
                    "TOKENIZERS_PARALLELISM": "true",
                    "NCCL_DEBUG": "WARN",
                    "VLLM_LOGGING_LEVEL": "WARN",
                }
            }
        )

    # We define the remote function as a sub function so that the `@ray.remote`
    # decorator is only run if this function is run. This function should only be run
    # if ray is installed, preventing errors when it isn't.
    @ray.remote
    def _run_verl_train(
        partial_trainer: Callable[[], BaseTrainer], checkpoint_location: Optional[str]
    ):
        trainer = partial_trainer()
        trainer.train(resume_from_checkpoint=checkpoint_location)

        logger.info("Training is Complete.")

    ray.get(_run_verl_train.remote(partial_trainer, checkpoint_location))
    logger.info(
        "\n\n» We're always looking for feedback. "
        "What's one thing we can improve? https://oumi.ai/feedback"
    )


def train(
    config: TrainingConfig,
    additional_model_kwargs: Optional[dict[str, Any]] = None,
    additional_trainer_kwargs: Optional[dict[str, Any]] = None,
) -> None:
    """Trains a model using the provided configuration."""
    _START_TIME = time.time()

    _create_training_dirs(config)
    _log_training_info(config)

    # Configure logging to file
    log_dir = Path(config.training.output_dir) / "logs"
    for logger_name in ("oumi", "oumi.telemetry"):
        configure_logger(logger_name, level=config.training.log_level, log_dir=log_dir)

    telemetry_dir = config.training.telemetry_dir

    config = _finalize_training_config(config)

    if is_local_process_zero():
        logger.info(f"TrainingConfig:\n{pformat(config)}")
        if telemetry_dir and is_world_process_zero():

Review comment (refactor):
Make this a helper method like `log_feedback_request` so you don't duplicate the message for each trainer.

Code after:
        raise RuntimeError(
            "ray is not installed. Please install it with `pip install 'oumi[gpu]'`."
        )
    if not ray.is_initialized():
        logger.info("Initializing Ray cluster...")
        ray.init(
            runtime_env={
                "env_vars": {
                    "TOKENIZERS_PARALLELISM": "true",
                    "NCCL_DEBUG": "WARN",
                    "VLLM_LOGGING_LEVEL": "WARN",
                }
            }
        )

    # We define the remote function as a sub function so that the `@ray.remote`
    # decorator is only run if this function is run. This function should only be run
    # if ray is installed, preventing errors when it isn't.
    @ray.remote
    def _run_verl_train(
        partial_trainer: Callable[[], BaseTrainer], checkpoint_location: Optional[str]
    ):
        trainer = partial_trainer()
        trainer.train(resume_from_checkpoint=checkpoint_location)

        logger.info("Training is Complete.")

    ray.get(_run_verl_train.remote(partial_trainer, checkpoint_location))
    _log_feedback_request()


def train(
    config: TrainingConfig,
    additional_model_kwargs: Optional[dict[str, Any]] = None,
    additional_trainer_kwargs: Optional[dict[str, Any]] = None,
) -> None:
    """Trains a model using the provided configuration."""
    _START_TIME = time.time()

    _create_training_dirs(config)
    _log_training_info(config)

    # Configure logging to file
    log_dir = Path(config.training.output_dir) / "logs"
    for logger_name in ("oumi", "oumi.telemetry"):
        configure_logger(logger_name, level=config.training.log_level, log_dir=log_dir)

    telemetry_dir = config.training.telemetry_dir

    config = _finalize_training_config(config)