feature: support the GridSearch strategy for hyperparameter optimization (#3439)

Co-authored-by: Anton Repushko <[email protected]>

Author: 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,7 @@
 HYPERPARAMETER_TUNING_JOB_NAME = "HyperParameterTuningJobName"
 PARENT_HYPERPARAMETER_TUNING_JOBS = "ParentHyperParameterTuningJobs"
 WARM_START_TYPE = "WarmStartType"
+GRID_SEARCH = "GridSearch"
 
 logger = logging.getLogger(__name__)
 
@@ -165,7 +174,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):
@@ -219,7 +229,7 @@ def __init__(
         metric_definitions: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
         strategy: Union[str, PipelineVariable] = "Bayesian",
         objective_type: Union[str, PipelineVariable] = "Maximize",
-        max_jobs: Union[int, PipelineVariable] = 1,
+        max_jobs: Union[int, PipelineVariable] = None,
         max_parallel_jobs: Union[int, PipelineVariable] = 1,
         tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
         base_tuning_job_name: Optional[str] = None,
@@ -258,7 +268,8 @@ def __init__(
                 evaluating training jobs. This value can be either 'Minimize' or
                 'Maximize' (default: 'Maximize').
             max_jobs (int or PipelineVariable): Maximum total number of training jobs to start for
-                the hyperparameter tuning job (default: 1).
+                the hyperparameter tuning job. The default value is unspecified fot the GridSearch
+                strategy and the default value is 1 for all others strategies (default: None).
             max_parallel_jobs (int or PipelineVariable): Maximum number of parallel training jobs to
                 start (default: 1).
             tags (list[dict[str, str] or list[dict[str, PipelineVariable]]): List of tags for
@@ -311,7 +322,12 @@ def __init__(
 
         self.strategy = strategy
         self.objective_type = objective_type
+        # For the GridSearch strategy we expect the max_jobs equals None and recalculate it later.
+        # For all other strategies for the backward compatibility we keep
+        # the default value as 1 (previous default value).
         self.max_jobs = max_jobs
+        if max_jobs is None and strategy is not GRID_SEARCH:
+            self.max_jobs = 1
         self.max_parallel_jobs = max_parallel_jobs
 
         self.tags = tags
@@ -373,7 +389,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 +451,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 +549,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 +573,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 +992,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 +1185,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
 
@@ -1301,7 +1338,7 @@ def create(
         base_tuning_job_name=None,
         strategy="Bayesian",
         objective_type="Maximize",
-        max_jobs=1,
+        max_jobs=None,
         max_parallel_jobs=1,
         tags=None,
         warm_start_config=None,
@@ -1351,7 +1388,8 @@ def create(
             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
-                tuning job (default: 1).
+                tuning job. The default value is unspecified fot the GridSearch strategy
+                and the value is 1 for all others strategies (default: None).
             max_parallel_jobs (int): Maximum number of parallel training jobs to start
                 (default: 1).
             tags (list[dict]): List of tags for labeling the tuning job (default: None). For more,

tests/unit/test_tuner.py

@@ -1774,3 +1774,22 @@ def test_no_tags_prefixes_non_jumpstart_models(
     assert sagemaker_session.create_model.call_args_list[0][1]["tags"] == []
 
     assert sagemaker_session.endpoint_from_production_variants.call_args_list[0][1]["tags"] == []
+
+
+def test_create_tuner_with_grid_search_strategy():
+    tuner = HyperparameterTuner.create(
+        base_tuning_job_name=BASE_JOB_NAME,
+        estimator_dict={ESTIMATOR_NAME: ESTIMATOR},
+        objective_metric_name_dict={ESTIMATOR_NAME: OBJECTIVE_METRIC_NAME},
+        hyperparameter_ranges_dict={ESTIMATOR_NAME: HYPERPARAMETER_RANGES},
+        metric_definitions_dict={ESTIMATOR_NAME: METRIC_DEFINITIONS},
+        strategy="GridSearch",
+        objective_type="Minimize",
+        max_parallel_jobs=1,
+        tags=TAGS,
+        warm_start_config=WARM_START_CONFIG,
+        early_stopping_type="Auto",
+    )
+
+    assert tuner is not None
+    assert tuner.max_jobs is None