orchestrator

orchard.optimization.orchestrator

Optuna Study Orchestrator Package.

High-level coordination for hyperparameter optimization studies. Provides a modular architecture for study creation, execution, visualization, and result export.

Key Components

OptunaOrchestrator: Primary study lifecycle manager. run_optimization: Convenience function for full pipeline execution (study creation → trial loop → artifact export). export_best_config, export_study_summary, export_top_trials: Post-study artifact generators.

Typical Usage

from orchard.optimization.orchestrator import run_optimization study = run_optimization(cfg=config, device=device, paths=paths)

TrialData(number, value, params, datetime_start=None, datetime_complete=None, state=None, duration_seconds=None) dataclass

Immutable snapshot of Optuna trial metadata for serialization.

Attributes:

Name	Type	Description
number	int	Trial number within the study.
value	float | None	Objective value (None for incomplete trials).
params	dict[str, Any]	Hyperparameter values sampled for this trial.
datetime_start	str | None	ISO-formatted start timestamp.
datetime_complete	str | None	ISO-formatted completion timestamp.
state	str | None	Trial state name (COMPLETE, PRUNED, FAIL, etc.).
duration_seconds	float | None	Wall-clock duration in seconds.

from_trial(trial) classmethod

Build from an Optuna FrozenTrial, computing duration if timestamps are available.

Parameters:

Name	Type	Description	Default
trial	FrozenTrial	Frozen trial from study.	required

Returns:

Type	Description
TrialData	Immutable trial snapshot with computed duration.

Source code in orchard/optimization/orchestrator/exporters.py

@classmethod
def from_trial(cls, trial: optuna.trial.FrozenTrial) -> TrialData:
    """
    Build from an Optuna FrozenTrial, computing duration if timestamps are available.

    Args:
        trial: Frozen trial from study.

    Returns:
        Immutable trial snapshot with computed duration.
    """
    duration = None
    if trial.datetime_complete and trial.datetime_start:
        duration = (trial.datetime_complete - trial.datetime_start).total_seconds()
    return cls(
        number=trial.number,
        value=trial.value,
        params=trial.params,
        state=trial.state.name,
        datetime_start=trial.datetime_start.isoformat() if trial.datetime_start else None,
        datetime_complete=(
            trial.datetime_complete.isoformat() if trial.datetime_complete else None
        ),
        duration_seconds=duration,
    )

to_dict()

Serialize to plain dictionary for JSON export.

Returns:

Type	Description
dict[str, Any]	Dictionary representation with all fields.

Source code in orchard/optimization/orchestrator/exporters.py

def to_dict(self) -> dict[str, Any]:
    """
    Serialize to plain dictionary for JSON export.

    Returns:
        Dictionary representation with all fields.
    """
    return asdict(self)

OptunaOrchestrator(cfg, device, paths, tracker=None)

High-level manager for Optuna hyperparameter optimization studies.

Coordinates the complete optimization lifecycle: study creation, trial execution, and post-processing artifact generation. Integrates with Orchard ML's Config and RunPaths infrastructure, delegating specialized tasks (sampler/pruner building, visualization, export) to focused submodules.

This orchestrator serves as the entry point for hyperparameter tuning, wrapping Optuna's API with Orchard ML-specific configuration and output management.

Attributes:

Name	Type	Description
cfg	Config	Template configuration that will be overridden per trial
device	device	Hardware target for training (CPU/CUDA/MPS)
paths	RunPaths	Output directory structure for artifacts and results

Example

orchestrator = OptunaOrchestrator(cfg=config, device=device, paths=paths) study = orchestrator.optimize() print(f"Best AUC: {study.best_value:.3f}")

Artifacts saved to paths.figures/ and paths.exports/

Initialize orchestrator.

Parameters:

Name	Type	Description	Default
cfg	Config	Base Config to override per trial	required
device	device	PyTorch device for training	required
paths	RunPaths	Root directory for outputs	required
tracker	TrackerProtocol | None	Optional experiment tracker for nested trial logging	None

