Move TrialStatus out to its own file (preliminary cleanup) (faceboo…

…k#3325)

Summary: Pull Request resolved: facebook#3325

Reviewed By: saitcakmak

Differential Revision: D69223588

ax/analysis/healthcheck/tests/test_can_generate_candidates.py

@@ -13,7 +13,7 @@
     CanGenerateCandidatesAnalysis,
 )
 from ax.analysis.healthcheck.healthcheck_analysis import HealthcheckStatus
-from ax.core.base_trial import TrialStatus
+from ax.core.trial_status import TrialStatus
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_branin_experiment
 from pandas import testing as pdt

ax/benchmark/benchmark.py

@@ -35,11 +35,11 @@
 from ax.benchmark.benchmark_test_function import BenchmarkTestFunction
 from ax.benchmark.methods.sobol import get_sobol_generation_strategy
 from ax.core.arm import Arm
-from ax.core.base_trial import TrialStatus
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.search_space import SearchSpace
+from ax.core.trial_status import TrialStatus
 from ax.core.types import TParameterization, TParamValue
 from ax.core.utils import get_model_times
 from ax.service.scheduler import Scheduler

ax/benchmark/tests/test_benchmark_runner.py

@@ -25,11 +25,11 @@
     Jenatton,
 )
 from ax.core.arm import Arm
-from ax.core.base_trial import TrialStatus
 from ax.core.batch_trial import BatchTrial
 from ax.core.experiment import Experiment
 from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
