Source code for ax.storage.sqa_store.reduced_state
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
from ax.storage.sqa_store.sqa_classes import SQAGeneratorRun
from sqlalchemy.orm import defaultload, lazyload, strategy_options
from sqlalchemy.orm.attributes import InstrumentedAttribute
GR_LARGE_MODEL_ATTRS: list[InstrumentedAttribute] = [ # pyre-ignore[9]
SQAGeneratorRun.model_kwargs,
SQAGeneratorRun.bridge_kwargs,
SQAGeneratorRun.model_state_after_gen,
SQAGeneratorRun.gen_metadata,
]
GR_PARAMS_METRICS_COLS = [
"parameters",
"parameter_constraints",
"metrics",
]
[docs]
def get_query_options_to_defer_immutable_duplicates() -> list[strategy_options.Load]:
"""Returns the query options that defer loading of attributes that are duplicated
on each trial (like search space attributes and metrics). These attributes do not
need to be loaded for experiments with immutable search space and optimization
configuration.
"""
options = [lazyload(f"generator_runs.{col}") for col in GR_PARAMS_METRICS_COLS]
return options
[docs]
def get_query_options_to_defer_large_model_cols() -> list[strategy_options.Load]:
"""Returns the query options that defer loading of model-state-related columns
of generator runs, which can be large and are not needed on every generator run
when loading experiment and generation strategy in reduced state.
"""
return [
defaultload("generator_runs").defer(col.key) for col in GR_LARGE_MODEL_ATTRS
]