"""Native Optuna hyperparameter optimization executor. Provides direct Optuna integration without requiring Ray Tune as an intermediary. Runs trials sequentially on the local machine using Ludwig's standard training API. Supports AutoSampler (auto-selects best algorithm), GPSampler (Bayesian optimization), TPE, CMA-ES, and other Optuna samplers. For distributed execution, use the Ray executor with OptunaSearch instead. Usage in Ludwig config: hyperopt: executor: type: optuna num_samples: 50 sampler: auto # auto, gp, tpe, cmaes, random """ import copy import logging import os import traceback from typing import Any from ludwig.api import LudwigModel from ludwig.constants import MAXIMIZE, TEST, TRAINING, VALIDATION from ludwig.hyperopt.results import HyperoptResults, TrialResults from ludwig.hyperopt.utils import substitute_parameters from ludwig.utils.defaults import default_random_seed logger = logging.getLogger(__name__) def _create_sampler(sampler_type: str): """Create an Optuna sampler from type string.""" import optuna if sampler_type == "auto": try: return optuna.samplers.AutoSampler() except AttributeError: logger.info("AutoSampler not available, falling back to TPE") return optuna.samplers.TPESampler() elif sampler_type == "gp": try: return optuna.samplers.GPSampler() except AttributeError: logger.info("GPSampler not available, falling back to TPE") return optuna.samplers.TPESampler() elif sampler_type == "tpe": return optuna.samplers.TPESampler() elif sampler_type == "cmaes": return optuna.samplers.CmaEsSampler() elif sampler_type == "random": return optuna.samplers.RandomSampler() else: raise ValueError(f"Unknown sampler: {sampler_type}. Options: auto, gp, tpe, cmaes, random") def _suggest_params(trial, parameters: dict) -> dict[str, Any]: """Suggest parameter values for a trial based on the search space definition.""" params = {} for param_name, space_def in parameters.items(): space_type = space_def.get("space", "uniform") if space_type == "uniform": params[param_name] = trial.suggest_float(param_name, space_def["lower"], space_def["upper"]) elif space_type == "loguniform": params[param_name] = trial.suggest_float(param_name, space_def["lower"], space_def["upper"], log=True) elif space_type in ("int", "randint", "qrandint"): params[param_name] = trial.suggest_int(param_name, int(space_def["lower"]), int(space_def["upper"])) elif space_type in ("choice", "categorical"): params[param_name] = trial.suggest_categorical(param_name, space_def["categories"]) elif space_type == "grid_search": params[param_name] = trial.suggest_categorical(param_name, space_def["values"]) else: raise ValueError(f"Unknown search space type: {space_type} for parameter {param_name}") return params class OptunaExecutor: """Native Optuna hyperparameter optimization executor. Runs trials sequentially on the local machine. Each trial trains a full Ludwig model with parameters suggested by Optuna, then reports the validation metric back. """ def __init__( self, parameters: dict, output_feature: str, metric: str, goal: str, split: str, search_alg: dict | None = None, num_samples: int = 10, sampler: str = "auto", pruner: str | None = None, study_name: str | None = None, storage: str | None = None, **kwargs, ) -> None: try: import optuna # noqa: F401 except ImportError: raise ImportError("Optuna is required for the optuna executor. Install with: pip install optuna") self.parameters = parameters self.output_feature = output_feature self.metric = metric self.goal = goal self.split = split self.num_samples = num_samples self.sampler_type = sampler self.pruner_type = pruner self.study_name = study_name or "ludwig_hyperopt" self.storage = storage def execute( self, config, dataset=None, training_set=None, validation_set=None, test_set=None, training_set_metadata=None, data_format=None, experiment_name="hyperopt", model_name="run", resume=None, skip_save_training_description=False, skip_save_training_statistics=False, skip_save_model=False, skip_save_progress=False, skip_save_log=False, skip_save_processed_input=True, skip_save_unprocessed_output=False, skip_save_predictions=False, skip_save_eval_stats=False, output_directory="results", gpus=None, gpu_memory_limit=None, allow_parallel_threads=True, callbacks=None, tune_callbacks=None, backend=None, random_seed=default_random_seed, debug=False, hyperopt_log_verbosity=3, **kwargs, ) -> HyperoptResults: import optuna sampler_obj = _create_sampler(self.sampler_type) pruner_obj = None if self.pruner_type == "median": pruner_obj = optuna.pruners.MedianPruner() elif self.pruner_type == "hyperband": pruner_obj = optuna.pruners.HyperbandPruner() direction = "minimize" if self.goal != MAXIMIZE else "maximize" study = optuna.create_study( study_name=self.study_name, direction=direction, sampler=sampler_obj, pruner=pruner_obj, storage=self.storage, load_if_exists=True, ) trial_results = [] output_dir = os.path.join(output_directory, experiment_name) os.makedirs(output_dir, exist_ok=True) def objective(trial): sampled_params = _suggest_params(trial, self.parameters) for cb in callbacks or []: cb.on_hyperopt_trial_start(sampled_params) # Substitute sampled parameters into config trial_config = copy.deepcopy(config) substitute_parameters(trial_config, sampled_params) trial_dir = os.path.join(output_dir, f"trial_{trial.number}") os.makedirs(trial_dir, exist_ok=True) try: model = LudwigModel( config=trial_config, backend=backend, gpus=gpus, gpu_memory_limit=gpu_memory_limit, allow_parallel_threads=allow_parallel_threads, callbacks=callbacks, ) train_result = model.train( dataset=dataset, training_set=training_set, validation_set=validation_set, test_set=test_set, training_set_metadata=training_set_metadata, data_format=data_format, experiment_name=f"trial_{trial.number}", model_name=model_name, skip_save_training_description=skip_save_training_description, skip_save_training_statistics=skip_save_training_statistics, skip_save_model=skip_save_model, skip_save_progress=skip_save_progress, skip_save_log=skip_save_log, skip_save_processed_input=skip_save_processed_input, output_directory=trial_dir, random_seed=random_seed + trial.number, ) train_stats = train_result.train_stats preprocessed_data = train_result.preprocessed_data # Evaluate on the target split eval_split = self.split eval_dataset = None if eval_split == TRAINING: eval_dataset = preprocessed_data.training_set elif eval_split == VALIDATION: eval_dataset = preprocessed_data.validation_set elif eval_split == TEST: eval_dataset = preprocessed_data.test_set eval_stats = {} if eval_dataset is not None: eval_stats_list, _, _ = model.evaluate( dataset=eval_dataset, skip_save_unprocessed_output=True, skip_save_predictions=True, skip_save_eval_stats=True, callbacks=callbacks, ) eval_stats = eval_stats_list # Extract the target metric metric_value = None if self.output_feature in eval_stats: feature_stats = eval_stats[self.output_feature] if self.metric in feature_stats: metric_value = feature_stats[self.metric] elif "combined" in eval_stats and self.metric in eval_stats["combined"]: metric_value = eval_stats["combined"][self.metric] if metric_value is None: raise ValueError( f"Could not find metric '{self.metric}' for output feature " f"'{self.output_feature}' in evaluation stats: {list(eval_stats.keys())}" ) trial_results.append( TrialResults( parameters=sampled_params, metric_score=metric_value, training_stats=train_stats, eval_stats=eval_stats, ) ) for cb in callbacks or []: cb.on_hyperopt_trial_end(sampled_params) logger.info( f"Trial {trial.number}: {self.output_feature}.{self.metric} = {metric_value:.6f} " f"(params: {sampled_params})" ) return metric_value except Exception as e: logger.error(f"Trial {trial.number} failed: {e}\n{traceback.format_exc()}") for cb in callbacks or []: cb.on_hyperopt_trial_end(sampled_params) raise optuna.TrialPruned(f"Trial failed: {e}") logger.info( f"Starting Optuna hyperopt: {self.num_samples} trials, " f"{self.goal} {self.output_feature}.{self.metric}, sampler={self.sampler_type}" ) study.optimize(objective, n_trials=self.num_samples) # Sort results by metric score trial_results.sort(key=lambda t: t.metric_score, reverse=(self.goal == MAXIMIZE)) logger.info( f"Optuna hyperopt complete. Best {self.output_feature}.{self.metric}: " f"{study.best_value:.6f}, params: {study.best_params}" ) return HyperoptResults(ordered_trials=trial_results, experiment_analysis=study)