fix: Update HPO TrainingInputMode with s3_input InputMode

Author: imujjwal96

src/sagemaker/estimator.py

@@ -571,6 +571,8 @@ def start_new(cls, estimator, inputs):
 
         if isinstance(inputs, s3_input):
             if 'InputMode' in inputs.config:
+                logging.debug('Selecting s3_input\'s input_mode ({}) for TrainingInputMode.'
+                              .format(inputs.config['InputMode']))
                 train_args['input_mode'] = inputs.config['InputMode']
 
         if estimator.enable_network_isolation():

src/sagemaker/tuner.py

@@ -15,6 +15,7 @@
 import importlib
 import inspect
 import json
+import logging
 from enum import Enum
 
 import sagemaker
@@ -26,6 +27,7 @@
 from sagemaker.parameter import (CategoricalParameter, ContinuousParameter,
                                  IntegerParameter, ParameterRange)
 from sagemaker.session import Session
+from sagemaker.session import s3_input
 from sagemaker.utils import base_name_from_image, name_from_base, to_str
 
 AMAZON_ESTIMATOR_MODULE = 'sagemaker'
@@ -640,6 +642,12 @@ def start_new(cls, tuner, inputs):
         tuner_args['warm_start_config'] = warm_start_config_req
         tuner_args['early_stopping_type'] = tuner.early_stopping_type
 
+        if isinstance(inputs, s3_input):
+            if 'InputMode' in inputs.config:
+                logging.debug('Selecting s3_input\'s input_mode ({}) for TrainingInputMode.'
+                              .format(inputs.config['InputMode']))
+                tuner_args['input_mode'] = inputs.config['InputMode']
+
         if isinstance(tuner.estimator, sagemaker.algorithm.AlgorithmEstimator):
             tuner_args['algorithm_arn'] = tuner.estimator.algorithm_arn
         else:

tests/unit/test_tuner.py

@@ -28,6 +28,7 @@
 from sagemaker.tuner import (_TuningJob, create_identical_dataset_and_algorithm_tuner,
                              create_transfer_learning_tuner, HyperparameterTuner, WarmStartConfig,
                              WarmStartTypes)
+from sagemaker.session import s3_input
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
 MODEL_DATA = "s3://bucket/model.tar.gz"
@@ -286,6 +287,31 @@ def test_fit_mxnet_with_vpc_config(sagemaker_session, tuner):
     assert tune_kwargs['vpc_config'] == {'Subnets': subnets, 'SecurityGroupIds': security_group_ids}
 
 
+def test_s3_input_mode(sagemaker_session, tuner):
+    expected_input_mode = 'Pipe'
+
+    script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'failure_script.py')
+    mxnet = MXNet(entry_point=script_path,
+                  role=ROLE,
+                  framework_version=FRAMEWORK_VERSION,
+                  train_instance_count=TRAIN_INSTANCE_COUNT,
+                  train_instance_type=TRAIN_INSTANCE_TYPE,
+                  sagemaker_session=sagemaker_session)
+    tuner.estimator = mxnet
+
+    tags = [{'Name': 'some-tag-without-a-value'}]
+    tuner.tags = tags
+
+    hyperparameter_ranges = {'num_components': IntegerParameter(2, 4),
+                             'algorithm_mode': CategoricalParameter(['regular', 'randomized'])}
+    tuner._hyperparameter_ranges = hyperparameter_ranges
+
+    tuner.fit(inputs=s3_input('s3://mybucket/train_manifest', input_mode=expected_input_mode))
+
+    actual_input_mode = sagemaker_session.method_calls[1][2]['input_mode']
+    assert actual_input_mode == expected_input_mode
+
+
 def test_fit_pca_with_inter_container_traffic_encryption_flag(sagemaker_session, tuner):
     pca = PCA(ROLE, TRAIN_INSTANCE_COUNT, TRAIN_INSTANCE_TYPE, NUM_COMPONENTS,
               base_job_name='pca', sagemaker_session=sagemaker_session,