Optimization Engine

engine

Async optimization engine with retries, timeouts, caching, and sync wrapper.

Optimizer

Optimizer(*, adapter: LangChainAdapter | LangGraphAdapter | GoogleADKAdapter, strategy: SearchStrategy, evaluators: Sequence[Evaluator], objective: Objective, store: SQLiteStore | None = None, cache: Cache | None = None, settings: Settings | None = None, stop_conditions: Sequence[StopCondition] | None = None, retry_policy: RetryPolicy | None = None, callbacks: Sequence[EventCallback] = (), sleeper: Callable[[float], Awaitable[None]] = sleep, clock: Callable[[], float] = monotonic, duplicate_suggestion_limit: int = 100)

Tune pipelines through official LangChain, LangGraph, or Google ADK adapters.

Source code in src/tunarag/engine.py

def __init__(
    self,
    *,
    adapter: LangChainAdapter | LangGraphAdapter | GoogleADKAdapter,
    strategy: SearchStrategy,
    evaluators: Sequence[Evaluator],
    objective: Objective,
    store: SQLiteStore | None = None,
    cache: Cache | None = None,
    settings: Settings | None = None,
    stop_conditions: Sequence[StopCondition] | None = None,
    retry_policy: RetryPolicy | None = None,
    callbacks: Sequence[EventCallback] = (),
    sleeper: Callable[[float], Awaitable[None]] = asyncio.sleep,
    clock: Callable[[], float] = time.monotonic,
    duplicate_suggestion_limit: int = 100,
) -> None:
    if not isinstance(adapter, (LangChainAdapter, LangGraphAdapter, GoogleADKAdapter)):
        raise TypeError(
            "optimizer requires LangChainAdapter, LangGraphAdapter, or GoogleADKAdapter; "
            "custom RAG adapters are unsupported"
        )
    if not evaluators:
        raise ValueError("optimizer requires at least one evaluator")
    if duplicate_suggestion_limit < 1:
        raise ValueError("duplicate suggestion limit must be positive")
    self._adapter = adapter
    self._strategy = strategy
    self._evaluators = tuple(evaluators)
    self._objective = objective
    self._store = SQLiteStore() if store is None else store
    self._owns_store = store is None
    self._cache = cache
    self._settings = Settings.resolve(environ={}) if settings is None else settings
    self._stop_conditions = (
        (MaxTrials(self._settings.max_trials),)
        if stop_conditions is None
        else tuple(stop_conditions)
    )
    if not self._stop_conditions:
        raise ValueError("optimizer requires at least one stop condition")
    self._retry_policy = retry_policy or RetryPolicy(
        max_retries=self._settings.max_retries,
        sample_timeout_seconds=self._settings.sample_timeout_seconds,
        trial_timeout_seconds=self._settings.trial_timeout_seconds,
    )
    self._callbacks = tuple(callbacks)
    self._event_cursors: dict[str, int] = {}
    self._event_locks: dict[str, asyncio.Lock] = {}
    self._sleeper = sleeper
    self._clock = clock
    self._duplicate_suggestion_limit = duplicate_suggestion_limit

store property

store: SQLiteStore

Return the experiment store used by this optimizer.

aoptimize async

aoptimize(dataset: Dataset, *, study_name: str) -> OptimizationResult

Run a new study and return terminal results.

Source code in src/tunarag/engine.py

async def aoptimize(self, dataset: Dataset, *, study_name: str) -> OptimizationResult:
    """Run a new study and return terminal results."""

    self._validate_dataset(dataset)
    study = self._store.create_study(study_name, self._study_configuration(dataset))
    study = self._store.transition_study(study.id, StudyStatus.RUNNING)
    self._event_cursors[study.id] = -1
    self._event_locks[study.id] = asyncio.Lock()
    await self._dispatch_events(study.id)
    return await self._continue_study(dataset, study, [], [])

aresume async

aresume(dataset: Dataset, *, study_id: str) -> OptimizationResult

Resume a compatible durable study after replaying committed successes.

Source code in src/tunarag/engine.py

async def aresume(self, dataset: Dataset, *, study_id: str) -> OptimizationResult:
    """Resume a compatible durable study after replaying committed successes."""

    self._validate_dataset(dataset)
    try:
        study = self._store.get_study(study_id)
    except KeyError as error:
        raise ResumeError(
            "cannot resume an unknown study",
            stage="resume.load",
            details={"study_id": study_id},
        ) from error
    expected_hash = content_hash(self._study_configuration(dataset), namespace="study:v1")
    if study.config_hash != expected_hash:
        raise ResumeError(
            "study configuration is incompatible with the current optimizer",
            stage="resume.validate",
            details={"study_id": study_id},
        )
    if study.status in {StudyStatus.FAILED, StudyStatus.CANCELLED}:
        raise ResumeError(
            "failed or cancelled studies are terminal and cannot resume",
            stage="resume.validate",
            details={"study_id": study_id, "status": study.status.value},
        )

    self._event_cursors[study.id] = self._latest_event_sequence(study.id)
    self._event_locks[study.id] = asyncio.Lock()
    summaries = self._load_terminal_summaries(study.id)
    if study.status is StudyStatus.COMPLETED:
        await self._replay_strategy(study.id, summaries)
        return _build_result(
            study.id, study.status, self._persisted_stop_reason(study.id), summaries
        )
    if study.status is StudyStatus.CREATED:
        study = self._store.transition_study(study.id, StudyStatus.RUNNING)

    abandoned = self._store.abandon_expired_attempts(study_id=study.id)
    recoverable_ids = {trial.id for trial in self._store.list_recoverable_trials(study.id)}
    unfinished = [
        trial
        for trial in self._store.list_trials(study.id)
        if trial.status is TrialStatus.PENDING or trial.id in recoverable_ids
    ]
    active = [
        trial.id
        for trial in self._store.list_trials(study.id)
        if trial.status is TrialStatus.RUNNING and trial.id not in recoverable_ids
    ]
    if active:
        await self._dispatch_events(study.id)
        raise ResumeError(
            "study has attempts with active recovery leases",
            stage="resume.recover",
            retryable=True,
            details={"study_id": study.id, "trial_ids": active},
        )
    await self._replay_strategy(study.id, summaries)
    self._store.append_event(
        study.id,
        "study.resumed",
        {
            "terminal_trials": len(summaries),
            "unfinished_trials": len(unfinished),
            "abandoned_attempts": len(abandoned),
        },
    )
    await self._dispatch_events(study.id)
    return await self._continue_study(dataset, study, summaries, unfinished)

optimize

optimize(dataset: Dataset, *, study_name: str) -> OptimizationResult

Run aoptimize when no event loop is active.

Source code in src/tunarag/engine.py

def optimize(self, dataset: Dataset, *, study_name: str) -> OptimizationResult:
    """Run `aoptimize` when no event loop is active."""

    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(self.aoptimize(dataset, study_name=study_name))
    raise SyncInAsyncContextError()

resume

resume(dataset: Dataset, *, study_id: str) -> OptimizationResult

Run aresume when no event loop is active.

Source code in src/tunarag/engine.py

def resume(self, dataset: Dataset, *, study_id: str) -> OptimizationResult:
    """Run `aresume` when no event loop is active."""

    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(self.aresume(dataset, study_id=study_id))
    raise SyncInAsyncContextError()

close

close() -> None

Close an internally created store; caller-owned resources remain open.

Source code in src/tunarag/engine.py

def close(self) -> None:
    """Close an internally created store; caller-owned resources remain open."""

    if self._owns_store:
        self._store.close()