feature: support the Hyperband strategy with the StrategyConfig

Author: Anton Repushko

src/sagemaker/tuner.py

@@ -33,7 +33,10 @@
 from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.inputs import TrainingInput, FileSystemInput
 from sagemaker.job import _Job
-from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
+from sagemaker.jumpstart.utils import (
+    add_jumpstart_tags,
+    get_jumpstart_base_name_if_jumpstart_model,
+)
 from sagemaker.parameter import (
     CategoricalParameter,
     ContinuousParameter,
@@ -44,7 +47,12 @@
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 
 from sagemaker.session import Session
-from sagemaker.utils import base_from_name, base_name_from_image, name_from_base, to_string
+from sagemaker.utils import (
+    base_from_name,
+    base_name_from_image,
+    name_from_base,
+    to_string,
+)
 
 AMAZON_ESTIMATOR_MODULE = "sagemaker"
 AMAZON_ESTIMATOR_CLS_NAMES = {
@@ -60,6 +68,9 @@
 HYPERPARAMETER_TUNING_JOB_NAME = "HyperParameterTuningJobName"
 PARENT_HYPERPARAMETER_TUNING_JOBS = "ParentHyperParameterTuningJobs"
 WARM_START_TYPE = "WarmStartType"
+HYPERBAND_STRATEGY_CONFIG = "HyperbandStrategyConfig"
+HYPERBAND_MIN_RESOURCE = "MinResource"
+HYPERBAND_MAX_RESOURCE = "MaxResource"
 
 logger = logging.getLogger(__name__)
 
@@ -165,7 +176,8 @@ def from_job_desc(cls, warm_start_config):
             parents.append(parent[HYPERPARAMETER_TUNING_JOB_NAME])
 
         return cls(
-            warm_start_type=WarmStartTypes(warm_start_config[WARM_START_TYPE]), parents=parents
+            warm_start_type=WarmStartTypes(warm_start_config[WARM_START_TYPE]),
+            parents=parents,
         )
 
     def to_input_req(self):
@@ -197,6 +209,179 @@ def to_input_req(self):
         }
 
 
+class HyperbandStrategyConfig(object):
+    """The configuration for Hyperband, a multi-fidelity based hyperparameter tuning strategy.
+
+    Hyperband uses the final and intermediate results of a training job to dynamically allocate
+    resources to hyperparameter configurations being evaluated while automatically stopping
+    under-performing configurations. This parameter should be provided only if Hyperband is
+    selected as the Strategy under the HyperParameterTuningJobConfig.
+
+    Examples:
+        >>> hyperband_strategy_config = HyperbandStrategyConfig(
+        >>>                                 max_resource=10, min_resource = 1)
+        >>> hyperband_strategy_config.max_resource
+        10
+        >>> hyperband_strategy_config.min_resource
+        1
+    """
+
+    def __init__(self, max_resource: int, min_resource: int):
+        """Creates a ``HyperbandStrategyConfig`` with provided `min_resource`` and ``max_resource``.
+
+        Args:
+            max_resource (int): The maximum number of resources (such as epochs) that can be used
+            by a training job launched by a hyperparameter tuning job.
+                Once a job reaches the MaxResource value, it is stopped.
+                If a value for MaxResource is not provided, and Hyperband is selected as the
+                hyperparameter tuning strategy, HyperbandTrainingJ attempts to infer MaxResource
+                from the following keys (if present) in StaticsHyperParameters:
+                    epochs
+                    numepochs
+                    n-epochs
+                    n_epochs
+                    num_epochs
+                If HyperbandStrategyConfig is unable to infer a value for MaxResource, it generates
+                a validation error.
+                The maximum value is 20,000 epochs. All metrics that correspond to an objective
+                metric are used to derive early stopping decisions.
+                For distributed training jobs, ensure that duplicate metrics are not printed in the
+                logs across the individual nodes in a training job.
+                If multiple nodes are publishing duplicate or incorrect metrics, hyperband
+                optimisation algorithm may make an incorrect stopping decision and stop the job
+                prematurely.
+            min_resource (int): The minimum number of resources (such as epochs)
+                that can be used by a training job launched by a hyperparameter tuning job.
+                If the value for MinResource has not been reached, the training job will not be
+                stopped by Hyperband.
+        """
+        self.min_resource = min_resource
+        self.max_resource = max_resource
+
+    @classmethod
+    def from_job_desc(cls, hyperband_strategy_config):
+        """Creates a ``HyperbandStrategyConfig`` from a hyperband strategy configuration response.
+
+        This is the Hyperband strategy configuration from the DescribeTuningJob response.
+
+        Examples:
+            >>> hyperband_strategy_config =
+            >>>     HyperbandStrategyConfig.from_job_desc(hyperband_strategy_config={
+            >>>         "MaxResource": 10,
+            >>>         "MinResource": 1
+            >>>     })
+            >>> hyperband_strategy_config.max_resource
+            10
+            >>> hyperband_strategy_config.min_resource
+            1
+
+        Args:
+            hyperband_strategy_config (dict): The expected format of the
+                ``hyperband_strategy_config`` contains two first-class fields
+
+        Returns:
+            sagemaker.tuner.HyperbandStrategyConfig: De-serialized instance of
+            HyperbandStrategyConfig containing the max_resource and min_resource provided as part of
+            ``hyperband_strategy_config``.
+        """
+        return cls(
+            min_resource=hyperband_strategy_config[HYPERBAND_MIN_RESOURCE],
+            max_resource=hyperband_strategy_config[HYPERBAND_MAX_RESOURCE],
+        )
+
+    def to_input_req(self):
+        """Converts the ``self`` instance to the desired input request format.
+
+        Examples:
+            >>> hyperband_strategy_config = HyperbandStrategyConfig (
+                max_resource=10,
+                min_resource=1
+            )
+            >>> hyperband_strategy_config.to_input_req()
+            {
+                "MaxResource":10,
+                "MinResource": 1
+            }
+
+        Returns:
+            dict: Containing the "MaxResource" and
+            "MinResource" as the first class fields.
+        """
+        return {
+            HYPERBAND_MIN_RESOURCE: self.min_resource,
+            HYPERBAND_MAX_RESOURCE: self.max_resource,
+        }
+
+
+class StrategyConfig(object):
+    """The configuration for a training job launched by a hyperparameter tuning job.
+
+    Choose Bayesian for Bayesian optimization, and Random for random search optimization.
+    For more advanced use cases, use Hyperband, which evaluates objective metrics for training jobs
+    after every epoch.
+    """
+
+    def __init__(
+        self,
+        hyperband_strategy_config: HyperbandStrategyConfig,
+    ):
+        """Creates a ``StrategyConfig`` with provided ``HyperbandStrategyConfig``.
+
+        Args:
+            hyperband_strategy_config (sagemaker.tuner.HyperbandStrategyConfig): The configuration
+            for the object that specifies the Hyperband strategy.
+                This parameter is only supported for the Hyperband selection for Strategy within
+                the HyperParameterTuningJobConfig.
+        """
+
+        self.hyperband_strategy_config = hyperband_strategy_config
+
+    @classmethod
+    def from_job_desc(cls, strategy_config):
+        """Creates a ``HyperbandStrategyConfig`` from a hyperband strategy configuration response.
+
+        This is the hyper band strategy configuration from the DescribeTuningJob response.
+
+        Args:
+            strategy_config (dict): The expected format of the
+                ``strategy_config`` contains one first-class field
+
+        Returns:
+            sagemaker.tuner.StrategyConfig: De-serialized instance of
+            StrategyConfig containing the strategy configuration.
+        """
+        return cls(
+            hyperband_strategy_config=HyperbandStrategyConfig.from_job_desc(
+                strategy_config[HYPERBAND_STRATEGY_CONFIG]
+            )
+        )
+
+    def to_input_req(self):
+        """Converts the ``self`` instance to the desired input request format.
+
+        Examples:
+            >>> strategy_config = StrategyConfig(
+                HyperbandStrategyConfig(
+                    max_resource=10,
+                    min_resource=1
+                )
+            )
+            >>> strategy_config.to_input_req()
+            {
+                "HyperbandStrategyConfig": {
+                    "MaxResource":10,
+                    "MinResource": 1
+                }
+            }
+
+        Returns:
+            dict: Containing the strategy configurations.
+        """
+        return {
+            HYPERBAND_STRATEGY_CONFIG: self.hyperband_strategy_config.to_input_req(),
+        }
+
+
 class HyperparameterTuner(object):
     """Defines interaction with Amazon SageMaker hyperparameter tuning jobs.
 
@@ -224,6 +409,7 @@ def __init__(
         tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
         base_tuning_job_name: Optional[str] = None,
         warm_start_config: Optional[WarmStartConfig] = None,
+        strategy_config: Optional[StrategyConfig] = None,
         early_stopping_type: Union[str, PipelineVariable] = "Off",
         estimator_name: Optional[str] = None,
     ):
@@ -272,6 +458,8 @@ def __init__(
             warm_start_config (sagemaker.tuner.WarmStartConfig): A
                 ``WarmStartConfig`` object that has been initialized with the
                 configuration defining the nature of warm start tuning job.
+            strategy_config (sagemaker.tuner.StrategyConfig): A configuration for "Hyperparameter"
+            tuning job optimisation strategy.
             early_stopping_type (str or PipelineVariable): Specifies whether early stopping is
                 enabled for the job. Can be either 'Auto' or 'Off' (default:
                 'Off'). If set to 'Off', early stopping will not be attempted.
@@ -310,6 +498,7 @@ def __init__(
         self._validate_parameter_ranges(estimator, hyperparameter_ranges)
 
         self.strategy = strategy
+        self.strategy_config = strategy_config
         self.objective_type = objective_type
         self.max_jobs = max_jobs
         self.max_parallel_jobs = max_parallel_jobs
@@ -373,7 +562,8 @@ def _prepare_job_name_for_tuning(self, job_name=None):
                     self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
                 )
                 base_name = base_name_from_image(
-                    estimator.training_image_uri(), default_base_name=EstimatorBase.JOB_CLASS_NAME
+                    estimator.training_image_uri(),
+                    default_base_name=EstimatorBase.JOB_CLASS_NAME,
                 )
 
                 jumpstart_base_name = get_jumpstart_base_name_if_jumpstart_model(
@@ -434,7 +624,15 @@ def _prepare_static_hyperparameters(
     def fit(
         self,
         inputs: Optional[
-            Union[str, Dict, List, TrainingInput, FileSystemInput, RecordSet, FileSystemRecordSet]
+            Union[
+                str,
+                Dict,
+                List,
+                TrainingInput,
+                FileSystemInput,
+                RecordSet,
+                FileSystemRecordSet,
+            ]
         ] = None,
         job_name: Optional[str] = None,
         include_cls_metadata: Union[bool, Dict[str, bool]] = False,
@@ -524,7 +722,9 @@ def _fit_with_estimator_dict(self, inputs, job_name, include_cls_metadata, estim
             allowed_keys=estimator_names,
         )
         self._validate_dict_argument(
-            name="estimator_kwargs", value=estimator_kwargs, allowed_keys=estimator_names
+            name="estimator_kwargs",
+            value=estimator_kwargs,
+            allowed_keys=estimator_names,
         )
 
         for (estimator_name, estimator) in self.estimator_dict.items():
@@ -546,7 +746,13 @@ def _prepare_estimator_for_tuning(cls, estimator, inputs, job_name, **kwargs):
             estimator._prepare_for_training(job_name)
 
     @classmethod
-    def attach(cls, tuning_job_name, sagemaker_session=None, job_details=None, estimator_cls=None):
+    def attach(
+        cls,
+        tuning_job_name,
+        sagemaker_session=None,
+        job_details=None,
+        estimator_cls=None,
+    ):
         """Attach to an existing hyperparameter tuning job.
 
         Create a HyperparameterTuner bound to an existing hyperparameter
@@ -959,7 +1165,8 @@ def _prepare_estimator_cls(cls, estimator_cls, training_details):
 
         # Default to the BYO estimator
         return getattr(
-            importlib.import_module(cls.DEFAULT_ESTIMATOR_MODULE), cls.DEFAULT_ESTIMATOR_CLS_NAME
+            importlib.import_module(cls.DEFAULT_ESTIMATOR_MODULE),
+            cls.DEFAULT_ESTIMATOR_CLS_NAME,
         )
 
     @classmethod
@@ -1151,7 +1358,10 @@ def _validate_parameter_ranges(self, estimator, hyperparameter_ranges):
 
     def _validate_parameter_range(self, value_hp, parameter_range):
         """Placeholder docstring"""
-        for (parameter_range_key, parameter_range_value) in parameter_range.__dict__.items():
+        for (
+            parameter_range_key,
+            parameter_range_value,
+        ) in parameter_range.__dict__.items():
             if parameter_range_key == "scaling_type":
                 continue
 
@@ -1258,6 +1468,7 @@ def _create_warm_start_tuner(self, additional_parents, warm_start_type, estimato
                 objective_metric_name=self.objective_metric_name,
                 hyperparameter_ranges=self._hyperparameter_ranges,
                 strategy=self.strategy,
+                strategy_config=self.strategy_config,
                 objective_type=self.objective_type,
                 max_jobs=self.max_jobs,
                 max_parallel_jobs=self.max_parallel_jobs,
@@ -1284,6 +1495,7 @@ def _create_warm_start_tuner(self, additional_parents, warm_start_type, estimato
             hyperparameter_ranges_dict=self._hyperparameter_ranges_dict,
             metric_definitions_dict=self.metric_definitions_dict,
             strategy=self.strategy,
+            strategy_config=self.strategy_config,
             objective_type=self.objective_type,
             max_jobs=self.max_jobs,
             max_parallel_jobs=self.max_parallel_jobs,
@@ -1300,6 +1512,7 @@ def create(
         metric_definitions_dict=None,
         base_tuning_job_name=None,
         strategy="Bayesian",
+        strategy_config=None,
         objective_type="Maximize",
         max_jobs=1,
         max_parallel_jobs=1,
@@ -1343,11 +1556,13 @@ def create(
                 metric from the logs. This should be defined only for hyperparameter tuning jobs
                 that don't use an Amazon algorithm.
             base_tuning_job_name (str): Prefix for the hyperparameter tuning job name when the
-                :meth:`~sagemaker.tuner.HyperparameterTuner.fit` method launches. If not specified,
-                a default job name is generated, based on the training image name and current
-                timestamp.
+                :meth:`~sagemaker.tuner.HyperparameterTuner.fit` method launches.
+                If not specified, a default job name is generated,
+                based on the training image name and current timestamp.
             strategy (str): Strategy to be used for hyperparameter estimations
                 (default: 'Bayesian').
+            strategy_config (dict): The configuration for a training job launched by a
+            hyperparameter tuning job.
             objective_type (str): The type of the objective metric for evaluating training jobs.
                 This value can be either 'Minimize' or 'Maximize' (default: 'Maximize').
             max_jobs (int): Maximum total number of training jobs to start for the hyperparameter
@@ -1394,6 +1609,7 @@ def create(
             hyperparameter_ranges=hyperparameter_ranges_dict[first_estimator_name],
             metric_definitions=metric_definitions,
             strategy=strategy,
+            strategy_config=strategy_config,
             objective_type=objective_type,
             max_jobs=max_jobs,
             max_parallel_jobs=max_parallel_jobs,
@@ -1551,6 +1767,9 @@ def _get_tuner_args(cls, tuner, inputs):
             "early_stopping_type": tuner.early_stopping_type,
         }
 
+        if tuner.strategy_config is not None:
+            tuning_config["strategy_config"] = tuner.strategy_config
+
         if tuner.objective_metric_name is not None:
             tuning_config["objective_type"] = tuner.objective_type
             tuning_config["objective_metric_name"] = tuner.objective_metric_name