chore: import upstream snapshot with attribution

python/ray/tune/schedulers/__init__.py

@@ -0,0 +1,96 @@
+import inspect
+from typing import Any
+
+from ray._common.utils import get_function_args
+from ray.tune.schedulers.async_hyperband import ASHAScheduler, AsyncHyperBandScheduler
+from ray.tune.schedulers.hb_bohb import HyperBandForBOHB
+from ray.tune.schedulers.hyperband import HyperBandScheduler
+from ray.tune.schedulers.median_stopping_rule import MedianStoppingRule
+from ray.tune.schedulers.pbt import (
+    PopulationBasedTraining,
+    PopulationBasedTrainingReplay,
+)
+from ray.tune.schedulers.resource_changing_scheduler import ResourceChangingScheduler
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.util import PublicAPI
+
+
+def _pb2_importer():
+    # PB2 is imported lazily since it has additional dependencies.
+    from ray.tune.schedulers.pb2 import PB2
+
+    return PB2
+
+
+# Values in this dictionary will be one two kinds:
+#    class of the scheduler object to create
+#    wrapper function to support a lazy import of the scheduler class
+SCHEDULER_IMPORT = {
+    "fifo": FIFOScheduler,
+    "async_hyperband": AsyncHyperBandScheduler,
+    "asynchyperband": AsyncHyperBandScheduler,
+    "median_stopping_rule": MedianStoppingRule,
+    "medianstopping": MedianStoppingRule,
+    "hyperband": HyperBandScheduler,
+    "hb_bohb": HyperBandForBOHB,
+    "pbt": PopulationBasedTraining,
+    "pbt_replay": PopulationBasedTrainingReplay,
+    "pb2": _pb2_importer,
+    "resource_changing": ResourceChangingScheduler,
+}
+
+
+@PublicAPI(stability="beta")
+def create_scheduler(
+    scheduler: str,
+    **kwargs: Any,
+) -> TrialScheduler:
+    """Instantiate a scheduler based on the given string.
+
+    This is useful for swapping between different schedulers.
+
+    Args:
+        scheduler: The scheduler to use.
+        **kwargs: Scheduler parameters.
+            These keyword arguments will be passed to the initialization
+            function of the chosen scheduler.
+    Returns:
+        ray.tune.schedulers.trial_scheduler.TrialScheduler: The scheduler.
+    Example:
+        >>> from ray import tune
+        >>> pbt_kwargs = {}
+        >>> scheduler = tune.create_scheduler('pbt', **pbt_kwargs) # doctest: +SKIP
+    """
+
+    scheduler = scheduler.lower()
+    if scheduler not in SCHEDULER_IMPORT:
+        raise ValueError(
+            f"The `scheduler` argument must be one of "
+            f"{list(SCHEDULER_IMPORT)}. "
+            f"Got: {scheduler}"
+        )
+
+    SchedulerClass = SCHEDULER_IMPORT[scheduler]
+
+    if inspect.isfunction(SchedulerClass):
+        # invoke the wrapper function to retrieve class
+        SchedulerClass = SchedulerClass()
+
+    scheduler_args = get_function_args(SchedulerClass)
+    trimmed_kwargs = {k: v for k, v in kwargs.items() if k in scheduler_args}
+
+    return SchedulerClass(**trimmed_kwargs)
+
+
+__all__ = [
+    "TrialScheduler",
+    "HyperBandScheduler",
+    "AsyncHyperBandScheduler",
+    "ASHAScheduler",
+    "MedianStoppingRule",
+    "FIFOScheduler",
+    "PopulationBasedTraining",
+    "PopulationBasedTrainingReplay",
+    "HyperBandForBOHB",
+    "ResourceChangingScheduler",
+]