Source code for ax.modelbridge.generation_strategy

#!/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 __future__ import annotations

from copy import deepcopy
from functools import wraps
from logging import Logger
from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar

import pandas as pd
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.generation_strategy_interface import GenerationStrategyInterface
from ax.core.generator_run import GeneratorRun
from ax.core.observation import ObservationFeatures
from ax.core.utils import (
    extend_pending_observations,
    get_pending_observation_features_based_on_trial_status,
)
from ax.exceptions.core import DataRequiredError, UnsupportedError, UserInputError
from ax.exceptions.generation_strategy import (
    GenerationStrategyCompleted,
    GenerationStrategyMisconfiguredException,
)
from ax.modelbridge.base import ModelBridge
from ax.modelbridge.generation_node import GenerationNode, GenerationStep
from ax.modelbridge.model_spec import FactoryFunctionModelSpec
from ax.modelbridge.modelbridge_utils import get_fixed_features_from_experiment
from ax.modelbridge.registry import _extract_model_state_after_gen, ModelRegistryBase
from ax.modelbridge.transition_criterion import TrialBasedCriterion
from ax.utils.common.logger import _round_floats_for_logging, get_logger
from ax.utils.common.typeutils import checked_cast, checked_cast_list, not_none
from pyre_extensions import none_throws

logger: Logger = get_logger(__name__)


MAX_CONDITIONS_GENERATED = 10000
MAX_GEN_DRAWS = 5
MAX_GEN_DRAWS_EXCEEDED_MESSAGE = (
    f"GenerationStrategy exceeded `MAX_GEN_DRAWS` of {MAX_GEN_DRAWS} while trying to "
    "generate a unique parameterization. This indicates that the search space has "
    "likely been fully explored, or that the sweep has converged."
)
T = TypeVar("T")


