Source code for ax.service.utils.storage

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

from typing import Tuple

from ax.core.experiment import Experiment
from ax.modelbridge.generation_strategy import GenerationStrategy
from ax.storage.sqa_store.db import init_engine_and_session_factory
from ax.storage.sqa_store.load import (
    _load_experiment,
    _load_generation_strategy_by_experiment_name,
)
from ax.storage.sqa_store.save import _save_experiment, _save_generation_strategy
from ax.storage.sqa_store.structs import DBSettings


"""Utilities for storing experiment to the database for AxClient."""


[docs]def load_experiment(name: str, db_settings: DBSettings) -> Experiment: """ Load experiment from the db. Service API only supports `Experiment`. Args: name: Experiment name. db_settings: Defines behavior for loading/saving experiment to/from db. Returns: ax.core.Experiment: Loaded experiment. """ init_engine_and_session_factory(creator=db_settings.creator, url=db_settings.url) experiment = _load_experiment(name, decoder=db_settings.decoder) if not isinstance(experiment, Experiment) or experiment.is_simple_experiment: raise ValueError("Service API only supports Experiment") return experiment
[docs]def save_experiment(experiment: Experiment, db_settings: DBSettings) -> None: """ Save experiment to db. Args: experiment: `Experiment` object. db_settings: Defines behavior for loading/saving experiment to/from db. """ init_engine_and_session_factory(creator=db_settings.creator, url=db_settings.url) _save_experiment(experiment, encoder=db_settings.encoder)
[docs]def load_experiment_and_generation_strategy( experiment_name: str, db_settings: DBSettings ) -> Tuple[Experiment, GenerationStrategy]: """Load experiment and the corresponding generation strategy from the DB. Args: name: Experiment name. db_settings: Defines behavior for loading/saving experiment to/from db. Returns: A tuple of the loaded experiment and generation strategy. """ experiment = load_experiment(name=experiment_name, db_settings=db_settings) generation_strategy = _load_generation_strategy_by_experiment_name( experiment_name=experiment_name, decoder=db_settings.decoder ) return experiment, generation_strategy
[docs]def save_experiment_and_generation_strategy( experiment: Experiment, generation_strategy: GenerationStrategy, db_settings: DBSettings, ) -> None: """Save experiment and generation strategy to DB. Args: experiment: `Experiment` object. generation_strategy: Corresponding generation strategy. db_settings: Defines behavior for loading/saving experiment to/from db. """ save_experiment(experiment=experiment, db_settings=db_settings) _save_generation_strategy( generation_strategy=generation_strategy, encoder=db_settings.encoder )