do not call create bucket if data location is provided (#96)

aws · Mar 14, 2018 · 17799ce · 17799ce
1 parent 515a6ed
commit 17799ce
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/sagemaker/amazon/amazon_estimator.py b/src/sagemaker/amazon/amazon_estimator.py
@@ -42,8 +42,8 @@ def __init__(self, role, train_instance_count, train_instance_type, data_locatio
         super(AmazonAlgorithmEstimatorBase, self).__init__(role, train_instance_count, train_instance_type,
                                                            **kwargs)
 
-        default_location = "s3://{}/sagemaker-record-sets/".format(self.sagemaker_session.default_bucket())
-        data_location = data_location or default_location
+        data_location = data_location or "s3://{}/sagemaker-record-sets/".format(
+            self.sagemaker_session.default_bucket())
         self.data_location = data_location
 
     def train_image(self):

diff --git a/tests/unit/test_amazon_estimator.py b/tests/unit/test_amazon_estimator.py
@@ -88,6 +88,13 @@ def test_data_location_validation(sagemaker_session):
         pca.data_location = "nots3://abcd/efgh"
 
 
+def test_data_location_does_not_call_default_bucket(sagemaker_session):
+    data_location = "s3://my-bucket/path/"
+    pca = PCA(num_components=2, sagemaker_session=sagemaker_session, data_location=data_location, **COMMON_ARGS)
+    assert pca.data_location == data_location
+    assert not sagemaker_session.default_bucket.called
+
+
 def test_pca_hyperparameters(sagemaker_session):
     pca = PCA(num_components=55, algorithm_mode='randomized',
               subtract_mean=True, extra_components=33, sagemaker_session=sagemaker_session,