Source code in orchard/optimization/orchestrator/orchestrator.py

def __init__(
    self,
    cfg: Config,
    device: torch.device,
    paths: RunPaths,
    tracker: TrackerProtocol | None = None,
) -> None:
    """
    Initialize orchestrator.

    Args:
        cfg (Config): Base Config to override per trial
        device (torch.device): PyTorch device for training
        paths (RunPaths): Root directory for outputs
        tracker (TrackerProtocol | None): Optional experiment tracker for nested trial logging
    """
    self.cfg = cfg
    self.device = device
    self.paths = paths
    self.tracker = tracker

create_study()

Create or load Optuna study with configured sampler and pruner.

Returns:

Type	Description
Study	Configured Optuna study instance

Source code in orchard/optimization/orchestrator/orchestrator.py

def create_study(self) -> optuna.Study:
    """
    Create or load Optuna study with configured sampler and pruner.

    Returns:
        Configured Optuna study instance
    """
    sampler = build_sampler(self.cfg.optuna)
    pruner = build_pruner(self.cfg.optuna)
    storage_url = self.cfg.optuna.get_storage_url(self.paths)

    study = optuna.create_study(
        study_name=self.cfg.optuna.study_name,
        direction=self.cfg.optuna.direction,
        sampler=sampler,
        pruner=pruner,
        storage=storage_url,
        load_if_exists=self.cfg.optuna.load_if_exists,
    )

    return study

optimize()

Execute hyperparameter optimization.

Returns:

Type	Description
Study	Completed study with trial results

Source code in orchard/optimization/orchestrator/orchestrator.py

def optimize(self) -> optuna.Study:
    """
    Execute hyperparameter optimization.

    Returns:
        Completed study with trial results
    """
    # Suppress Optuna's internal INFO logs (e.g. "A new study created in RDB")
    # before create_study(); our own header in phases.py is sufficient
    optuna.logging.set_verbosity(optuna.logging.WARNING)

    study = self.create_study()
    search_space = get_search_space(
        self.cfg.optuna.search_space_preset,
        resolution=self.cfg.dataset.resolution,
        include_models=self.cfg.optuna.enable_model_search,
        model_pool=self.cfg.optuna.model_pool,
        overrides=self.cfg.optuna.search_space_overrides,
    )

    objective = OptunaObjective(
        cfg=self.cfg,
        search_space=search_space,
        device=self.device,
        tracker=self.tracker,
    )

    # Configure callbacks and log our structured header
    log_optimization_header(self.cfg)

    callbacks = build_callbacks(self.cfg.optuna, self.cfg.training.monitor_metric)

    study.set_user_attr("n_trials", self.cfg.optuna.n_trials)

    interrupted = False
    try:
        study.optimize(
            objective,
            n_trials=self.cfg.optuna.n_trials,
            timeout=self.cfg.optuna.timeout,
            n_jobs=self.cfg.optuna.n_jobs,
            show_progress_bar=self.cfg.optuna.show_progress_bar,
            callbacks=callbacks,
        )
    except KeyboardInterrupt:
        interrupted = True
        logger.warning("Optimization interrupted by user. Saving partial results...")

    self._post_optimization_processing(study)

    if interrupted:
        logger.warning(
            "Continuing to training in 5 seconds... (Ctrl+C again to abort pipeline)"
        )
        time.sleep(5)  # grace period for the user to fully abort

    return study

export_best_config(study, cfg, paths)

Export best trial configuration as YAML file.

Creates a new Config instance with best hyperparameters applied, validates it, and saves to reports/best_config.yaml.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study with at least one successful trial	required
cfg	Config	Template configuration (used for non-optimized parameters)	required
paths	RunPaths	RunPaths instance for output location	required

Returns:

Type	Description
Path | None	Path to exported config file, or None if no completed trials

Note

Skips export with warning if no completed trials exist.

Example

export_best_config(study, cfg, paths)

Creates: {paths.reports}/best_config.yaml

Source code in orchard/optimization/orchestrator/exporters.py

def export_best_config(study: optuna.Study, cfg: Config, paths: RunPaths) -> Path | None:
    """
    Export best trial configuration as YAML file.

    Creates a new Config instance with best hyperparameters applied,
    validates it, and saves to reports/best_config.yaml.

    Args:
        study: Completed Optuna study with at least one successful trial
        cfg: Template configuration (used for non-optimized parameters)
        paths: RunPaths instance for output location

    Returns:
        Path to exported config file, or None if no completed trials

    Note:
        Skips export with warning if no completed trials exist.

    Example:
        >>> export_best_config(study, cfg, paths)
        # Creates: {paths.reports}/best_config.yaml
    """
    if not has_completed_trials(study):
        logger.warning("No completed trials. Cannot export best config.")
        return None

    # Build config dict with best parameters
    config_dict = build_best_config_dict(study.best_params, cfg)

    # Create and validate new config
    best_config = Config(**config_dict)

    # Save to YAML
    output_path = paths.reports / "best_config.yaml"
    save_config_as_yaml(best_config, output_path)

    return output_path

export_study_summary(study, paths)

Export complete study metadata to JSON.

Serializes all trials with parameters, values, states, timestamps, and durations. Handles studies with zero completed trials gracefully.

Parameters:

Name	Type	Description	Default
study	Study	Optuna study (may contain failed/pruned trials)	required
paths	RunPaths	RunPaths instance for output location	required

Output structure::

{
    "study_name": str,
    "direction": str,
    "n_trials": int,
    "n_completed": int,
    "best_trial": {...} or null,
    "trials": [...]
}

Example

export_study_summary(study, paths)

Creates: {paths.reports}/study_summary.json

Source code in orchard/optimization/orchestrator/exporters.py

def export_study_summary(study: optuna.Study, paths: RunPaths) -> None:
    """
    Export complete study metadata to JSON.

    Serializes all trials with parameters, values, states, timestamps,
    and durations. Handles studies with zero completed trials gracefully.

    Args:
        study: Optuna study (may contain failed/pruned trials)
        paths: RunPaths instance for output location

    Output structure::

        {
            "study_name": str,
            "direction": str,
            "n_trials": int,
            "n_completed": int,
            "best_trial": {...} or null,
            "trials": [...]
        }

    Example:
        >>> export_study_summary(study, paths)
        # Creates: {paths.reports}/study_summary.json
    """
    completed = get_completed_trials(study)

    # Build best trial data (may be None if no completed trials)
    best_trial_data = build_best_trial_data(study, completed)

    summary = {
        "study_name": study.study_name,
        "direction": study.direction.name,
        "n_trials": len(study.trials),
        "n_completed": len(completed),
        "best_trial": best_trial_data.to_dict() if best_trial_data else None,
        "trials": [TrialData.from_trial(trial).to_dict() for trial in study.trials],
    }

    output_path = paths.reports / "study_summary.json"
    with open(output_path, "w") as f:
        json.dump(summary, f, indent=2)

    logger.info(
        "%s%s %-22s: %s",
        LogStyle.INDENT,
        LogStyle.ARROW,
        "Study Summary",
        Path(output_path).name,
    )

export_top_trials(study, paths, metric_name, top_k=10)

Export top K trials to Excel spreadsheet with professional formatting.

Creates human-readable comparison table of best-performing trials with hyperparameters, metric values, and durations. Applies professional Excel styling matching TrainingReport format.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study with at least one successful trial	required
paths	RunPaths	RunPaths instance for output location	required
metric_name	str	Name of optimization metric (for column header)	required
top_k	int	Number of top trials to export (default: 10)	10

DataFrame Columns:

• Rank: 1-based ranking
• Trial: Trial number
• {METRIC_NAME}: Objective value
• {param_name}: Each hyperparameter
• Duration (s): Trial duration if available

Example

export_top_trials(study, paths, "auc", top_k=10)

Creates: {paths.reports}/top_10_trials.xlsx

Source code in orchard/optimization/orchestrator/exporters.py

def export_top_trials(
    study: optuna.Study, paths: RunPaths, metric_name: str, top_k: int = 10
) -> None:
    """
    Export top K trials to Excel spreadsheet with professional formatting.

    Creates human-readable comparison table of best-performing trials
    with hyperparameters, metric values, and durations. Applies professional
    Excel styling matching TrainingReport format.

    Args:
        study: Completed Optuna study with at least one successful trial
        paths: RunPaths instance for output location
        metric_name: Name of optimization metric (for column header)
        top_k: Number of top trials to export (default: 10)

    DataFrame Columns:

    - Rank: 1-based ranking
    - Trial: Trial number
    - {METRIC_NAME}: Objective value
    - {param_name}: Each hyperparameter
    - Duration (s): Trial duration if available

    Example:
        >>> export_top_trials(study, paths, "auc", top_k=10)
        # Creates: {paths.reports}/top_10_trials.xlsx
    """
    completed = get_completed_trials(study)
    if not completed:
        logger.warning("No completed trials. Cannot export top trials.")
        return

    reverse = study.direction == optuna.study.StudyDirection.MAXIMIZE
    # Filter out trials with None or NaN values before sorting
    valid_trials = [
        t
        for t in completed
        if t.value is not None and not (isinstance(t.value, float) and math.isnan(t.value))
    ]
    sorted_trials = sorted(valid_trials, key=lambda t: cast(float, t.value), reverse=reverse)[
        :top_k
    ]

    df = build_top_trials_dataframe(sorted_trials, metric_name)

    output_path = paths.reports / "top_10_trials.xlsx"

    wb = Workbook()
    ws = wb.active
    ws.title = "Top Trials"

    _write_styled_rows(ws, df)
    _auto_adjust_column_widths(ws)

    wb.save(output_path)
    logger.info(
        "%s%s %-22s: %s (%d trials)",
        LogStyle.INDENT,
        LogStyle.ARROW,
        "Top Trials",
        Path(output_path).name,
        len(sorted_trials),
    )

run_optimization(cfg, device, paths, tracker=None)

Convenience function to run complete optimization pipeline.

Parameters:

Name	Type	Description	Default
cfg	Config	Global configuration with optuna section	required
device	device	PyTorch device for training	required
paths	RunPaths	RunPaths instance for output management	required
tracker	TrackerProtocol | None	Optional experiment tracker for nested trial logging	None

Returns:

Type	Description
Study	Completed Optuna study with trial results

Example

study = run_optimization(cfg=config, device=torch.device("cuda"), paths=paths) print(f"Best AUC: {study.best_value:.3f}")

Source code in orchard/optimization/orchestrator/orchestrator.py

def run_optimization(
    cfg: Config,
    device: torch.device,
    paths: RunPaths,
    tracker: TrackerProtocol | None = None,
) -> optuna.Study:
    """
    Convenience function to run complete optimization pipeline.

    Args:
        cfg (Config): Global configuration with optuna section
        device (torch.device): PyTorch device for training
        paths (RunPaths): RunPaths instance for output management
        tracker (TrackerProtocol | None): Optional experiment tracker for nested trial logging

    Returns:
        Completed Optuna study with trial results

    Example:
        >>> study = run_optimization(cfg=config, device=torch.device("cuda"), paths=paths)
        >>> print(f"Best AUC: {study.best_value:.3f}")
    """
    orchestrator = OptunaOrchestrator(cfg=cfg, device=device, paths=paths, tracker=tracker)
    return orchestrator.optimize()