Move poll_trial_status and poll_available_capacity to TorchXRunner, remove TorchXScheduler

Summary:
Adjusting the TorchX setup following D31031589 and D31032567
Deprecate `torchx.runtime.hpo.ax.TorchXScheduler`

Differential Revision: D33062113

fbshipit-source-id: 41987a64ba9671e93a35f55088e7a5d77eafca78

Author: aivanou

docs/source/runtime/hpo.rst

@@ -14,5 +14,4 @@ Ax (Adaptive Experimentation)
 .. currentmodule:: torchx.runtime.hpo.ax
 
 .. autoclass:: TorchXRunner
-.. autoclass:: TorchXScheduler
-.. autoclass:: AppMetric
+.. autoclass:: AppMetric

torchx/runtime/hpo/__init__.py

@@ -54,13 +54,13 @@
     SearchSpace,
  )
  from ax.modelbridge.dispatch_utils import choose_generation_strategy
- from ax.service.scheduler import SchedulerOptions
+ from ax.service.scheduler import SchedulerOptions, Scheduler
  from ax.service.utils.best_point import get_best_parameters
  from ax.service.utils.report_utils import exp_to_df
  from ax.utils.common.constants import Keys
  from pyre_extensions import none_throws
  from torchx.components import utils
- from torchx.runtime.hpo.ax import AppMetric, TorchXRunner, TorchXScheduler
+ from torchx.runtime.hpo.ax import AppMetric, TorchXRunner
 
  # Run HPO on the booth function (https://en.wikipedia.org/wiki/Test_functions_for_optimization)
 
@@ -100,7 +100,7 @@
      properties={Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True},
  )
 
- scheduler = TorchXScheduler(
+ scheduler = Scheduler(
      experiment=experiment,
      generation_strategy=(
          choose_generation_strategy(

torchx/runtime/hpo/ax.py

@@ -6,16 +6,15 @@
 # LICENSE file in the root directory of this source tree.
 
 import inspect
-from typing import Any, Callable, Dict, Mapping, Optional, Set, cast
+from typing import Any, Callable, Dict, Mapping, Optional, Set, cast, Iterable
 
 import pandas as pd
 from ax.core import Trial
 from ax.core.base_trial import BaseTrial
 from ax.core.data import Data
 from ax.core.metric import Metric
 from ax.core.runner import Runner as ax_Runner
-from ax.service.scheduler import Scheduler as ax_Scheduler, TrialStatus
-from ax.utils.common.typeutils import not_none
+from ax.service.scheduler import TrialStatus
 from pyre_extensions import none_throws
 from torchx.runner import Runner, get_runner
 from torchx.runtime.tracking import FsspecResultTracker
@@ -206,25 +205,12 @@ def run(self, trial: BaseTrial) -> Dict[str, Any]:
             _TORCHX_TRACKER_BASE: self._tracker_base,
         }
 
-
-class TorchXScheduler(ax_Scheduler):
-    """
-    An implementation of an `Ax Scheduler <https://ax.dev/tutorials/scheduler.html>`_
-    that works with Experiments hooked up with the ``TorchXRunner``.
-
-    This scheduler is not a real scheduler but rather a facade scheduler
-    that delegates to scheduler clients for various remote/local schedulers.
-    For a list of supported schedulers please refer to TorchX
-    `scheduler docs <https://pytorch.org/torchx/latest/schedulers.html>`_.
-
-    """
-
     def poll_trial_status(
-        self, poll_all_trial_statuses: bool = False
+        self, trials: Iterable[BaseTrial]
     ) -> Dict[TrialStatus, Set[int]]:
         trial_statuses: Dict[TrialStatus, Set[int]] = {}
 
-        for trial in self.running_trials:
+        for trial in trials:
             app_handle: str = trial.run_metadata[_TORCHX_APP_HANDLE]
             torchx_runner = trial.run_metadata[_TORCHX_RUNNER]
             app_status: AppStatus = torchx_runner.status(app_handle)
@@ -250,8 +236,4 @@ def poll_available_capacity(self) -> int:
 
         """
 
-        return (
-            -1
-            if self.generation_strategy._curr.max_parallelism is None
-            else not_none(self.generation_strategy._curr.max_parallelism)
-        )
+        return super().poll_available_capacity()

torchx/runtime/hpo/test/ax_test.py

@@ -22,16 +22,16 @@
     SearchSpace,
 )
 from ax.modelbridge.dispatch_utils import choose_generation_strategy
-from ax.service.scheduler import SchedulerOptions
+from ax.service.scheduler import SchedulerOptions, Scheduler
 from ax.service.utils.best_point import get_best_parameters
 from ax.service.utils.report_utils import exp_to_df
 from ax.utils.common.constants import Keys
 from pyre_extensions import none_throws
 from torchx.components import utils
-from torchx.runtime.hpo.ax import AppMetric, TorchXRunner, TorchXScheduler
+from torchx.runtime.hpo.ax import AppMetric, TorchXRunner
 
 
-class TorchXSchedulerTest(unittest.TestCase):
+class TorchXAxTest(unittest.TestCase):
     def setUp(self) -> None:
         self.test_dir = tempfile.mkdtemp("torchx_runtime_hpo_ax_test")
 
@@ -86,7 +86,7 @@ def test_run_experiment_locally(self) -> None:
 
         # maybe add-on cfg into SchedulerOption?
         # so that we can pass it from one place
-        scheduler = TorchXScheduler(
+        scheduler = Scheduler(
             experiment=experiment,
             generation_strategy=(
                 choose_generation_strategy(
@@ -131,7 +131,7 @@ def test_run_experiment_locally_in_batches(self) -> None:
 
         # maybe add-on cfg into SchedulerOption?
         # so that we can pass it from one place
-        scheduler = TorchXScheduler(
+        scheduler = Scheduler(
             experiment=experiment,
             generation_strategy=(
                 choose_generation_strategy(