Skip to content

Improve Auxillary Exp check #4824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mgarrard wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Improve Auxillary Exp check #4824

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
af94693
Validated uniqueness of TC repr method (#4787)
mgarrard Jan 27, 2026
76b40cf
Replace data fetch with lookup (#4788)
mgarrard Jan 27, 2026
0345779
Increase restritiveness of data requirements for TrialBasedCriteria (…
mgarrard Jan 27, 2026
dfd73e3
Make `transition_to` non-optional for TC
mgarrard Jan 27, 2026
aa057b4
Validated generatorspec unique id
mgarrard Jan 27, 2026
542296f
Remove completion criteria
mgarrard Jan 27, 2026
8df87e1
Improve Auxillary Exp check
mgarrard Jan 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ax/core/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -528,7 +528,7 @@ def get_target_trial_index(
return None

# trial indices that have data for required metrics
trial_indices_with_required_metrics = _get_trial_indices_with_required_metrics(
trial_indices_with_required_metrics = get_trial_indices_with_required_metrics(
experiment=experiment,
df=df,
require_data_for_all_metrics=require_data_for_all_metrics,
Expand Down Expand Up @@ -593,7 +593,7 @@ def get_target_trial_index(
return None


def _get_trial_indices_with_required_metrics(
def get_trial_indices_with_required_metrics(
experiment: Experiment,
df: "pd.DataFrame",
require_data_for_all_metrics: bool,
Expand Down
19 changes: 7 additions & 12 deletions ax/generation_strategy/generation_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -992,8 +992,6 @@ class GenerationStep:
generator_gen_kwargs: Each call to `generation_strategy.gen` performs a call
to the step's adapter's `gen` under the hood; `generator_gen_kwargs` will be
passed to the adapter's `gen` like: `adapter.gen(**generator_gen_kwargs)`.
completion_criteria: List of TransitionCriterion. All `is_met` must evaluate
True for the GenerationStrategy to move on to the next Step
index: Index of this generation step, for use internally in `Generation
Strategy`. Do not assign as it will be reassigned when instantiating
`GenerationStrategy` with a list of its steps.
Expand Down Expand Up @@ -1024,7 +1022,6 @@ def __new__(
num_trials: int,
generator_kwargs: dict[str, Any] | None = None,
generator_gen_kwargs: dict[str, Any] | None = None,
completion_criteria: Sequence[TransitionCriterion] | None = None,
min_trials_observed: int = 0,
max_parallelism: int | None = None,
enforce_num_trials: bool = True,
Expand Down Expand Up @@ -1061,7 +1058,6 @@ def __new__(

generator_kwargs = generator_kwargs or {}
generator_gen_kwargs = generator_gen_kwargs or {}
completion_criteria = completion_criteria or []

if (
enforce_num_trials
Expand Down Expand Up @@ -1101,47 +1097,46 @@ def __new__(
# is set in `GenerationStrategy` constructor, because only then is the order
# of the generation steps actually known.
transition_criteria: list[TransitionCriterion] = []
# Placeholder - will be overwritten in _validate_and_set_step_sequence in GS
placeholder_transition_to = f"GenerationStep_{str(index)}"

if num_trials != -1:
transition_criteria.append(
MinTrials(
threshold=num_trials,
transition_to=placeholder_transition_to,
not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
block_gen_if_met=enforce_num_trials,
block_transition_if_unmet=True,
use_all_trials_in_exp=use_all_trials_in_exp,
transition_to=None, # Re-set in GS constructor.
)
)

if min_trials_observed > 0:
transition_criteria.append(
MinTrials(
threshold=min_trials_observed,
transition_to=placeholder_transition_to,
only_in_statuses=[
TrialStatus.COMPLETED,
TrialStatus.EARLY_STOPPED,
],
threshold=min_trials_observed,
block_gen_if_met=False,
block_transition_if_unmet=True,
use_all_trials_in_exp=use_all_trials_in_exp,
transition_to=None, # Re-set in GS constructor.
)
)
if max_parallelism is not None:
transition_criteria.append(
MaxGenerationParallelism(
threshold=max_parallelism,
transition_to=placeholder_transition_to,
only_in_statuses=[TrialStatus.RUNNING],
block_gen_if_met=True,
block_transition_if_unmet=False,
# MaxParallelism transitions to self,
# this will be confirmed in GS init
transition_to=f"GenerationStep_{str(index)}",
)
)

transition_criteria += list(completion_criteria)

# Create and return a GenerationNode instance
node = GenerationNode(
# NOTE: This name is a placeholder that will be overwritten in
Expand Down
1 change: 0 additions & 1 deletion ax/generation_strategy/generator_spec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -382,5 +382,4 @@ def __eq__(self, other: GeneratorSpec) -> bool:
@property
def _unique_id(self) -> str:
"""Returns the unique ID of this model spec"""
# TODO @mgarrard verify that this is unique enough
return str(hash(self))
170 changes: 0 additions & 170 deletions ax/generation_strategy/tests/test_aepsych_criterion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,170 +0,0 @@
# 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 unittest.mock import patch

import pandas as pd
from ax.adapter.registry import Generators
from ax.core.data import Data
from ax.core.trial_status import TrialStatus
from ax.generation_strategy.generation_strategy import (
GenerationStep,
GenerationStrategy,
)
from ax.generation_strategy.transition_criterion import (
MinimumPreferenceOccurances,
MinTrials,
)
from ax.utils.common.testutils import TestCase
from ax.utils.testing.core_stubs import get_experiment


class TestAEPsychCriterion(TestCase):
"""
This test class tests the AEPsych usecase implementation. Previously, AEPsych
used `CompletionCriterion` to determine when to move to the next generation.
However, `CompletionCriterion` is deprecated and replaced by `TransitionCriterion`.
We still want to test the bespoke TransitionCriterion used by AEPsych
"""

def test_single_criterion(self) -> None:
criterion = MinimumPreferenceOccurances(metric_signature="m1", threshold=3)

experiment = get_experiment()

generation_strategy = GenerationStrategy(
name="SOBOL+MBM::default",
steps=[
GenerationStep(
generator=Generators.SOBOL,
num_trials=-1,
completion_criteria=[criterion],
),
GenerationStep(
generator=Generators.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
),
],
)
generation_strategy.experiment = experiment

# Has not seen enough of each preference
self.assertFalse(
generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
)
# check the transition_to is being set
self.assertEqual(
generation_strategy._curr.transition_criteria[0].transition_to,
"GenerationStep_1_BoTorch",
)

data = Data(
df=pd.DataFrame(
{
"trial_index": range(6),
"arm_name": [f"{i}_0" for i in range(6)],
"metric_name": ["m1" for _ in range(6)],
"mean": [0, 0, 0, 1, 1, 1],
"sem": [0 for _ in range(6)],
"metric_signature": ["m1" for _ in range(6)],
}
)
)
with patch.object(experiment, "fetch_data", return_value=data):
move_to_next_node = generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
# We have seen three "yes" and three "no"
self.assertTrue(move_to_next_node)

self.assertEqual(
generation_strategy._curr.generator_spec_to_gen_from.generator_enum,
Generators.BOTORCH_MODULAR,
)

def test_many_criteria(self) -> None:
criteria = [
MinimumPreferenceOccurances(metric_signature="m1", threshold=3),
MinTrials(only_in_statuses=[TrialStatus.COMPLETED], threshold=5),
]

experiment = get_experiment()

generation_strategy = GenerationStrategy(
name="SOBOL+MBM::default",
steps=[
GenerationStep(
generator=Generators.SOBOL,
num_trials=-1,
completion_criteria=criteria,
),
GenerationStep(
generator=Generators.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
),
],
)
generation_strategy.experiment = experiment

# Has not seen enough of each preference
self.assertFalse(
generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
)

data = Data(
df=pd.DataFrame(
{
"trial_index": range(6),
"arm_name": [f"{i}_0" for i in range(6)],
"metric_name": ["m1" for _ in range(6)],
"mean": [0, 0, 0, 1, 1, 1],
"sem": [0 for _ in range(6)],
"metric_signature": ["m1" for _ in range(6)],
}
)
)
with patch.object(experiment, "fetch_data", return_value=data):
move_to_next_node = generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
# We have seen three "yes" and three "no", but not enough trials are
# completed
self.assertFalse(move_to_next_node)

for _i in range(6):
experiment.new_trial(
generation_strategy.gen_single_trial(experiment=experiment)
)
for trial in experiment.trials.values():
trial._status = TrialStatus.COMPLETED

# Enough trials are completed but we have not seen three "yes" and three
# "no"
self.assertFalse(
generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
)

with patch.object(experiment, "fetch_data", return_value=data):
move_to_next_node = generation_strategy._maybe_transition_to_next_node(
raise_data_required_error=False
)
# Enough trials are completed but we have not seen three "yes" and three
# "no"
self.assertTrue(move_to_next_node)

self.assertEqual(
generation_strategy._curr.generator_spec_to_gen_from.generator_enum,
Generators.BOTORCH_MODULAR,
)
13 changes: 10 additions & 3 deletions ax/generation_strategy/tests/test_generation_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -321,7 +321,11 @@ def test_node_string_representation(self) -> None:
self.mbm_generator_spec,
],
transition_criteria=[
MinTrials(threshold=5, only_in_statuses=[TrialStatus.RUNNING])
MinTrials(
threshold=5,
transition_to="next_node",
only_in_statuses=[TrialStatus.RUNNING],
)
],
)
string_rep = str(node)
Expand All @@ -331,7 +335,7 @@ def test_node_string_representation(self) -> None:
"GenerationNode(name='test', "
"generator_specs=[GeneratorSpec(generator_enum=BoTorch, "
"generator_key_override=None)], "
"transition_criteria=[MinTrials(transition_to='None')])",
"transition_criteria=[MinTrials(transition_to='next_node')])",
)

def test_single_fixed_features(self) -> None:
Expand Down Expand Up @@ -439,6 +443,7 @@ def test_init(self) -> None:
[
MinTrials(
threshold=5,
transition_to="GenerationStep_-1", # overwritten during GS init
not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
block_gen_if_met=True,
block_transition_if_unmet=True,
Expand All @@ -464,17 +469,19 @@ def test_init(self) -> None:
[
MinTrials(
threshold=5,
transition_to="GenerationStep_-1", # overwritten during GS init
not_in_statuses=[TrialStatus.FAILED, TrialStatus.ABANDONED],
block_gen_if_met=False,
block_transition_if_unmet=True,
use_all_trials_in_exp=True,
),
MinTrials(
threshold=3,
transition_to="GenerationStep_-1", # overwritten during GS init
only_in_statuses=[
TrialStatus.COMPLETED,
TrialStatus.EARLY_STOPPED,
],
threshold=3,
block_gen_if_met=False,
block_transition_if_unmet=True,
use_all_trials_in_exp=True,
Expand Down
Loading