+from ax.core.trial_status import TrialStatus
 from ax.exceptions.core import UnsupportedError
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.benchmark_stubs import (

ax/core/base_trial.py

@@ -12,7 +12,6 @@
 from collections.abc import Callable
 from copy import deepcopy
 from datetime import datetime, timedelta
-from enum import Enum
 from typing import Any, TYPE_CHECKING
 
 from ax.core.arm import Arm
@@ -23,6 +22,7 @@
 from ax.core.map_metric import MapMetric
 from ax.core.metric import Metric, MetricFetchResult
 from ax.core.runner import Runner
+from ax.core.trial_status import TrialStatus
 from ax.core.types import TCandidateMetadata, TEvaluationOutcome
 from ax.exceptions.core import UnsupportedError
 from ax.utils.common.base import SortableBase
@@ -34,144 +34,6 @@
     from ax import core  # noqa F401
 
 
-class TrialStatus(int, Enum):
-    """Enum of trial status.
-
-    General lifecycle of a trial is:::
-
-        CANDIDATE --> STAGED --> RUNNING --> COMPLETED
-                  ------------->         --> FAILED (retryable)
-                                         --> EARLY_STOPPED (deemed unpromising)
-                  -------------------------> ABANDONED (non-retryable)
-
-    Trial is marked as a ``CANDIDATE`` immediately upon its creation.
-
-    Trials may be abandoned at any time prior to completion or failure.
-    The difference between abandonment and failure is that the ``FAILED`` state
-    is meant to express a possibly transient or retryable error, so trials in
-    that state may be re-run and arm(s) in them may be resuggested by Ax models
-    to be added to new trials.
-
-    ``ABANDONED`` trials on the other end, indicate
-    that the trial (and arms(s) in it) should not be rerun or added to new
-    trials. A trial might be marked ``ABANDONED`` as a result of human-initiated
-    action (if some trial in experiment is poorly-performing, deterministically
-    failing etc., and should not be run again in the experiment). It might also
-    be marked ``ABANDONED`` in an automated way if the trial's execution
-    encounters an error that indicates that the arm(s) in the trial should bot
-    be evaluated in the experiment again (e.g. the parameterization in a given
-    arm deterministically causes trial evaluation to fail). Note that it's also
-    possible to abandon a single arm in a `BatchTrial` via
-    ``batch.mark_arm_abandoned``.
-
-    Early-stopped refers to trials that were deemed
-    unpromising by an early-stopping strategy and therefore terminated.
-
-    Additionally, when trials are deployed, they may be in an intermediate
-    staged state (e.g. scheduled but waiting for resources) or immediately
-    transition to running. Note that ``STAGED`` trial status is not always
-    applicable and depends on the ``Runner`` trials are deployed with
-    (and whether a ``Runner`` is present at all; for example, in Ax Service
-    API, trials are marked as ``RUNNING`` immediately when generated from
-    ``get_next_trial``, skipping the ``STAGED`` status).
-
-    NOTE: Data for abandoned trials (or abandoned arms in batch trials) is
-    not passed to the model as part of training data, unless ``fit_abandoned``
-    option is specified to model bridge. Additionally, data from MapMetrics is
-    typically excluded unless the corresponding trial is completed.
-    """
-
-    CANDIDATE = 0
-    STAGED = 1
-    FAILED = 2
-    COMPLETED = 3
-    RUNNING = 4
-    ABANDONED = 5
-    DISPATCHED = 6  # Deprecated.
-    EARLY_STOPPED = 7
-
-    @property
-    def is_terminal(self) -> bool:
-        """True if trial is completed."""
-        return (
-            self == TrialStatus.ABANDONED
-            or self == TrialStatus.COMPLETED
-            or self == TrialStatus.FAILED
-            or self == TrialStatus.EARLY_STOPPED
-        )
-
-    @property
-    def expecting_data(self) -> bool:
-        """True if trial is expecting data."""
-        return self in STATUSES_EXPECTING_DATA
-
-    @property
-    def is_deployed(self) -> bool:
-        """True if trial has been deployed but not completed."""
-        return self == TrialStatus.STAGED or self == TrialStatus.RUNNING
-
-    @property
-    def is_failed(self) -> bool:
-        """True if this trial is a failed one."""
-        return self == TrialStatus.FAILED
-
-    @property
-    def is_abandoned(self) -> bool:
-        """True if this trial is an abandoned one."""
-        return self == TrialStatus.ABANDONED
-
-    @property
-    def is_candidate(self) -> bool:
-        """True if this trial is a candidate."""
-        return self == TrialStatus.CANDIDATE
-
-    @property
-    def is_completed(self) -> bool:
-        """True if this trial is a successfully completed one."""
-        return self == TrialStatus.COMPLETED
-
-    @property
-    def is_running(self) -> bool:
-        """True if this trial is a running one."""
-        return self == TrialStatus.RUNNING
-
-    @property
-    def is_early_stopped(self) -> bool:
-        """True if this trial is an early stopped one."""
-        return self == TrialStatus.EARLY_STOPPED
-
-    def __format__(self, fmt: str) -> str:
-        """Define `__format__` to avoid pulling the `__format__` from the `int`
-        mixin (since its better for statuses to show up as `RUNNING` than as
-        just an int that is difficult to interpret).
-
-        E.g. batch trial representation with the overridden method is:
-        "BatchTrial(experiment_name='test', index=0, status=TrialStatus.CANDIDATE)".
-
-        Docs on enum formatting: https://docs.python.org/3/library/enum.html#others.
-        """
-        return f"{self!s}"
-
-    def __repr__(self) -> str:
-        return f"{self.__class__}.{self.name}"
-
-
-DEFAULT_STATUSES_TO_WARM_START: list[TrialStatus] = [
-    TrialStatus.RUNNING,
-    TrialStatus.COMPLETED,
-    TrialStatus.ABANDONED,
-    TrialStatus.EARLY_STOPPED,
-]
-
-NON_ABANDONED_STATUSES: set[TrialStatus] = set(TrialStatus) - {TrialStatus.ABANDONED}
-
-STATUSES_EXPECTING_DATA: list[TrialStatus] = [
-    TrialStatus.RUNNING,
-    TrialStatus.COMPLETED,
-    TrialStatus.EARLY_STOPPED,
-]
-
-
 def immutable_once_run(func: Callable) -> Callable:
     """Decorator for methods that should throw Error when
     trial is running or has ever run and immutable.

ax/core/experiment.py

@@ -22,12 +22,7 @@
 import pandas as pd
 from ax.core.arm import Arm
 from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
-from ax.core.base_trial import (
-    BaseTrial,
-    DEFAULT_STATUSES_TO_WARM_START,
-    STATUSES_EXPECTING_DATA,
-    TrialStatus,
-)
+from ax.core.base_trial import BaseTrial
 from ax.core.batch_trial import BatchTrial, LifecycleStage
 from ax.core.data import Data
 from ax.core.formatting_utils import DATA_TYPE_LOOKUP, DataType
@@ -41,6 +36,11 @@
 from ax.core.runner import Runner
 from ax.core.search_space import HierarchicalSearchSpace, SearchSpace
 from ax.core.trial import Trial
+from ax.core.trial_status import (
+    DEFAULT_STATUSES_TO_WARM_START,
+    STATUSES_EXPECTING_DATA,
+    TrialStatus,
+)
 from ax.core.types import ComparisonOp, TParameterization
 from ax.exceptions.core import (
     AxError,

ax/core/observation.py

@@ -19,11 +19,11 @@
 import numpy.typing as npt
 import pandas as pd
 from ax.core.arm import Arm
-from ax.core.base_trial import NON_ABANDONED_STATUSES, TrialStatus
 from ax.core.batch_trial import BatchTrial
 from ax.core.data import Data
 from ax.core.map_data import MapData
 from ax.core.map_metric import MapMetric
+from ax.core.trial_status import NON_ABANDONED_STATUSES, TrialStatus
 from ax.core.types import TCandidateMetadata, TParameterization
 from ax.utils.common.base import Base
 from ax.utils.common.constants import Keys

ax/core/tests/test_batch_trial.py

@@ -13,12 +13,12 @@
 
 import numpy as np
 from ax.core.arm import Arm
-from ax.core.base_trial import TrialStatus
 from ax.core.batch_trial import BatchTrial, GeneratorRunStruct
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun, GeneratorRunType
 from ax.core.parameter import FixedParameter, ParameterType
 from ax.core.search_space import SearchSpace
+from ax.core.trial_status import TrialStatus
 from ax.exceptions.core import UnsupportedError
 from ax.runners.synthetic import SyntheticRunner
 from ax.utils.common.testutils import TestCase

ax/core/tests/test_observation.py

@@ -12,7 +12,6 @@
 import numpy as np
 import pandas as pd
 from ax.core.arm import Arm
-from ax.core.base_trial import TrialStatus
 from ax.core.batch_trial import BatchTrial
 from ax.core.data import Data
 from ax.core.generator_run import GeneratorRun
@@ -29,6 +28,7 @@
     separate_observations,
 )
 from ax.core.trial import Trial
+from ax.core.trial_status import TrialStatus
 from ax.core.types import TParameterization
 from ax.utils.common.testutils import TestCase
 from pyre_extensions import none_throws

ax/core/tests/test_utils.py

@@ -12,14 +12,14 @@
 import numpy as np
 import pandas as pd
 from ax.core.arm import Arm
-from ax.core.base_trial import TrialStatus
 from ax.core.data import Data
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
 from ax.core.objective import Objective
 from ax.core.observation import ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.outcome_constraint import OutcomeConstraint
+from ax.core.trial_status import TrialStatus
 from ax.core.types import ComparisonOp
 from ax.core.utils import (
     best_feasible_objective,