Relax type hints on compute() (#2662)

Summary:
Pull Request resolved: #2662

Scheduler does not have a GenerationStrategy, it has a GenerationStrategyInterface. Relaxing this type hint is important for creating Scheduler.compute_analyses(...)

Reviewed By: Cesar-Cardoso

Differential Revision: D61339432

ax/analysis/analysis.py

@@ -11,7 +11,7 @@
 
 import pandas as pd
 from ax.core.experiment import Experiment
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.utils.common.base import Base
 from ax.utils.common.logger import get_logger
 from ax.utils.common.result import Err, ExceptionE, Ok, Result
@@ -87,7 +87,7 @@ class Analysis(Protocol):
     def compute(
         self,
         experiment: Optional[Experiment] = None,
-        generation_strategy: Optional[GenerationStrategy] = None,
+        generation_strategy: Optional[GenerationStrategyInterface] = None,
     ) -> AnalysisCard:
         # Note: when implementing compute always prefer experiment.lookup_data() to
         # experiment.fetch_data() to avoid unintential data fetching within the report
@@ -97,7 +97,7 @@ def compute(
     def compute_result(
         self,
         experiment: Optional[Experiment] = None,
-        generation_strategy: Optional[GenerationStrategy] = None,
+        generation_strategy: Optional[GenerationStrategyInterface] = None,
     ) -> Result[AnalysisCard, ExceptionE]:
         """
         Utility method to compute an AnalysisCard as a Result. This can be useful for

ax/analysis/markdown/markdown_analysis.py

@@ -9,7 +9,7 @@
 
 from ax.analysis.analysis import Analysis, AnalysisCard
 from ax.core.experiment import Experiment
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.core.generation_strategy_interface import GenerationStrategyInterface
 
 
 class MarkdownAnalysisCard(AnalysisCard):
@@ -27,5 +27,5 @@ class MarkdownAnalysis(Analysis):
     def compute(
         self,
         experiment: Optional[Experiment] = None,
-        generation_strategy: Optional[GenerationStrategy] = None,
+        generation_strategy: Optional[GenerationStrategyInterface] = None,
     ) -> MarkdownAnalysisCard: ...

ax/analysis/plotly/parallel_coordinates/parallel_coordinates.py

@@ -13,9 +13,9 @@
 
 from ax.analysis.plotly.plotly_analysis import PlotlyAnalysis, PlotlyAnalysisCard
 from ax.core.experiment import Experiment
+from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from ax.core.objective import MultiObjective, ScalarizedObjective
 from ax.exceptions.core import UnsupportedError, UserInputError
-from ax.modelbridge.generation_strategy import GenerationStrategy
 from plotly import graph_objects as go, io as pio
 
 
@@ -45,7 +45,7 @@ def __init__(self, metric_name: Optional[str] = None) -> None:
     def compute(
         self,
         experiment: Optional[Experiment] = None,
-        generation_strategy: Optional[GenerationStrategy] = None,
+        generation_strategy: Optional[GenerationStrategyInterface] = None,
     ) -> PlotlyAnalysisCard:
         if experiment is None:
             raise UserInputError("ParallelCoordinatesPlot requires an Experiment")

ax/analysis/plotly/plotly_analysis.py

@@ -9,7 +9,7 @@
 
 from ax.analysis.analysis import Analysis, AnalysisCard
 from ax.core.experiment import Experiment
-from ax.modelbridge.generation_strategy import GenerationStrategy
+from ax.core.generation_strategy_interface import GenerationStrategyInterface
 from plotly import graph_objects as go, io as pio
 
 
@@ -28,5 +28,5 @@ class PlotlyAnalysis(Analysis):
     def compute(
         self,
         experiment: Optional[Experiment] = None,
-        generation_strategy: Optional[GenerationStrategy] = None,
+        generation_strategy: Optional[GenerationStrategyInterface] = None,
     ) -> PlotlyAnalysisCard: ...