Add support for automatic model tuner's warm start jobs (#489)

Author: uditbhatia

CHANGELOG.rst

@@ -10,6 +10,7 @@ CHANGELOG
 * feature: Add APIs to export Airflow training and tuning config
 * doc-fix: Fix typos in tensorflow serving documentation
 * doc-fix: Add estimator base classes to API docs
+* feature: HyperparameterTuner: add support for Automatic Model Tuning's Warm Start Jobs
 
 1.14.2
 ======

src/sagemaker/session.py

@@ -279,7 +279,8 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
     def tune(self, job_name, strategy, objective_type, objective_metric_name,
              max_jobs, max_parallel_jobs, parameter_ranges,
              static_hyperparameters, image, input_mode, metric_definitions,
-             role, input_config, output_config, resource_config, stop_condition, tags):
+             role, input_config, output_config, resource_config, stop_condition, tags,
+             warm_start_config):
         """Create an Amazon SageMaker hyperparameter tuning job
 
         Args:
@@ -323,6 +324,8 @@ def tune(self, job_name, strategy, objective_type, objective_metric_name,
             stop_condition (dict): When training should finish, e.g. ``MaxRuntimeInSeconds``.
             tags (list[dict]): List of tags for labeling the tuning job. For more, see
                 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.
         """
         tune_request = {
             'HyperParameterTuningJobName': job_name,
@@ -352,6 +355,9 @@ def tune(self, job_name, strategy, objective_type, objective_metric_name,
             }
         }
 
+        if warm_start_config:
+            tune_request['WarmStartConfig'] = warm_start_config
+
         if metric_definitions is not None:
             tune_request['TrainingJobDefinition']['AlgorithmSpecification']['MetricDefinitions'] = metric_definitions
 

src/sagemaker/tuner.py

@@ -15,6 +15,7 @@
 import importlib
 import inspect
 import json
+from enum import Enum
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, RecordSet
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
@@ -35,6 +36,9 @@
     'knn': 'KNN',
     'object2vec': 'Object2Vec',
 }
+HYPERPARAMETER_TUNING_JOB_NAME = 'HyperParameterTuningJobName'
+PARENT_HYPERPARAMETER_TUNING_JOBS = 'ParentHyperParameterTuningJobs'
+WARM_START_TYPE = 'WarmStartType'
 
 
 class _ParameterRange(object):
@@ -134,6 +138,113 @@ class IntegerParameter(_ParameterRange):
     __name__ = 'Integer'
 
 
