56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
"""Lazy Optuna study integration used for durable trial history and pruning metadata."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class OptunaStudies:
|
|
def __init__(self, root: Path):
|
|
self.path = (root / "optuna.sqlite3").resolve()
|
|
self.url = f"sqlite:///{self.path}"
|
|
|
|
def _study(self, session_id: str):
|
|
import optuna
|
|
|
|
optuna.logging.set_verbosity(optuna.logging.WARNING)
|
|
return optuna.create_study(
|
|
study_name=f"reward-tuning-{session_id}",
|
|
storage=self.url,
|
|
direction="maximize",
|
|
load_if_exists=True,
|
|
pruner=optuna.pruners.SuccessiveHalvingPruner(
|
|
min_resource=300, reduction_factor=3, min_early_stopping_rate=0
|
|
),
|
|
)
|
|
|
|
def record(
|
|
self,
|
|
session_id: str,
|
|
reward_config: dict[str, Any],
|
|
score: float,
|
|
eligible: bool,
|
|
rung: int,
|
|
) -> int:
|
|
"""Record an externally proposed Agent config through Optuna ask/tell."""
|
|
import optuna
|
|
|
|
study = self._study(session_id)
|
|
trial = study.ask()
|
|
trial.set_user_attr("reward_config", reward_config)
|
|
trial.set_user_attr("eligible", eligible)
|
|
trial.set_user_attr("rung", rung)
|
|
state = optuna.trial.TrialState.COMPLETE if eligible else optuna.trial.TrialState.PRUNED
|
|
study.tell(trial, score if eligible else None, state=state)
|
|
return trial.number
|
|
|
|
def summary(self, session_id: str) -> dict[str, Any]:
|
|
study = self._study(session_id)
|
|
completed = [trial for trial in study.trials if trial.value is not None]
|
|
return {
|
|
"studyName": study.study_name,
|
|
"trialCount": len(study.trials),
|
|
"bestValue": max((trial.value for trial in completed), default=None),
|
|
}
|