tuning ¶

def get_args() -> argparse.Namespace:
    """Get the arguments when using from the commandline.

    Returns:
        Parsed command line arguments.
    """
    parser = argparse.ArgumentParser(description="Launch check_model.")
    parser.add_argument("-d", "--data", type=str, required=True, metavar="FILE", help="Path to input csv file.")
    parser.add_argument("-m", "--model", type=str, required=True, metavar="FILE", help="Path to model file.")
    parser.add_argument(
        "-e",
        "--data_config",
        type=str,
        required=True,
        metavar="FILE",
        help="Path to data config file.",
    )
    parser.add_argument(
        "-c",
        "--model_config",
        type=str,
        required=True,
        metavar="FILE",
        help="Path to yaml config training file.",
    )
    parser.add_argument(
        "-w",
        "--initial_weights",
        type=str,
        required=False,
        nargs="?",
        const=None,
        default=None,
        metavar="FILE",
        help="The path to the initial weights (optional).",
    )
    parser.add_argument(
        "--ray_results_dirpath",
        type=str,
        required=False,
        nargs="?",
        const=None,
        default=None,
        metavar="DIR_PATH",
        help="Location where ray_results output dir should be written. If None, uses ~/ray_results.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        required=False,
        nargs="?",
        const="best_model.pt",
        default="best_model.pt",
        metavar="FILE",
        help="The output file path to write the trained model to",
    )
    parser.add_argument(
        "-bm",
        "--best_metrics",
        type=str,
        required=False,
        nargs="?",
        const="best_metrics.csv",
        default="best_metrics.csv",
        metavar="FILE",
        help="The path to write the best metrics to",
    )
    parser.add_argument(
        "-bc",
        "--best_config",
        type=str,
        required=False,
        nargs="?",
        const="best_config.yaml",
        default="best_config.yaml",
        metavar="FILE",
        help="The path to write the best config to",
    )
    parser.add_argument(
        "-bo",
        "--best_optimizer",
        type=str,
        required=False,
        nargs="?",
        const="best_optimizer.pt",
        default="best_optimizer.pt",
        metavar="FILE",
        help="The path to write the best optimizer to",
    )
    parser.add_argument(
        "--tune_run_name",
        type=str,
        required=False,
        nargs="?",
        const=None,
        default=None,
        metavar="CUSTOM_RUN_NAME",
        help=(
            "Tells ray tune what the 'experiment_name' (i.e. the given tune_run name) should be. "
            "If set, the subdirectory of ray_results is named with this value and its train dir is prefixed accordingly. "
            "Default None means that ray will generate such a name on its own."
        ),
    )
    parser.add_argument(
        "--debug_mode",
        action="store_true",
        help="Activate debug mode for tuning. Default false, no debug.",
    )
    return parser.parse_args()

def main(
    model_path: str,
    data_path: str,
    data_config_path: str,
    model_config_path: str,
    initial_weights: str | None = None,  # noqa: ARG001
    ray_results_dirpath: str | None = None,
    output_path: str | None = None,
    best_optimizer_path: str | None = None,
    best_metrics_path: str | None = None,
    best_config_path: str | None = None,
    *,
    debug_mode: bool = False,
) -> None:
    """Run the main model checking pipeline.

    Args:
        data_path: Path to input data file.
        model_path: Path to model file.
        data_config_path: Path to data config file.
        model_config_path: Path to model config file.
        initial_weights: Optional path to initial weights.
        ray_results_dirpath: Directory for ray results.
        debug_mode: Whether to run in debug mode.
        output_path: Path to write the best model to.
        best_optimizer_path: Path to write the best optimizer to.
        best_metrics_path: Path to write the best metrics to.
        best_config_path: Path to write the best config to.
    """
    # Convert data config to proper type
    with open(data_config_path) as file:
        data_config_dict: dict[str, Any] = yaml.safe_load(file)
    data_config: yaml_data.YamlSubConfigDict = yaml_data.YamlSubConfigDict(**data_config_dict)

    with open(model_config_path) as file:
        model_config_dict: dict[str, Any] = yaml.safe_load(file)
    model_config: yaml_model_schema.Model = yaml_model_schema.Model(**model_config_dict)

    encoder_loader = loaders.EncoderLoader()
    encoder_loader.initialize_column_encoders_from_config(column_config=data_config.columns)

    model_class = launch_utils.import_class_from_file(model_path)

    ray_config_loader = yaml_model_schema.YamlRayConfigLoader(model=model_config)
    ray_config_model = ray_config_loader.get_config()

    tuner = raytune_learner.TuneWrapper(
        model_config=ray_config_model,
        data_config_path=data_config_path,
        model_class=model_class,
        data_path=data_path,
        encoder_loader=encoder_loader,
        seed=42,
        ray_results_dir=ray_results_dirpath,
        debug=debug_mode,
    )

    # Ensure output_path is provided
    if output_path is None:
        raise ValueError("output_path must not be None")
    try:
        grid_results = tuner.tune()
        if not grid_results:
            _raise_empty_grid()

        # Initialize parser with results
        parser = raytune_parser.TuneParser(result=grid_results)

        # Ensure output directory exists
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)

        # Save outputs using proper Result object API
        parser.save_best_model(output=output_path)
        parser.save_best_optimizer(output=best_optimizer_path)
        parser.save_best_metrics_dataframe(output=best_metrics_path)
        parser.save_best_config(output=best_config_path)

    except RuntimeError:
        logger.exception("Tuning failed")
        raise
    except KeyError:
        logger.exception("Missing expected result key")
        raise
    finally:
        if debug_mode:
            logger.info("Debug mode - preserving Ray results directory")
        elif ray_results_dirpath:
            shutil.rmtree(ray_results_dirpath, ignore_errors=True)

def run() -> None:
    """Run the model checking script."""
    args = get_args()
    main(
        data_path=args.data,
        model_path=args.model,
        data_config_path=args.data_config,
        model_config_path=args.model_config,
        initial_weights=args.initial_weights,
        ray_results_dirpath=args.ray_results_dirpath,
        output_path=args.output,
        best_optimizer_path=args.best_optimizer,
        best_metrics_path=args.best_metrics,
        best_config_path=args.best_config,
        debug_mode=args.debug_mode,
    )