+class WarmStartTypes(Enum):
+    """Warm Start Configuration type. There can be two types of warm start jobs:
+        * IdenticalDataAndAlgorithm: Type of warm start that allows users to reuse training results from existing
+        tuning jobs that have the same algorithm code and datasets.
+        * TransferLearning: Type of warm start that allows users to reuse training results from existing tuning jobs
+        that have similar algorithm code and datasets.
+    """
+    IDENTICAL_DATA_AND_ALGORITHM = "IdenticalDataAndAlgorithm"
+    TRANSFER_LEARNING = "TransferLearning"
+
+
+class WarmStartConfig(object):
+    """Warm Start Configuration which defines the nature of the warm start ``HyperparameterTuner``, with type and
+    parents for warm start.
+
+    Examples:
+        >>> warm_start_config = WarmStartConfig(type=WarmStartTypes.TransferLearning, parents={"p1","p2"})
+        >>> warm_start_config.type
+        "TransferLearning"
+        >>> warm_start_config.parents
+        {"p1","p2"}
+    """
+
+    def __init__(self, warm_start_type, parents):
+        """Initializes the ``WarmStartConfig`` with the provided ``WarmStartTypes`` and parents.
+
+        Args:
+            warm_start_type (sagemaker.tuner.WarmStartTypes): This should be one of the supported warm start types
+            in WarmStartType
+            parents (set{str}): Set of parent tuning jobs which will be used to warm start the new tuning job.
+        """
+
+        if warm_start_type not in WarmStartTypes:
+            raise ValueError(
+                "Invalid type: {}, valid warm start types are: [{}]".format(warm_start_type,
+                                                                            [t for t in WarmStartTypes]))
+
+        if not parents:
+            raise ValueError("Invalid parents: {}, parents should not be None/empty".format(parents))
+
+        self.type = warm_start_type
+        self.parents = set(parents)
+
+    @classmethod
+    def from_job_desc(cls, warm_start_config):
+        """Creates an instance of ``WarmStartConfig`` class, from warm start configuration response from
+        DescribeTrainingJob.
+
+        Args:
+            warm_start_config (dict): The expected format of the ``warm_start_config`` contains two first-class
+            fields:
+                * "type": Type of warm start tuner, currently two supported types - "IdenticalDataAndAlgorithm" and
+                "TransferLearning".
+                * "parents": List of tuning job names from which the warm start should be done.
+
+        Returns:
+            sagemaker.tuner.WarmStartConfig: De-serialized instance of WarmStartConfig containing the type and parents
+            provided as part of ``warm_start_config``.
+
+        Examples:
+            >>> warm_start_config = WarmStartConfig.from_job_desc(warm_start_config={
+            >>>    "WarmStartType":"TransferLearning",
+            >>>    "ParentHyperParameterTuningJobs": [
+            >>>        {'HyperParameterTuningJobName': "p1"},
+            >>>        {'HyperParameterTuningJobName': "p2"},
+            >>>    ]
+            >>>})
+            >>> warm_start_config.type
+            "TransferLearning"
+            >>> warm_start_config.parents
+            ["p1","p2"]
+        """
+        if not warm_start_config or \
+                WARM_START_TYPE not in warm_start_config or \
+                PARENT_HYPERPARAMETER_TUNING_JOBS not in warm_start_config:
+            return None
+
+        parents = []
+        for parent in warm_start_config[PARENT_HYPERPARAMETER_TUNING_JOBS]:
+            parents.append(parent[HYPERPARAMETER_TUNING_JOB_NAME])
+
+        return cls(warm_start_type=WarmStartTypes(warm_start_config[WARM_START_TYPE]),
+                   parents=parents)
+
+    def to_input_req(self):
+        """Converts the ``self`` instance to the desired input request format.
+
+        Returns:
+            dict: Containing the "WarmStartType" and "ParentHyperParameterTuningJobs" as the first class fields.
+
+        Examples:
+            >>> warm_start_config = WarmStartConfig(warm_start_type=WarmStartTypes.TransferLearning,parents=["p1,p2"])
+            >>> warm_start_config.to_input_req()
+            {
+                "WarmStartType":"TransferLearning",
+                "ParentHyperParameterTuningJobs": [
+                    {'HyperParameterTuningJobName': "p1"},
+                    {'HyperParameterTuningJobName': "p2"},
+                ]
+            }
+        """
+        return {
+            WARM_START_TYPE: self.type.value,
+            PARENT_HYPERPARAMETER_TUNING_JOBS: [{HYPERPARAMETER_TUNING_JOB_NAME: parent} for parent in self.parents]
+        }
+
+
 class HyperparameterTuner(object):
     """A class for creating and interacting with Amazon SageMaker hyperparameter tuning jobs, as well as
     deploying the resulting model(s).
@@ -148,7 +259,7 @@ class HyperparameterTuner(object):
 
     def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metric_definitions=None,
                  strategy='Bayesian', objective_type='Maximize', max_jobs=1, max_parallel_jobs=1,
-                 tags=None, base_tuning_job_name=None):
+                 tags=None, base_tuning_job_name=None, warm_start_config=None):
         """Initialize a ``HyperparameterTuner``. It takes an estimator to obtain configuration information
         for training jobs that are created as the result of a hyperparameter tuning job.
 
@@ -175,6 +286,8 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
             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 generaged, based on the training image name and current timestamp.
+            warm_start_config (sagemaker.tuner.WarmStartConfig): A ``WarmStartConfig`` object that has been initialized
+                with the configuration defining the nature of warm start tuning job.
         """
         self._hyperparameter_ranges = hyperparameter_ranges
         if self._hyperparameter_ranges is None or len(self._hyperparameter_ranges) == 0:
@@ -194,6 +307,7 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
         self.base_tuning_job_name = base_tuning_job_name
         self._current_job_name = None
         self.latest_tuning_job = None
+        self.warm_start_config = warm_start_config
 
     def _prepare_for_training(self, job_name=None, include_cls_metadata=True):
         if job_name is not None:
@@ -419,6 +533,7 @@ def _prepare_init_params_from_job_description(cls, job_details):
             'strategy': tuning_config['Strategy'],
             'max_jobs': tuning_config['ResourceLimits']['MaxNumberOfTrainingJobs'],
             'max_parallel_jobs': tuning_config['ResourceLimits']['MaxParallelTrainingJobs'],
+            'warm_start_config': WarmStartConfig.from_job_desc(job_details.get('WarmStartConfig', None))
         }
 
     @classmethod
@@ -489,6 +604,82 @@ def _validate_parameter_ranges(self):
                     except KeyError:
                         pass
 
+    def transfer_learning_tuner(self, additional_parents=None, estimator=None):
+        """Creates a new ``HyperparameterTuner`` by copying the request fields from the provided parent to the new
+        instance of ``HyperparameterTuner``. Followed by addition of warm start configuration with the type as
+        "TransferLearning" and parents as the union of provided list of ``additional_parents`` and the ``self``.
+        Also, training image in the new tuner's estimator is updated with the provided ``training_image``.
+
+        Args:
+            additional_parents (set{str}): Set of additional parents along with the self to be used in warm starting
+            the transfer learning tuner.
+            estimator (sagemaker.estimator.EstimatorBase): An estimator object that has been initialized with
+                the desired configuration. There does not need to be a training job associated with this instance.
+
+        Returns:
+            sagemaker.tuner.HyperparameterTuner: ``HyperparameterTuner`` instance which can be used to launch transfer
+            learning tuning job.
+
+        Examples:
+            >>> parent_tuner = HyperparameterTuner.attach(tuning_job_name="parent-job-1")
+            >>> transfer_learning_tuner = parent_tuner.transfer_learning_tuner(additional_parents={"parent-job-2"})
+            Later On:
+            >>> transfer_learning_tuner.fit(inputs={})
+        """
+
+        return self._create_warm_start_tuner(additional_parents=additional_parents,
+                                             warm_start_type=WarmStartTypes.TRANSFER_LEARNING,
+                                             estimator=estimator)
+
+    def identical_dataset_and_algorithm_tuner(self, additional_parents=None):
+        """Creates a new ``HyperparameterTuner`` by copying the request fields from the provided parent to the new
+        instance of ``HyperparameterTuner``. Followed by addition of warm start configuration with the type as
+        "IdenticalDataAndAlgorithm" and parents as the union of provided list of ``additional_parents`` and the ``self``
+
+        Args:
+            additional_parents (set{str}): Set of additional parents along with the self to be used in warm starting
+            the identical dataset and algorithm tuner.
+
+        Returns:
+            sagemaker.tuner.HyperparameterTuner: HyperparameterTuner instance which can be used to launch identical
+            dataset and algorithm tuning job.
+
+        Examples:
+            >>> parent_tuner = HyperparameterTuner.attach(tuning_job_name="parent-job-1")
+            >>> identical_dataset_algo_tuner = parent_tuner.identical_dataset_and_algorithm_tuner(
+            >>>                                                             additional_parents={"parent-job-2"})
+            Later On:
+            >>> identical_dataset_algo_tuner.fit(inputs={})
+        """
+
+        return self._create_warm_start_tuner(additional_parents=additional_parents,
+                                             warm_start_type=WarmStartTypes.IDENTICAL_DATA_AND_ALGORITHM)
+
+    def _create_warm_start_tuner(self, additional_parents, warm_start_type, estimator=None):
+        """Creates a new ``HyperparameterTuner`` with ``WarmStartConfig``, where type will be equal to
+        ``warm_start_type`` and``parents`` would be equal to union of ``additional_parents`` and self.
+
+        Args:
+            additional_parents (set{str}): Additional parents along with self, to be used for warm starting.
+            warm_start_type (sagemaker.tuner.WarmStartTypes): Type of warm start job.
+
+        Returns:
+            sagemaker.tuner.HyperparameterTuner: Instance with the request fields copied from self along with the
+            warm start configuration
+        """
+        all_parents = {self.latest_tuning_job.name}
+        if additional_parents:
+            all_parents = all_parents.union(additional_parents)
+
+        return HyperparameterTuner(estimator=estimator if estimator else self.estimator,
+                                   objective_metric_name=self.objective_metric_name,
+                                   hyperparameter_ranges=self._hyperparameter_ranges,
+                                   objective_type=self.objective_type,
+                                   max_jobs=self.max_jobs,
+                                   max_parallel_jobs=self.max_parallel_jobs,
+                                   warm_start_config=WarmStartConfig(warm_start_type=warm_start_type,
+                                                                     parents=all_parents))
+
 
 class _TuningJob(_Job):
     @classmethod
@@ -504,6 +695,10 @@ def start_new(cls, tuner, inputs):
         """
         config = _Job._load_config(inputs, tuner.estimator)
 
+        warm_start_config_req = None
+        if tuner.warm_start_config:
+            warm_start_config_req = tuner.warm_start_config.to_input_req()
+
         tuner.estimator.sagemaker_session.tune(job_name=tuner._current_job_name, strategy=tuner.strategy,
                                                objective_type=tuner.objective_type,
                                                objective_metric_name=tuner.objective_metric_name,
@@ -516,7 +711,8 @@ def start_new(cls, tuner, inputs):
                                                role=(config['role']), input_config=(config['input_config']),
                                                output_config=(config['output_config']),
                                                resource_config=(config['resource_config']),
-                                               stop_condition=(config['stop_condition']), tags=tuner.tags)
+                                               stop_condition=(config['stop_condition']), tags=tuner.tags,
+                                               warm_start_config=warm_start_config_req)
 
         return cls(tuner.sagemaker_session, tuner._current_job_name)
 
@@ -525,3 +721,50 @@ def stop(self):
 
     def wait(self):
         self.sagemaker_session.wait_for_tuning_job(self.name)
+
+
+def create_identical_dataset_and_algorithm_tuner(parent, additional_parents=None, sagemaker_session=None):
+    """Creates a new tuner by copying the request fields from the provided parent to the new instance of
+        ``HyperparameterTuner`` followed by addition of warm start configuration with the type as
+        "IdenticalDataAndAlgorithm" and ``parents`` as the union of provided list of ``additional_parents`` and the
+        ``parent``.
+
+    Args:
+        parent (str): Primary parent tuning job's name from which the Tuner and Estimator configuration has to be copied
+        additional_parents (set{str}): Set of additional parent tuning job's names along with the primary parent tuning
+            job name to be used in warm starting the transfer learning tuner.
+        sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
+                Amazon SageMaker APIs and any other AWS services needed. If not specified, one is created
+                using the default AWS configuration chain.
+
+    Returns:
+        sagemaker.tuner.HyperparameterTuner: a new ``HyperparameterTuner`` object for the warm-started
+        hyperparameter tuning job
+    """
+
+    parent_tuner = HyperparameterTuner.attach(tuning_job_name=parent, sagemaker_session=sagemaker_session)
+    return parent_tuner.identical_dataset_and_algorithm_tuner(additional_parents=additional_parents)
+
+
+def create_transfer_learning_tuner(parent, additional_parents=None, estimator=None, sagemaker_session=None):
+    """Creates a new ``HyperParameterTuner`` by copying the request fields from the provided parent to the new instance
+        of ``HyperparameterTuner`` followed by addition of warm start configuration with the type as "TransferLearning"
+        and ``parents`` as the union of provided list of ``additional_parents`` and the ``parent``.
+
+    Args:
+        parent (str): Primary parent tuning job's name from which the Tuner and Estimator configuration has to be copied
+        additional_parents (set{str}): Set of additional parent tuning job's names along with the primary parent tuning
+            job name to be used in warm starting the identical dataset and algorithm tuner.
+        estimator (sagemaker.estimator.EstimatorBase): An estimator object that has been initialized with
+                the desired configuration. There does not need to be a training job associated with this instance.
+        sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
+                Amazon SageMaker APIs and any other AWS services needed. If not specified, one is created
+                using the default AWS configuration chain.
+
+    Returns:
+        sagemaker.tuner.HyperparameterTuner: New instance of warm started HyperparameterTuner
+    """
+
+    parent_tuner = HyperparameterTuner.attach(tuning_job_name=parent, sagemaker_session=sagemaker_session)
+    return parent_tuner.transfer_learning_tuner(additional_parents=additional_parents,
+                                                estimator=estimator)