From c08f30bb04c74809b1e764e5af5e5148c0108089 Mon Sep 17 00:00:00 2001
From: Mia Garrard <mgarrard@meta.com>
Date: Mon, 29 Jan 2024 12:36:09 -0800
Subject: [PATCH] Update all instances of MinimumPreferenceOccurances to
 directly use TransitionCriterion (#2135)

Summary:
Pull Request resolved: https://github.com/facebook/Ax/pull/2135

X-link: https://github.com/facebookresearch/aepsych/pull/327

This diff updates MinimumPreferenceOccurances to directly inherit from its source in TransitionCriterion file

In following diffs we will:
- Completely remove the completion criterion file
- update all four completion criterion defined in aepsych code here: https://www.internalfb.com/code/fbsource/[409e3dfb01ec5c613d34e58c491d63e8051d10d9]/fbcode/frl/ae/aepsych/tests/generators/test_completion_criteria.py?lines=12-15
- revisit storage
- remove all todos in gennode, genstrat, and transitioncriterion classes related to maintaining this deprecated code
- update AEPsych GSs as needed
- determine if run indefinetly can be replaced by simply having gen_unlimited_trials = true

Reviewed By: lena-kashtelyan

Differential Revision: D52852317

fbshipit-source-id: 357ae081a0d2ceb83bc7376cb9144141f6b13f58
---
 ax/modelbridge/completion_criterion.py         | 18 +-----------------
 .../tests/test_completion_criterion.py         |  3 +--
 .../tests/test_transition_criterion.py         |  3 ++-
 ax/modelbridge/transition_criterion.py         |  2 ++
 4 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/ax/modelbridge/completion_criterion.py b/ax/modelbridge/completion_criterion.py
index 902b4cbbcec..8e2dee001b2 100644
--- a/ax/modelbridge/completion_criterion.py
+++ b/ax/modelbridge/completion_criterion.py
@@ -5,10 +5,7 @@
 
 from logging import Logger
 
-from ax.modelbridge.transition_criterion import (
-    MinimumPreferenceOccurances,
-    TransitionCriterion,
-)
+from ax.modelbridge.transition_criterion import TransitionCriterion
 from ax.utils.common.logger import get_logger
 
 logger: Logger = get_logger(__name__)
@@ -24,16 +21,3 @@ class CompletionCriterion(TransitionCriterion):
         "CompletionCriterion is deprecated, please use TransitionCriterion instead."
     )
     pass
-
-
-class MinimumPreferenceOccurances(MinimumPreferenceOccurances):
-    """
-    Deprecated child class that has been replaced by `MinimumPreferenceOccurances`
-    in `TransitionCriterion`, and will be fully reaped in a future release.
-    """
-
-    logger.warning(
-        "CompletionCriterion, which MinimumPreferenceOccurance inherits from, is"
-        " deprecated. Please use TransitionCriterion instead."
-    )
-    pass
diff --git a/ax/modelbridge/tests/test_completion_criterion.py b/ax/modelbridge/tests/test_completion_criterion.py
index cf3ca2d78d3..e0d60b80f89 100644
--- a/ax/modelbridge/tests/test_completion_criterion.py
+++ b/ax/modelbridge/tests/test_completion_criterion.py
@@ -8,10 +8,9 @@
 import pandas as pd
 from ax.core.base_trial import TrialStatus
 from ax.core.data import Data
-from ax.modelbridge.completion_criterion import MinimumPreferenceOccurances
 from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
 from ax.modelbridge.registry import Models
-from ax.modelbridge.transition_criterion import MinTrials
+from ax.modelbridge.transition_criterion import MinimumPreferenceOccurances, MinTrials
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_experiment
 
diff --git a/ax/modelbridge/tests/test_transition_criterion.py b/ax/modelbridge/tests/test_transition_criterion.py
index 9fbde23aeca..d9f0faad62e 100644
--- a/ax/modelbridge/tests/test_transition_criterion.py
+++ b/ax/modelbridge/tests/test_transition_criterion.py
@@ -409,7 +409,8 @@ def test_repr(self) -> None:
         self.assertEqual(
             str(minimum_preference_occurances_criterion),
             "MinimumPreferenceOccurances({'metric_name': 'm1', 'threshold': 3, "
-            + "'transition_to': None, 'block_gen_if_met': False})",
+            + "'transition_to': None, 'block_gen_if_met': False, "
+            "'block_transition_if_unmet': True})",
         )
         deprecated_min_trials_criterion = MinimumTrialsInStatus(
             status=TrialStatus.COMPLETED, threshold=3
diff --git a/ax/modelbridge/transition_criterion.py b/ax/modelbridge/transition_criterion.py
index 022ca6c3406..f964a3cbcd3 100644
--- a/ax/modelbridge/transition_criterion.py
+++ b/ax/modelbridge/transition_criterion.py
@@ -340,12 +340,14 @@ def __init__(
         threshold: int,
         transition_to: Optional[str] = None,
         block_gen_if_met: Optional[bool] = False,
+        block_transition_if_unmet: Optional[bool] = True,
     ) -> None:
         self.metric_name = metric_name
         self.threshold = threshold
         super().__init__(
             transition_to=transition_to,
             block_gen_if_met=block_gen_if_met,
+            block_transition_if_unmet=block_transition_if_unmet,
         )
 
     def is_met(