[tune] Preparation for deadline schedulers (ray-project#22006)

minaek · Feb 22, 2022 · 4a15c6f · 4a15c6f
1 parent dfe4706
commit 4a15c6f
Show file tree

Hide file tree

Showing 19 changed files with 1,315 additions and 123 deletions.
diff --git a/doc/source/tune/api_docs/schedulers.rst b/doc/source/tune/api_docs/schedulers.rst
@@ -245,15 +245,15 @@ An example of this in use can be found here: :doc:`/tune/examples/xgboost_dynami
 
 .. autoclass:: ray.tune.schedulers.ResourceChangingScheduler
 
-evenly_distribute_cpus_gpus
-~~~~~~~~~~~~~~~~~~~~~~~~~~~
+DistributeResources
+~~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: ray.tune.schedulers.resource_changing_scheduler.evenly_distribute_cpus_gpus
+.. autoclass:: ray.tune.schedulers.resource_changing_scheduler.DistributeResources
 
-evenly_distribute_cpus_gpus_distributed
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+DistributeResourcesToTopJob
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: ray.tune.schedulers.resource_changing_scheduler.evenly_distribute_cpus_gpus_distributed
+.. autoclass:: ray.tune.schedulers.resource_changing_scheduler.DistributeResourcesToTopJob
 
 FIFOScheduler
 -------------

diff --git a/python/ray/tune/BUILD b/python/ray/tune/BUILD
@@ -334,6 +334,14 @@ py_test(
     tags = ["team:ml", "exclusive", "tests_dir_T"],
 )
 
+py_test(
+    name = "test_trial_scheduler_resource_changing",
+    size = "small",
+    srcs = ["tests/test_trial_scheduler_resource_changing.py"],
+    deps = [":tune_lib"],
+    tags = ["team:ml", "exclusive", "tests_dir_T"],
+)
+
 py_test(
     name = "test_tune_restore_warm_start",
     size = "large",

diff --git a/python/ray/tune/examples/xgboost_dynamic_resources_example.py b/python/ray/tune/examples/xgboost_dynamic_resources_example.py
@@ -181,7 +181,7 @@ def example_resources_allocation_function(
         trial has currently, the scheduler will skip the update process
         internally (same with None).
 
-        See :func:`evenly_distribute_cpus_gpus` for a more complex,
+        See :class:`DistributeResources` for a more complex,
         robust approach.
 
         Args:
@@ -221,15 +221,15 @@ def example_resources_allocation_function(
         return PlacementGroupFactory([{"CPU": cpu_to_use, "GPU": 0}])
 
     # You can either define your own resources_allocation_function, or
-    # use the default one - evenly_distribute_cpus_gpus
+    # use the default one - DistributeResources
 
     # from ray.tune.schedulers.resource_changing_scheduler import \
-    #    evenly_distribute_cpus_gpus
+    #    DistributeResources
 
     scheduler = ResourceChangingScheduler(
         base_scheduler=base_scheduler,
         resources_allocation_function=example_resources_allocation_function
-        # resources_allocation_function=evenly_distribute_cpus_gpus  # default
+        # resources_allocation_function=DistributeResources()  # default
     )
 
     if use_class_trainable:

diff --git a/python/ray/tune/experiment.py b/python/ray/tune/experiment.py
@@ -76,7 +76,7 @@ class Experiment:
     """
 
     # Keys that will be present in `public_spec` dict.
-    PUBLIC_KEYS = {"stop", "num_samples"}
+    PUBLIC_KEYS = {"stop", "num_samples", "time_budget_s"}
 
     def __init__(
         self,
@@ -188,6 +188,7 @@ def __init__(
         spec = {
             "run": self._run_identifier,
             "stop": stopping_criteria,
+            "time_budget_s": time_budget_s,
             "config": config,
             "resources_per_trial": resources_per_trial,
             "num_samples": num_samples,

diff --git a/python/ray/tune/result.py b/python/ray/tune/result.py
@@ -102,6 +102,7 @@
     "iterations_since_restore",
     "timesteps_since_restore",
     "config",
+    "warmup_time",
 )
 
 # __duplicate__ is a magic keyword used internally to

diff --git a/python/ray/tune/schedulers/async_hyperband.py b/python/ray/tune/schedulers/async_hyperband.py
@@ -81,7 +81,9 @@ def __init__(
             self._metric_op = -1.0
         self._time_attr = time_attr
 
-    def set_search_properties(self, metric: Optional[str], mode: Optional[str]) -> bool:
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
         if self._metric and metric:
             return False
         if self._mode and mode:

diff --git a/python/ray/tune/schedulers/hyperband.py b/python/ray/tune/schedulers/hyperband.py
@@ -120,7 +120,9 @@ def __init__(
         self._time_attr = time_attr
         self._stop_last_trials = stop_last_trials
 
-    def set_search_properties(self, metric: Optional[str], mode: Optional[str]) -> bool:
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
         if self._metric and metric:
             return False
         if self._mode and mode:

diff --git a/python/ray/tune/schedulers/median_stopping_rule.py b/python/ray/tune/schedulers/median_stopping_rule.py
@@ -73,7 +73,9 @@ def __init__(
         self._last_pause = collections.defaultdict(lambda: float("-inf"))
         self._results = collections.defaultdict(list)
 
-    def set_search_properties(self, metric: Optional[str], mode: Optional[str]) -> bool:
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
         if self._metric and metric:
             return False
         if self._mode and mode:

diff --git a/python/ray/tune/schedulers/pbt.py b/python/ray/tune/schedulers/pbt.py
@@ -314,7 +314,9 @@ def __init__(
         self._num_checkpoints = 0
         self._num_perturbations = 0
 
-    def set_search_properties(self, metric: Optional[str], mode: Optional[str]) -> bool:
+    def set_search_properties(
+        self, metric: Optional[str], mode: Optional[str], **spec
+    ) -> bool:
         if self._metric and metric:
             return False
         if self._mode and mode: