Add HyperparameterTuner.attach() (aws#33)

laurenyu · web-flow · commit 320e7d493d39 · 2018-05-25T16:53:39.000-07:00
diff --git a/src/sagemaker/tuner.py b/src/sagemaker/tuner.py
@@ -12,14 +12,29 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import importlib
 import inspect
 import json
 
 from sagemaker.analytics import HyperparameterTuningJobAnalytics
 from sagemaker.estimator import Framework
 from sagemaker.job import _Job
+from sagemaker.session import Session
 from sagemaker.utils import base_name_from_image, name_from_base
 
+# TODO: probably move these somewhere to Amazon Estimator land after
+# the circular dependency issue is resolved
+AMAZON_ESTIMATOR_MODULE = 'sagemaker'
+AMAZON_ESTIMATOR_CLS_NAMES = {
+    'factorization-machines': 'FactorizationMachines',
+    'kmeans': 'KMeans',
+    'lda': 'LDA',
+    'linear-learner': 'LinearLearner',
+    'ntm': 'NTM',
+    'pca': 'PCA',
+    'randomcutforest': 'RandomCutForest',
+}
+
 
 class _ParameterRange(object):
     __all_types__ = ['Continuous', 'Categorical', 'Integer']
@@ -66,8 +81,11 @@ def __init__(self, min_value, max_value):
 
 
 class HyperparameterTuner(object):
-    SAGEMAKER_ESTIMATOR_CLASS_NAME = 'sagemaker_estimator_class_name'
     SAGEMAKER_ESTIMATOR_MODULE = 'sagemaker_estimator_module'
+    SAGEMAKER_ESTIMATOR_CLASS_NAME = 'sagemaker_estimator_class_name'
+
+    DEFAULT_ESTIMATOR_MODULE = 'sagemaker.estimator'
+    DEFAULT_ESTIMATOR_CLS_NAME = 'Estimator'
 
     def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metric_definitions, strategy='Bayesian',
                  objective_type='Maximize', max_jobs=1, max_parallel_jobs=1, base_tuning_job_name=None):
@@ -100,8 +118,8 @@ def prepare_for_training(self):
         from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 
         if not isinstance(self.estimator, AmazonAlgorithmEstimatorBase):
-            self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_CLASS_NAME] = self.estimator.__class__.__name__
-            self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_MODULE] = self.estimator.__module__
+            self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_CLASS_NAME] = json.dumps(self.estimator.__class__.__name__)
+            self.static_hyperparameters[self.SAGEMAKER_ESTIMATOR_MODULE] = json.dumps(self.estimator.__module__)
 
     def fit(self, inputs, job_name=None, **kwargs):
         """Start a hyperparameter tuning job.
@@ -124,6 +142,24 @@ def fit(self, inputs, job_name=None, **kwargs):
         self.prepare_for_training()
         self.latest_tuning_job = _TuningJob.start_new(self, inputs)
 
+    @classmethod
+    def attach(cls, tuning_job_name, sagemaker_session=None, job_details=None, estimator_cls=None):
+        sagemaker_session = sagemaker_session or Session()
+
+        if job_details is None:
+            job_details = sagemaker_session.sagemaker_client\
+                .describe_hyper_parameter_tuning_job(HyperParameterTuningJobName=tuning_job_name)
+
+        estimator_cls = cls._prepare_estimator_cls(estimator_cls, job_details['TrainingJobDefinition'])
+        estimator = cls._prepare_estimator_from_job_description(estimator_cls, job_details['TrainingJobDefinition'],
+                                                                sagemaker_session)
+        init_params = cls._prepare_init_params_from_job_description(job_details)
+
+        tuner = cls(estimator=estimator, **init_params)
+        tuner.latest_tuning_job = _TuningJob(sagemaker_session=sagemaker_session, tuning_job_name=tuning_job_name)
+
+        return tuner
+
     def deploy(self, initial_instance_count, instance_type, endpoint_name=None, **kwargs):
         """Deploy the best trained or user specified model to an Amazon SageMaker endpoint and return a
         ``sagemaker.RealTimePredictor``
@@ -182,6 +218,75 @@ def _ensure_last_tuning_job(self):
         if self.latest_tuning_job is None:
             raise ValueError('No tuning job available')
 
+    @classmethod
+    def _prepare_estimator_cls(cls, estimator_cls, training_details):
+        # Check for customer-specified estimator first
+        if estimator_cls is not None:
+            module, cls_name = estimator_cls.rsplit('.', 1)
+            return getattr(importlib.import_module(module), cls_name)
+
+        # Then check for estimator class in hyperparameters
+        hyperparameters = training_details['StaticHyperParameters']
+        if cls.SAGEMAKER_ESTIMATOR_CLASS_NAME in hyperparameters and cls.SAGEMAKER_ESTIMATOR_MODULE in hyperparameters:
+            module = hyperparameters.get(cls.SAGEMAKER_ESTIMATOR_MODULE)
+            cls_name = hyperparameters.get(cls.SAGEMAKER_ESTIMATOR_CLASS_NAME)
+            return getattr(importlib.import_module(json.loads(module)), json.loads(cls_name))
+
+        # Then try to derive the estimator from the image name for 1P algorithms
+        image_name = training_details['AlgorithmSpecification']['TrainingImage']
+        algorithm = image_name[image_name.find('/')+1:image_name.find(':')]
+        if algorithm in AMAZON_ESTIMATOR_CLS_NAMES:
+            cls_name = AMAZON_ESTIMATOR_CLS_NAMES[algorithm]
+            return getattr(importlib.import_module(AMAZON_ESTIMATOR_MODULE), cls_name)
+
+        # Default to the BYO estimator
+        return getattr(importlib.import_module(cls.DEFAULT_ESTIMATOR_MODULE), cls.DEFAULT_ESTIMATOR_CLS_NAME)
+
+    @classmethod
+    def _prepare_estimator_from_job_description(cls, estimator_cls, training_details, sagemaker_session):
+        # Swap name for static hyperparameters to what an estimator would expect
+        training_details['HyperParameters'] = training_details['StaticHyperParameters']
+        del training_details['StaticHyperParameters']
+
+        # Remove hyperparameter reserved by SageMaker for tuning jobs
+        del training_details['HyperParameters']['_tuning_objective_metric']
+
+        # Add items expected by the estimator (but aren't needed otherwise)
+        training_details['TrainingJobName'] = ''
+        if 'KmsKeyId' not in training_details['OutputDataConfig']:
+            training_details['OutputDataConfig']['KmsKeyId'] = ''
+
+        estimator_init_params = estimator_cls._prepare_init_params_from_job_description(training_details)
+        return estimator_cls(sagemaker_session=sagemaker_session, **estimator_init_params)
+
+    @classmethod
+    def _prepare_init_params_from_job_description(cls, job_details):
+        tuning_config = job_details['HyperParameterTuningJobConfig']
+        return {
+            'metric_definitions': job_details['TrainingJobDefinition']['AlgorithmSpecification']['MetricDefinitions'],
+            'objective_metric_name': tuning_config['HyperParameterTuningJobObjective']['MetricName'],
+            'objective_type': tuning_config['HyperParameterTuningJobObjective']['Type'],
+            'hyperparameter_ranges': cls._prepare_parameter_ranges(tuning_config['ParameterRanges']),
+            'strategy': tuning_config['Strategy'],
+            'max_jobs': tuning_config['ResourceLimits']['MaxNumberOfTrainingJobs'],
+            'max_parallel_jobs': tuning_config['ResourceLimits']['MaxParallelTrainingJobs'],
+        }
+
+    @classmethod
+    def _prepare_parameter_ranges(cls, parameter_ranges):
+        ranges = {}
+
+        for parameter in parameter_ranges['CategoricalParameterRanges']:
+            ranges[parameter['Name']] = CategoricalParameter(parameter['Values'])
+
+        for parameter in parameter_ranges['ContinuousParameterRanges']:
+            ranges[parameter['Name']] = ContinuousParameter(float(parameter['MinValue']), float(parameter['MaxValue']))
+
+        for parameter in parameter_ranges['IntegerParameterRanges']:
+            ranges[parameter['Name']] = IntegerParameter(int(parameter['MinValue']), int(parameter['MaxValue']))
+
+        return ranges
+
     def hyperparameter_ranges(self):
         """Return collections of ``ParameterRanges``
 
diff --git a/tests/integ/test_tuner.py b/tests/integ/test_tuner.py
@@ -20,7 +20,6 @@
 from sagemaker.amazon.kmeans import KMeans
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.tuner import IntegerParameter, ContinuousParameter, CategoricalParameter, HyperparameterTuner
-from sagemaker.session import s3_input
 from tests.integ import DATA_DIR
 from tests.integ.timeout import timeout
 
diff --git a/tests/unit/test_tuner.py b/tests/unit/test_tuner.py
@@ -12,6 +12,9 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import copy
+import json
+
 import pytest
 from mock import Mock
 
@@ -44,6 +47,80 @@
                          'blank': CategoricalParameter([0, 5])}
 METRIC_DEFINTIONS = 'mock_metric_definitions'
 
+TUNING_JOB_DETAILS = {
+    'HyperParameterTuningJobConfig': {
+        'ResourceLimits': {
+            'MaxParallelTrainingJobs': 1,
+            'MaxNumberOfTrainingJobs': 1
+        },
+        'HyperParameterTuningJobObjective': {
+            'MetricName': OBJECTIVE_METRIC_NAME,
+            'Type': 'Minimize'
+        },
+        'Strategy': 'Bayesian',
+        'ParameterRanges': {
+            'CategoricalParameterRanges': [],
+            'ContinuousParameterRanges': [],
+            'IntegerParameterRanges': [
+                {
+                    'MaxValue': '100',
+                    'Name': 'mini_batch_size',
+                    'MinValue': '10',
+                },
+            ]
+        }
+    },
+    'HyperParameterTuningJobName': JOB_NAME,
+    'TrainingJobDefinition': {
+        'RoleArn': ROLE,
+        'StaticHyperParameters': {
+            'num_components': '1',
+            '_tuning_objective_metric': 'train:throughput',
+            'feature_dim': '784',
+            'sagemaker_estimator_module': '"sagemaker.amazon.pca"',
+            'sagemaker_estimator_class_name': '"PCA"',
+        },
+        'ResourceConfig': {
+            'VolumeSizeInGB': 30,
+            'InstanceType': 'ml.c4.xlarge',
+            'InstanceCount': 1
+        },
+        'AlgorithmSpecification': {
+            'TrainingImage': IMAGE_NAME,
+            'TrainingInputMode': 'File',
+            'MetricDefinitions': METRIC_DEFINTIONS,
+        },
+        'InputDataConfig': [
+            {
+                'ChannelName': 'train',
+                'DataSource': {
+                    'S3DataSource': {
+                        'S3DataDistributionType': 'ShardedByS3Key',
+                        'S3Uri': INPUTS,
+                        'S3DataType': 'ManifestFile'
+                    }
+                }
+            }
+        ],
+        'StoppingCondition': {
+            'MaxRuntimeInSeconds': 86400
+        },
+        'OutputDataConfig': {
+            'S3OutputPath': BUCKET_NAME,
+        }
+    },
+    'TrainingJobCounters': {
+        'ClientError': 0,
+        'Completed': 1,
+        'InProgress': 0,
+        'Fault': 0,
+        'Stopped': 0
+    },
+    'HyperParameterTuningEndTime': 1526605831.0,
+    'CreationTime': 1526605605.0,
+    'HyperParameterTuningJobArn': 'arn:tuning_job',
+}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -73,8 +150,11 @@ def test_prepare_for_training(tuner):
 
     assert len(tuner.static_hyperparameters) == 3
     assert tuner.static_hyperparameters['another_one'] == '0'
-    assert tuner.static_hyperparameters['sagemaker_estimator_class_name'] == tuner.estimator.__class__.__name__
-    assert tuner.static_hyperparameters['sagemaker_estimator_module'] == tuner.estimator.__module__
+
+    class_name = json.dumps(tuner.estimator.__class__.__name__)
+    assert tuner.static_hyperparameters['sagemaker_estimator_class_name'] == class_name
+    module = json.dumps(tuner.estimator.__module__)
+    assert tuner.static_hyperparameters['sagemaker_estimator_module'] == module
 
 
 def test_validate_parameter_ranges_number_validation_error(sagemaker_session):
@@ -127,6 +207,68 @@ def test_fit_1p(sagemaker_session, tuner):
     assert tune_kwargs['job_name'].startswith('pca')
 
 
+def test_attach_tuning_job_with_estimator_from_hyperparameters(sagemaker_session):
+    job_details = copy.deepcopy(TUNING_JOB_DETAILS)
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(name='describe_tuning_job',
+                                                                                  return_value=job_details)
+    tuner = HyperparameterTuner.attach(JOB_NAME, sagemaker_session=sagemaker_session)
+
+    assert tuner.latest_tuning_job.name == JOB_NAME
+    assert tuner.objective_metric_name == OBJECTIVE_METRIC_NAME
+    assert tuner.max_jobs == 1
+    assert tuner.max_parallel_jobs == 1
+    assert tuner.metric_definitions == METRIC_DEFINTIONS
+    assert tuner.strategy == 'Bayesian'
+    assert tuner.objective_type == 'Minimize'
+
+    assert isinstance(tuner.estimator, PCA)
+    assert tuner.estimator.role == ROLE
+    assert tuner.estimator.train_instance_count == 1
+    assert tuner.estimator.train_max_run == 24 * 60 * 60
+    assert tuner.estimator.input_mode == 'File'
+    assert tuner.estimator.output_path == BUCKET_NAME
+    assert tuner.estimator.output_kms_key == ''
+
+    assert '_tuning_objective_metric' not in tuner.estimator.hyperparameters()
+    assert tuner.estimator.hyperparameters()['num_components'] == '1'
+
+
+def test_attach_tuning_job_with_job_details(sagemaker_session):
+    job_details = copy.deepcopy(TUNING_JOB_DETAILS)
+    HyperparameterTuner.attach(JOB_NAME, sagemaker_session=sagemaker_session, job_details=job_details)
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job.assert_not_called
+
+
+def test_attach_tuning_job_with_estimator_from_image(sagemaker_session):
+    job_details = copy.deepcopy(TUNING_JOB_DETAILS)
+    job_details['TrainingJobDefinition']['AlgorithmSpecification']['TrainingImage'] = '1111.amazonaws.com/pca:1'
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(name='describe_tuning_job',
+                                                                                  return_value=job_details)
+
+    tuner = HyperparameterTuner.attach(JOB_NAME, sagemaker_session=sagemaker_session)
+    assert isinstance(tuner.estimator, PCA)
+
+
+def test_attach_tuning_job_with_estimator_from_kwarg(sagemaker_session):
+    job_details = copy.deepcopy(TUNING_JOB_DETAILS)
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(name='describe_tuning_job',
+                                                                                  return_value=job_details)
+    tuner = HyperparameterTuner.attach(JOB_NAME, sagemaker_session=sagemaker_session,
+                                       estimator_cls='sagemaker.estimator.Estimator')
+    assert isinstance(tuner.estimator, Estimator)
+
+
+def test_attach_with_no_specified_estimator(sagemaker_session):
+    job_details = copy.deepcopy(TUNING_JOB_DETAILS)
+    del job_details['TrainingJobDefinition']['StaticHyperParameters']['sagemaker_estimator_module']
+    del job_details['TrainingJobDefinition']['StaticHyperParameters']['sagemaker_estimator_class_name']
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(name='describe_tuning_job',
+                                                                                  return_value=job_details)
+
+    tuner = HyperparameterTuner.attach(JOB_NAME, sagemaker_session=sagemaker_session)
+    assert isinstance(tuner.estimator, Estimator)
+
+
 def test_serialize_parameter_ranges(tuner):
     hyperparameter_ranges = tuner.hyperparameter_ranges()
 
