feature: add support of the intelligent stopping in the tuner

Author: Anton Repushko

src/sagemaker/session.py

@@ -2178,6 +2178,7 @@ def tune(  # noqa: C901
         objective_metric_name,
         max_jobs,
         max_parallel_jobs,
+        max_runtime_in_seconds,
         parameter_ranges,
         static_hyperparameters,
         input_mode,
@@ -2190,6 +2191,7 @@ def tune(  # noqa: C901
         tags,
         warm_start_config,
         strategy_config=None,
+        completion_criteria_config=None,
         enable_network_isolation=False,
         image_uri=None,
         algorithm_arn=None,
@@ -2215,6 +2217,8 @@ def tune(  # noqa: C901
             max_jobs (int): Maximum total number of training jobs to start for the hyperparameter
                 tuning job.
             max_parallel_jobs (int): Maximum number of parallel training jobs to start.
+            max_runtime_in_seconds (int or PipelineVariable): The maximum time in seconds
+                that a training job launched by a hyperparameter tuning job can run.
             parameter_ranges (dict): Dictionary of parameter ranges. These parameter ranges can be
                 one of three types: Continuous, Integer, or Categorical.
             static_hyperparameters (dict): Hyperparameters for model training. These
@@ -2255,6 +2259,8 @@ def tune(  # noqa: C901
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             warm_start_config (dict): Configuration defining the type of warm start and
                 other required configurations.
+            completion_criteria_config (sagemaker.tuner.TuningJobCompletionCriteriaConfig): A
+                configuration for the completion criteria.
             early_stopping_type (str): Specifies whether early stopping is enabled for the job.
                 Can be either 'Auto' or 'Off'. If set to 'Off', early stopping will not be
                 attempted. If set to 'Auto', early stopping of some training jobs may happen, but
@@ -2294,12 +2300,14 @@ def tune(  # noqa: C901
                 strategy=strategy,
                 max_jobs=max_jobs,
                 max_parallel_jobs=max_parallel_jobs,
+                max_runtime_in_seconds=max_runtime_in_seconds,
                 objective_type=objective_type,
                 objective_metric_name=objective_metric_name,
                 parameter_ranges=parameter_ranges,
                 early_stopping_type=early_stopping_type,
                 random_seed=random_seed,
                 strategy_config=strategy_config,
+                completion_criteria_config=completion_criteria_config,
             ),
             "TrainingJobDefinition": self._map_training_config(
                 static_hyperparameters=static_hyperparameters,
@@ -2452,12 +2460,14 @@ def _map_tuning_config(
         strategy,
         max_jobs,
         max_parallel_jobs,
+        max_runtime_in_seconds=None,
         early_stopping_type="Off",
         objective_type=None,
         objective_metric_name=None,
         parameter_ranges=None,
         random_seed=None,
         strategy_config=None,
+        completion_criteria_config=None,
     ):
         """Construct tuning job configuration dictionary.
 
@@ -2466,6 +2476,8 @@ def _map_tuning_config(
             max_jobs (int): Maximum total number of training jobs to start for the hyperparameter
                 tuning job.
             max_parallel_jobs (int): Maximum number of parallel training jobs to start.
+            max_runtime_in_seconds (int or PipelineVariable): The maximum time in seconds
+                that a training job launched by a hyperparameter tuning job can run.
             early_stopping_type (str): Specifies whether early stopping is enabled for the job.
                 Can be either 'Auto' or 'Off'. If set to 'Off', early stopping will not be
                 attempted. If set to 'Auto', early stopping of some training jobs may happen,
@@ -2480,6 +2492,8 @@ def _map_tuning_config(
                 produce more consistent configurations for the same tuning job.
             strategy_config (dict): A configuration for the hyperparameter tuning job optimisation
                 strategy.
+            completion_criteria_config (dict): A configuration
+                for the completion criteria.
 
         Returns:
             A dictionary of tuning job configuration. For format details, please refer to
@@ -2496,6 +2510,9 @@ def _map_tuning_config(
             "TrainingJobEarlyStoppingType": early_stopping_type,
         }
 
+        if max_runtime_in_seconds is not None:
+            tuning_config["ResourceLimits"]["MaxRuntimeInSeconds"] = max_runtime_in_seconds
+
         if random_seed is not None:
             tuning_config["RandomSeed"] = random_seed
 
@@ -2508,6 +2525,9 @@ def _map_tuning_config(
 
         if strategy_config is not None:
             tuning_config["StrategyConfig"] = strategy_config
+
+        if completion_criteria_config is not None:
+            tuning_config["TuningJobCompletionCriteria"] = completion_criteria_config
         return tuning_config
 
     @classmethod

src/sagemaker/tuner.py

@@ -72,6 +72,12 @@
 HYPERBAND_MIN_RESOURCE = "MinResource"
 HYPERBAND_MAX_RESOURCE = "MaxResource"
 GRID_SEARCH = "GridSearch"
+MAX_NUMBER_OF_TRAINING_JOBS_NOT_IMPROVING = "MaxNumberOfTrainingJobsNotImproving"
+BEST_OBJECTIVE_NOT_IMPROVING = "BestObjectiveNotImproving"
+CONVERGENCE_DETECTED = "ConvergenceDetected"
+COMPLETE_ON_CONVERGENCE_DETECTED = "CompleteOnConvergence"
+TARGET_OBJECTIVE_METRIC_VALUE = "TargetObjectiveMetricValue"
+MAX_RUNTIME_IN_SECONDS = "MaxRuntimeInSeconds"
 
 logger = logging.getLogger(__name__)
 
@@ -383,6 +389,116 @@ def to_input_req(self):
         }
 
 
+class TuningJobCompletionCriteriaConfig(object):
+    """The configuration for a job completion criteria."""
+
+    def __init__(
+        self,
+        max_number_of_training_jobs_not_improving: int = None,
+        complete_on_convergence: bool = None,
+        target_objective_metric_value: float = None,
+    ):
+        """Creates a ``TuningJobCompletionCriteriaConfig`` with provided criteria.
+
+        Args:
+            max_number_of_training_jobs_not_improving (int): The number of training jobs that do not
+                improve the best objective after which tuning job will stop.
+            complete_on_convergence (bool): A flag to stop your hyperparameter tuning job if
+                automatic model tuning (AMT) has detected that your model has converged as evaluated
+                against your objective function.
+            target_objective_metric_value (float): The value of the objective metric.
+        """
+
+        self.max_number_of_training_jobs_not_improving = max_number_of_training_jobs_not_improving
+        self.complete_on_convergence = complete_on_convergence
+        self.target_objective_metric_value = target_objective_metric_value
+
+    @classmethod
+    def from_job_desc(cls, completion_criteria_config):
+        """Creates a ``TuningJobCompletionCriteriaConfig`` from a configuration response.
+
+        This is the completion criteria configuration from the DescribeTuningJob response.
+        Args:
+            completion_criteria_config (dict): The expected format of the
+                ``completion_criteria_config`` contains three first-class fields
+
+        Returns:
+            sagemaker.tuner.TuningJobCompletionCriteriaConfig: De-serialized instance of
+            TuningJobCompletionCriteriaConfig containing the completion criteria.
+        """
+        complete_on_convergence = None
+        if CONVERGENCE_DETECTED in completion_criteria_config:
+            if completion_criteria_config[CONVERGENCE_DETECTED][COMPLETE_ON_CONVERGENCE_DETECTED]:
+                complete_on_convergence = bool(
+                    completion_criteria_config[CONVERGENCE_DETECTED][
+                        COMPLETE_ON_CONVERGENCE_DETECTED
+                    ]
+                    == "Enabled"
+                )
+
+        max_number_of_training_jobs_not_improving = None
+        if BEST_OBJECTIVE_NOT_IMPROVING in completion_criteria_config:
+            if completion_criteria_config[BEST_OBJECTIVE_NOT_IMPROVING][
+                MAX_NUMBER_OF_TRAINING_JOBS_NOT_IMPROVING
+            ]:
+                max_number_of_training_jobs_not_improving = completion_criteria_config[
+                    BEST_OBJECTIVE_NOT_IMPROVING
+                ][MAX_NUMBER_OF_TRAINING_JOBS_NOT_IMPROVING]
+
+        target_objective_metric_value = None
+        if TARGET_OBJECTIVE_METRIC_VALUE in completion_criteria_config:
+            target_objective_metric_value = completion_criteria_config[
+                TARGET_OBJECTIVE_METRIC_VALUE
+            ]
+
+        return cls(
+            max_number_of_training_jobs_not_improving=max_number_of_training_jobs_not_improving,
+            complete_on_convergence=complete_on_convergence,
+            target_objective_metric_value=target_objective_metric_value,
+        )
+
+    def to_input_req(self):
+        """Converts the ``self`` instance to the desired input request format.
+
+        Examples:
+            >>> completion_criteria_config = TuningJobCompletionCriteriaConfig(
+                max_number_of_training_jobs_not_improving=5
+                complete_on_convergence = True,
+                target_objective_metric_value = 0.42
+            )
+            >>> completion_criteria_config.to_input_req()
+            {
+                "BestObjectiveNotImproving": {
+                    "MaxNumberOfTrainingJobsNotImproving":5
+                },
+                "ConvergenceDetected": {
+                    "CompleteOnConvergence": "Enabled",
+                },
+                "TargetObjectiveMetricValue": 0.42
+            }
+
+        Returns:
+            dict: Containing the completion criteria configurations.
+        """
+        completion_criteria_config = {}
+        if self.max_number_of_training_jobs_not_improving is not None:
+            completion_criteria_config[BEST_OBJECTIVE_NOT_IMPROVING][
+                MAX_NUMBER_OF_TRAINING_JOBS_NOT_IMPROVING
+            ] = self.max_number_of_training_jobs_not_improving
+
+        if self.target_objective_metric_value is not None:
+            completion_criteria_config[
+                TARGET_OBJECTIVE_METRIC_VALUE
+            ] = self.target_objective_metric_value
+
+        if self.complete_on_convergence is not None:
+            completion_criteria_config[CONVERGENCE_DETECTED][COMPLETE_ON_CONVERGENCE_DETECTED] = (
+                "Enabled" if self.complete_on_convergence else "Disabled"
+            )
+
+        return completion_criteria_config
+
+
 class HyperparameterTuner(object):
     """Defines interaction with Amazon SageMaker hyperparameter tuning jobs.
 
@@ -407,10 +523,12 @@ def __init__(
         objective_type: Union[str, PipelineVariable] = "Maximize",
         max_jobs: Union[int, PipelineVariable] = None,
         max_parallel_jobs: Union[int, PipelineVariable] = 1,
+        max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None,
         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,
+        completion_criteria_config: Optional[TuningJobCompletionCriteriaConfig] = None,
         early_stopping_type: Union[str, PipelineVariable] = "Off",
         estimator_name: Optional[str] = None,
         random_seed: Optional[int] = None,
@@ -450,6 +568,8 @@ def __init__(
                 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).
+            max_runtime_in_seconds (int or PipelineVariable): The maximum time in seconds
+                that a training job launched by a hyperparameter tuning job can run.
             tags (list[dict[str, str] or list[dict[str, PipelineVariable]]): List of tags for
                 labeling the tuning job (default: None). For more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
@@ -463,6 +583,8 @@ def __init__(
                 configuration defining the nature of warm start tuning job.
             strategy_config (sagemaker.tuner.StrategyConfig): A configuration for "Hyperparameter"
                 tuning job optimisation strategy.
+            completion_criteria_config (sagemaker.tuner.TuningJobCompletionCriteriaConfig): A
+                configuration for the completion criteria.
             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.
@@ -505,6 +627,7 @@ def __init__(
 
         self.strategy = strategy
         self.strategy_config = strategy_config
+        self.completion_criteria_config = completion_criteria_config
         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
@@ -513,6 +636,7 @@ def __init__(
         if max_jobs is None and strategy is not GRID_SEARCH:
             self.max_jobs = 1
         self.max_parallel_jobs = max_parallel_jobs
+        self.max_runtime_in_seconds = max_runtime_in_seconds
 
         self.tags = tags
         self.base_tuning_job_name = base_tuning_job_name
@@ -1227,6 +1351,16 @@ def _prepare_init_params_from_job_description(cls, job_details):
             "base_tuning_job_name": base_from_name(job_details["HyperParameterTuningJobName"]),
         }
 
+        if "TuningJobCompletionCriteria" in tuning_config:
+            params["completion_criteria_config"] = TuningJobCompletionCriteriaConfig.from_job_desc(
+                tuning_config["TuningJobCompletionCriteria"]
+            )
+
+        if MAX_RUNTIME_IN_SECONDS in tuning_config["ResourceLimits"]:
+            params["max_runtime_in_seconds"] = tuning_config["ResourceLimits"][
+                MAX_RUNTIME_IN_SECONDS
+            ]
+
         if "RandomSeed" in tuning_config:
             params["random_seed"] = tuning_config["RandomSeed"]
 
@@ -1484,9 +1618,11 @@ def _create_warm_start_tuner(self, additional_parents, warm_start_type, estimato
                 hyperparameter_ranges=self._hyperparameter_ranges,
                 strategy=self.strategy,
                 strategy_config=self.strategy_config,
+                completion_criteria_config=self.completion_criteria_config,
                 objective_type=self.objective_type,
                 max_jobs=self.max_jobs,
                 max_parallel_jobs=self.max_parallel_jobs,
+                max_runtime_in_seconds=self.max_runtime_in_seconds,
                 warm_start_config=WarmStartConfig(
                     warm_start_type=warm_start_type, parents=all_parents
                 ),
@@ -1512,9 +1648,11 @@ def _create_warm_start_tuner(self, additional_parents, warm_start_type, estimato
             metric_definitions_dict=self.metric_definitions_dict,
             strategy=self.strategy,
             strategy_config=self.strategy_config,
+            completion_criteria_config=self.completion_criteria_config,
             objective_type=self.objective_type,
             max_jobs=self.max_jobs,
             max_parallel_jobs=self.max_parallel_jobs,
+            max_runtime_in_seconds=self.max_runtime_in_seconds,
             warm_start_config=WarmStartConfig(warm_start_type=warm_start_type, parents=all_parents),
             early_stopping_type=self.early_stopping_type,
             random_seed=self.random_seed,
@@ -1530,9 +1668,11 @@ def create(
         base_tuning_job_name=None,
         strategy="Bayesian",
         strategy_config=None,
+        completion_criteria_config=None,
         objective_type="Maximize",
         max_jobs=None,
         max_parallel_jobs=1,
+        max_runtime_in_seconds=None,
         tags=None,
         warm_start_config=None,
         early_stopping_type="Off",
@@ -1581,13 +1721,16 @@ def create(
                 (default: 'Bayesian').
             strategy_config (dict): The configuration for a training job launched by a
                 hyperparameter tuning job.
+            completion_criteria_config (dict): The configuration for tuning job completion criteria.
             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. 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).
+            max_runtime_in_seconds (int): The maximum time in seconds
+                that a training job launched by a hyperparameter tuning job can run.
             tags (list[dict]): List of tags for labeling the tuning job (default: None). For more,
                 see https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             warm_start_config (sagemaker.tuner.WarmStartConfig): A ``WarmStartConfig`` object that
@@ -1632,9 +1775,11 @@ def create(
             metric_definitions=metric_definitions,
             strategy=strategy,
             strategy_config=strategy_config,
+            completion_criteria_config=completion_criteria_config,
             objective_type=objective_type,
             max_jobs=max_jobs,
             max_parallel_jobs=max_parallel_jobs,
+            max_runtime_in_seconds=max_runtime_in_seconds,
             tags=tags,
             warm_start_config=warm_start_config,
             early_stopping_type=early_stopping_type,
@@ -1790,6 +1935,9 @@ def _get_tuner_args(cls, tuner, inputs):
             "early_stopping_type": tuner.early_stopping_type,
         }
 
+        if tuner.max_runtime_in_seconds is not None:
+            tuning_config["max_runtime_in_seconds"] = tuner.max_runtime_in_seconds
+
         if tuner.random_seed is not None:
             tuning_config["random_seed"] = tuner.random_seed
 
@@ -1804,6 +1952,11 @@ def _get_tuner_args(cls, tuner, inputs):
         if parameter_ranges is not None:
             tuning_config["parameter_ranges"] = parameter_ranges
 
+        if tuner.completion_criteria_config is not None:
+            tuning_config[
+                "completion_criteria_config"
+            ] = tuner.completion_criteria_config.to_input_req()
+
         tuner_args = {
             "job_name": tuner._current_job_name,
             "tuning_config": tuning_config,