[docs]def step_based_gs_only(f: Callable[..., T]) -> Callable[..., T]: """ For use as a decorator on functions only implemented for ``GenerationStep``-based ``GenerationStrategies``. Mainly useful for older ``GenerationStrategies``. """ @wraps(f) def impl( self: "GenerationStrategy", *args: List[Any], **kwargs: Dict[str, Any] ) -> T: if self.is_node_based: raise UnsupportedError( f"{f.__name__} is not supported for GenerationNode based" " GenerationStrategies." ) return f(self, *args, **kwargs) return impl
[docs]class GenerationStrategy(GenerationStrategyInterface): """GenerationStrategy describes which model should be used to generate new points for which trials, enabling and automating use of different models throughout the optimization process. For instance, it allows to use one model for the initialization trials, and another one for all subsequent trials. In the general case, this allows to automate use of an arbitrary number of models to generate an arbitrary numbers of trials described in the `trials_per_model` argument. Args: nodes: A list of `GenerationNode`. Each `GenerationNode` in the list represents a single node in a `GenerationStrategy` which, when composed of `GenerationNodes`, can be conceptualized as a graph instead of a linear list. `TransitionCriterion` defined in each `GenerationNode` represent the edges in the `GenerationStrategy` graph. `GenerationNodes` are more flexible than `GenerationSteps` and new `GenerationStrategies` should use nodes. Notably, either, but not both, of `nodes` and `steps` must be provided. steps: A list of `GenerationStep` describing steps of this strategy. name: An optional name for this generaiton strategy. If not specified, strategy's name will be names of its nodes' models joined with '+'. """ _nodes: List[GenerationNode] _curr: GenerationNode # Current node in the strategy. # Whether all models in this GS are in Models registry enum. _uses_registered_models: bool # All generator runs created through this generation strategy, in chronological # order. _generator_runs: List[GeneratorRun] # Experiment, for which this generation strategy has generated trials, if # it exists. _experiment: Optional[Experiment] = None # Trial indices as last seen by the model; updated in `_model` property setter. _model: Optional[ModelBridge] = None # Current model. def __init__( self, steps: Optional[List[GenerationStep]] = None, name: Optional[str] = None, nodes: Optional[List[GenerationNode]] = None, ) -> None: # Validate that one and only one of steps or nodes is provided if not ((steps is None) ^ (nodes is None)): raise GenerationStrategyMisconfiguredException( error_info="GenerationStrategy must contain either steps or nodes." ) # pyre-ignore[8] self._nodes = none_throws(nodes if steps is None else steps) # Validate correctness of steps list or nodes graph if isinstance(steps, list) and all( isinstance(s, GenerationStep) for s in steps ): self._validate_and_set_step_sequence(steps=self._nodes) elif isinstance(nodes, list) and self.is_node_based: self._validate_and_set_node_graph(nodes=nodes) else: # TODO[mgarrard]: Allow mix of nodes and steps raise GenerationStrategyMisconfiguredException( "`GenerationStrategy` inputs are:\n" "`steps` (list of `GenerationStep`) or\n" "`nodes` (list of `GenerationNode`)." ) # Log warning if the GS uses a non-registered (factory function) model. self._uses_registered_models = not any( isinstance(ms, FactoryFunctionModelSpec) for node in self._nodes for ms in node.model_specs ) if not self._uses_registered_models: logger.info( "Using model via callable function, " "so optimization is not resumable if interrupted." ) self._generator_runs = [] # Set name to an explicit value ahead of time to avoid # adding properties during equality checks super().__init__(name=name or self._make_default_name()) @property def is_node_based(self) -> bool: """Whether this strategy consists of GenerationNodes only. This is useful for determining initialization properties and other logic. """ return not any(isinstance(n, GenerationStep) for n in self._nodes) and all( isinstance(n, GenerationNode) for n in self._nodes ) @property def name(self) -> str: """Name of this generation strategy. Defaults to a combination of model names provided in generation steps, set at the time of the ``GenerationStrategy`` creation. """ return self._name @name.setter def name(self, name: str) -> None: """Set generation strategy name.""" self._name = name @property @step_based_gs_only def model_transitions(self) -> List[int]: """[DEPRECATED]List of trial indices where a transition happened from one model to another. """ raise DeprecationWarning( "`model_transitions` is no longer supported. Please refer to `model_key` " "field on generator runs for similar information if needed." ) @property def current_step(self) -> GenerationStep: """Current generation step.""" if not isinstance(self._curr, GenerationStep): raise TypeError( "The current object is not a GenerationStep, you may be looking " "for the current_node property." ) return self._curr @property def current_node(self) -> GenerationNode: """Current generation node.""" if not isinstance(self._curr, GenerationNode): raise TypeError( "The current object is not a GenerationNode, you may be looking for the" " current_step property." ) return self._curr @property def current_node_name(self) -> str: """Current generation node name.""" return self._curr.node_name @property @step_based_gs_only def current_step_index(self) -> int: """Returns the index of the current generation step. This attribute is replaced by node_name in newer GenerationStrategies but surfaced here for backward compatibility. """ node_names_for_all_steps = [step._node_name for step in self._nodes] assert ( self._curr.node_name in node_names_for_all_steps ), "The current step is not found in the list of steps" return node_names_for_all_steps.index(self._curr.node_name) @property def model(self) -> Optional[ModelBridge]: """Current model in this strategy. Returns None if no model has been set yet (i.e., if no generator runs have been produced from this GS). """ return self._curr._fitted_model @property def experiment(self) -> Experiment: """Experiment, currently set on this generation strategy.""" if self._experiment is None: raise ValueError("No experiment set on generation strategy.") return not_none(self._experiment) @experiment.setter def experiment(self, experiment: Experiment) -> None: """If there is an experiment set on this generation strategy as the experiment it has been generating generator runs for, check if the experiment passed in is the same as the one saved and log an information statement if its not. Set the new experiment on this generation strategy. """ if self._experiment is None or experiment._name == self.experiment._name: self._experiment = experiment else: raise ValueError( "This generation strategy has been used for experiment " f"{self.experiment._name} so far; cannot reset experiment" f" to {experiment._name}. If this is a new optimization, " "a new generation strategy should be created instead." ) @property def last_generator_run(self) -> Optional[GeneratorRun]: """Latest generator run produced by this generation strategy. Returns None if no generator runs have been produced yet. """ # Used to restore current model when decoding a serialized GS. return self._generator_runs[-1] if self._generator_runs else None @property def uses_non_registered_models(self) -> bool: """Whether this generation strategy involves models that are not registered and therefore cannot be stored.""" return not self._uses_registered_models @property def trials_as_df(self) -> Optional[pd.DataFrame]: """Puts information on individual trials into a data frame for easy viewing. For example for a GenerationStrategy composed of GenerationSteps: Gen. Step | Model | Trial Index | Trial Status | Arm Parameterizations 0 | Sobol | 0 | RUNNING | {"0_0":{"x":9.17...}} """ logger.info( "Note that parameter values in dataframe are rounded to 2 decimal " "points; the values in the dataframe are thus not the exact ones " "suggested by Ax in trials." ) if self._experiment is None or all( len(step.trials_from_node) == 0 for step in self._nodes ): return None step_or_node_col = ( "Generation Node" if self.is_node_based else "Generation Step" ) records = [ { step_or_node_col: node.node_name, "Generation Model": self._nodes[ node_idx ].model_spec_to_gen_from.model_key, "Trial Index": trial_idx, "Trial Status": self.experiment.trials[trial_idx].status.name, "Arm Parameterizations": { arm.name: _round_floats_for_logging(arm.parameters) for arm in self.experiment.trials[trial_idx].arms }, } for node_idx, node in enumerate(self._nodes) for trial_idx in node.trials_from_node ] return pd.DataFrame.from_records(records).reindex( columns=[ step_or_node_col, "Generation Model", "Trial Index", "Trial Status", "Arm Parameterizations", ] ) @property def optimization_complete(self) -> bool: """Checks whether all nodes are completed in the generation strategy.""" return all(node.is_completed for node in self._nodes) @property @step_based_gs_only def _steps(self) -> List[GenerationStep]: """List of generation steps.""" return self._nodes # pyre-ignore[7]
[docs] def gen( self, experiment: Experiment, data: Optional[Data] = None, n: int = 1, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, **kwargs: Any, ) -> GeneratorRun: """Produce the next points in the experiment. Additional kwargs passed to this method are propagated directly to the underlying model's `gen`, along with the `model_gen_kwargs` set on the current generation node. NOTE: Each generator run returned from this function must become a single trial on the experiment to comply with assumptions made in generation strategy. Do not split one generator run produced from generation strategy into multiple trials (never making a generator run into a trial is allowed). Args: experiment: Experiment, for which the generation strategy is producing a new generator run in the course of `gen`, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method. data: Optional data to be passed to the underlying model's `gen`, which is called within this method and actually produces the resulting generator run. By default, data is all data on the `experiment`. n: Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the `n` and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from `n`. pending_observations: A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated. """ return self._gen_multiple( experiment=experiment, num_generator_runs=1, data=data, n=n, pending_observations=pending_observations, **kwargs, )[0]
[docs] def gen_for_multiple_trials_with_multiple_models( self, experiment: Experiment, num_generator_runs: int, data: Optional[Data] = None, n: int = 1, ) -> List[List[GeneratorRun]]: """Produce GeneratorRuns for multiple trials at once with the possibility of ensembling, or using multiple models per trial, getting multiple GeneratorRuns per trial. NOTE: This method is in development. Please do not use it yet. Args: experiment: Experiment, for which the generation strategy is producing a new generator run in the course of `gen`, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method. data: Optional data to be passed to the underlying model's `gen`, which is called within this method and actually produces the resulting generator run. By default, data is all data on the `experiment`. n: Integer representing how many trials should be in the generator run produced by this method. NOTE: Some underlying models may ignore the ``n`` and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from ``n``. Returns: A list of lists of lists generator runs. Each outer list represents a trial being suggested and each inner list represents a generator run for that trial. """ grs = self._gen_multiple( experiment=experiment, num_generator_runs=num_generator_runs, data=data, n=n, pending_observations=get_pending_observation_features_based_on_trial_status( experiment=experiment ), fixed_features=get_fixed_features_from_experiment(experiment=experiment), ) return [[gr] for gr in grs]
[docs] def current_generator_run_limit( self, ) -> Tuple[int, bool]: """First check if we can move the generation strategy to the next node, which is safe, as the next call to ``gen`` will just pick up from there. Then determine how many generator runs this generation strategy can generate right now, assuming each one of them becomes its own trial, and whether optimization is completed. Returns: a two-item tuple of: - the number of generator runs that can currently be produced, with -1 meaning unlimited generator runs, - whether optimization is completed and the generation strategy cannot generate any more generator runs at all. """ try: self._maybe_move_to_next_step(raise_data_required_error=False) except GenerationStrategyCompleted: return 0, True # if the generation strategy is not complete, optimization is not complete return self._curr.generator_run_limit(), False
[docs] def clone_reset(self) -> GenerationStrategy: """Copy this generation strategy without it's state.""" cloned_nodes = deepcopy(self._nodes) for n in cloned_nodes: # Unset the generation strategy back-pointer, so the nodes are not # associated with any generation strategy. n._generation_strategy = None if self.is_node_based: return GenerationStrategy(name=self.name, nodes=cloned_nodes) return GenerationStrategy( name=self.name, steps=checked_cast_list(GenerationStep, cloned_nodes) )
def _unset_non_persistent_state_fields(self) -> None: """Utility for testing convenience: unset fields of generation strategy that are set during candidate generation; these fields are not persisted during storage. To compare a pre-storage and a reloaded generation strategies; call this utility on the pre-storage one first. The rest of the fields should be identical. """ self._model = None for s in self._nodes: s._model_spec_to_gen_from = None @step_based_gs_only def _validate_and_set_step_sequence(self, steps: List[GenerationStep]) -> None: """Initialize and validate the steps provided to this GenerationStrategy. Some GenerationStrategies are composed of GenerationStep objects, but we also need to initialize the correct GenerationNode representation for these steps. This function validates: 1. That only the last step has num_trials=-1, which indicates unlimited trial generation is possible. 2. That each step's num_trials attrivute is either positive or -1 3. That each step's max_parallelism attribute is either None or positive It then sets the corect TransitionCriterion and node_name attributes on the underlying GenerationNode objects. """ for idx, step in enumerate(steps): if step.num_trials == -1 and len(step.completion_criteria) < 1: if idx < len(self._steps) - 1: raise UserInputError( "Only last step in generation strategy can have " "`num_trials` set to -1 to indicate that the model in " "the step shouldbe used to generate new trials " "indefinitely unless completion critera present." ) elif step.num_trials < 1 and step.num_trials != -1: raise UserInputError( "`num_trials` must be positive or -1 (indicating unlimited) " "for all generation steps." ) if step.max_parallelism is not None and step.max_parallelism < 1: raise UserInputError( "Maximum parallelism should be None (if no limit) or " f"a positive number. Got: {step.max_parallelism} for " f"step {step.model_name}." ) step._node_name = f"GenerationStep_{str(idx)}" step.index = idx # Set transition_to field for all but the last step, which remains # null. if idx != len(self._steps): for transition_criteria in step.transition_criteria: if ( transition_criteria.criterion_class != "MaxGenerationParallelism" ): transition_criteria._transition_to = ( f"GenerationStep_{str(idx + 1)}" ) step._generation_strategy = self self._curr = steps[0] def _validate_and_set_node_graph(self, nodes: List[GenerationNode]) -> None: """Initialize and validate the node graph provided to this GenerationStrategy. This function validates: 1. That all nodes have unique names. 2. That there is at least one node with a transition_to field. 3. That all `transition_to` attributes on a TransitionCriterion point to another node in the same GenerationStrategy. 4. Warns if no nodes contain a transition criterion """ node_names = [] for node in self._nodes: # validate that all node names are unique if node.node_name in node_names: raise GenerationStrategyMisconfiguredException( error_info="All node names in a GenerationStrategy " + "must be unique." ) node_names.append(node.node_name) node._generation_strategy = self # validate `transition_criterion` contains_a_transition_to_argument = False for node in self._nodes: for transition_criteria in node.transition_criteria: if transition_criteria.transition_to is not None: contains_a_transition_to_argument = True if transition_criteria.transition_to not in node_names: raise GenerationStrategyMisconfiguredException( error_info=f"`transition_to` argument " f"{transition_criteria.transition_to} does not " "correspond to any node in this GenerationStrategy." ) # validate that at least one node has transition_to field if len(self._nodes) > 1 and not contains_a_transition_to_argument: logger.warning( "None of the nodes in this GenerationStrategy " "contain a `transition_to` argument in their transition_criteria. " "Therefore, the GenerationStrategy will not be able to " "move from one node to another. Please add a " "`transition_to` argument." ) self._curr = nodes[0] @step_based_gs_only def _step_repr(self, step_str_rep: str) -> str: """Return the string representation of the steps in a GenerationStrategy composed of GenerationSteps. """ step_str_rep += "steps=[" remaining_trials = "subsequent" if len(self._nodes) > 1 else "all" for step in self._nodes: num_trials = remaining_trials for criterion in step.transition_criteria: if criterion.criterion_class == "MaxTrials" and isinstance( criterion, TrialBasedCriterion ): num_trials = criterion.threshold try: model_name = step.model_spec_to_gen_from.model_key except TypeError: model_name = "model with unknown name" step_str_rep += f"{model_name} for {num_trials} trials, " step_str_rep = step_str_rep[:-2] step_str_rep += "])" return step_str_rep def _make_default_name(self) -> str: """Make a default name for this generation strategy; used when no name is passed to the constructor. For node-based generation strategies, the name is constructed by joining together the names of the nodes set on this generation strategy. For step-based generation strategies, the model keys of the underlying model specs are used. Note: This should only be called once the nodes are set. """ if not self._nodes: raise UnsupportedError( "Cannot make a default name for a generation strategy with no nodes " "set yet." ) # TODO: Simplify this after updating GStep names to represent underlying models. if self.is_node_based: node_names = (node.node_name for node in self._nodes) else: node_names = (node.model_spec_to_gen_from.model_key for node in self._nodes) # Trim the "get_" beginning of the factory function if it's there. node_names = (n[4:] if n[:4] == "get_" else n for n in node_names) return "+".join(node_names) def __repr__(self) -> str: """String representation of this generation strategy.""" gs_str = f"GenerationStrategy(name='{self.name}', " if not self.is_node_based: return self._step_repr(gs_str) gs_str += f"nodes={str(self._nodes)})" return gs_str # ------------------------- Candidate generation helpers. ------------------------- def _gen_multiple( self, experiment: Experiment, num_generator_runs: int, data: Optional[Data] = None, n: int = 1, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, **model_gen_kwargs: Any, ) -> List[GeneratorRun]: """Produce multiple generator runs at once, to be made into multiple trials on the experiment. NOTE: This is used to ensure that maximum paralellism and number of trials per node are not violated when producing many generator runs from this generation strategy in a row. Without this function, if one generates multiple generator runs without first making any of them into running trials, generation strategy cannot enforce that it only produces as many generator runs as are allowed by the paralellism limit and the limit on number of trials in current node. Args: experiment: Experiment, for which the generation strategy is producing a new generator run in the course of `gen`, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method. data: Optional data to be passed to the underlying model's `gen`, which is called within this method and actually produces the resulting generator run. By default, data is all data on the `experiment`. n: Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the ``n`` and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from ``n``. pending_observations: A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated. model_gen_kwargs: Keyword arguments that are passed through to ``GenerationNode.gen``, which will pass them through to ``ModelSpec.gen``, which will pass them to ``ModelBridge.gen``. """ self.experiment = experiment self._maybe_move_to_next_step() self._fit_current_model(data=data) # Get GeneratorRun limit that respects the node's transition criterion that # affect the number of generator runs that can be produced. gr_limit = self._curr.generator_run_limit(supress_generation_errors=False) if gr_limit == -1: num_generator_runs = max(num_generator_runs, 1) else: num_generator_runs = max(min(num_generator_runs, gr_limit), 1) generator_runs = [] pending_observations = deepcopy(pending_observations) or {} for _ in range(num_generator_runs): try: generator_run = self._curr.gen( n=n, pending_observations=pending_observations, arms_by_signature_for_deduplication=experiment.arms_by_signature, **model_gen_kwargs, ) except DataRequiredError as err: # Model needs more data, so we log the error and return # as many generator runs as we were able to produce, unless # no trials were produced at all (in which case its safe to raise). if len(generator_runs) == 0: raise logger.debug(f"Model required more data: {err}.") break self._generator_runs.append(generator_run) generator_runs.append(generator_run) # Extend the `pending_observation` with newly generated point(s) # to avoid repeating them. extend_pending_observations( experiment=experiment, pending_observations=pending_observations, generator_run=generator_run, ) return generator_runs # ------------------------- Model selection logic helpers. ------------------------- def _fit_current_model(self, data: Optional[Data]) -> None: """Fits or update the model on the current generation node (does not move between generation nodes). Args: data: Optional ``Data`` to fit or update with; if not specified, generation strategy will obtain the data via ``experiment.lookup_data``. """ data = self.experiment.lookup_data() if data is None else data # If last generator run's index matches the current node, extract # model state from last generator run and pass it to the model # being instantiated in this function. model_state_on_lgr = self._get_model_state_from_last_generator_run() if not data.df.empty: trial_indices_in_data = sorted(data.df["trial_index"].unique()) logger.debug(f"Fitting model with data for trials: {trial_indices_in_data}") self._curr.fit(experiment=self.experiment, data=data, **model_state_on_lgr) self._model = self._curr._fitted_model def _maybe_move_to_next_step(self, raise_data_required_error: bool = True) -> bool: """Moves this generation strategy to next node if the current node is completed, and it is not the last node in this generation strategy. This method is safe to use both when generating candidates or simply checking how many generator runs (to be made into trials) can currently be produced. NOTE: this method raises ``GenerationStrategyCompleted`` error if the current generation node is complete, but it is also the last in generation strategy. Args: raise_data_required_error: Whether to raise ``DataRequiredError`` in the maybe_step_completed method in GenerationNode class. Returns: Whether generation strategy moved to the next node. """ move_to_next_node, next_node = self._curr.should_transition_to_next_node( raise_data_required_error=raise_data_required_error ) if move_to_next_node: if self.optimization_complete: raise GenerationStrategyCompleted( f"Generation strategy {self} generated all the trials as " "specified in its nodes." ) if next_node is None: # If the last node did not specify which node to transition to, # move to the next node in the list. current_node_index = self._nodes.index(self._curr) next_node = self._nodes[current_node_index + 1].node_name for node in self._nodes: if node.node_name == next_node: self._curr = node # Moving to the next node also entails unsetting this GS's model # (since new node's model will be initialized for the first time; # this is done in `self._fit_current_model). self._model = None return move_to_next_node def _get_model_state_from_last_generator_run(self) -> Dict[str, Any]: lgr = self.last_generator_run # NOTE: This will not be easily compatible with `GenerationNode`; # will likely need to find last generator run per model. Not a problem # for now though as GS only allows `GenerationStep`-s for now. # Potential solution: store generator runs on `GenerationNode`-s and # split them per-model there. model_state_on_lgr = {} # Need to check if model_spec_to_gen_from is none to account for # ExternalGenerationNodes which leverage models from outside Ax. model_on_curr = ( self._curr.model_spec_to_gen_from.model_enum if self._curr.model_spec_to_gen_from else None ) if lgr is None: return model_state_on_lgr if not self.is_node_based: grs_equal = lgr._generation_step_index == self.current_step_index else: grs_equal = lgr._generation_node_name == self._curr.node_name if grs_equal and lgr._model_state_after_gen: if self.model or isinstance(model_on_curr, ModelRegistryBase): # TODO[drfreund]: Consider moving this to `GenerationStep` or # `GenerationNode`. model_cls = ( self.model.model.__class__ if self.model is not None # NOTE: This checked cast is save per the OR-statement in last line # of the IF-check above. else checked_cast(ModelRegistryBase, model_on_curr).model_class ) model_state_on_lgr = _extract_model_state_after_gen( generator_run=lgr, model_class=model_cls, ) return model_state_on_lgr