Allow code_location to have no key prefix (#227)

* Allow code_location to have no key prefix

Optionally, enable matching the behaviour when using sagemaker's default bucket where a key prefix is not permitted

* Add missing blank line in test_estimator.py

Author: humanzz

CHANGELOG.rst

@@ -7,6 +7,7 @@ CHANGELOG
 
 * bug-fix: Unit Tests: Improve unit test runtime
 * bug-fix: Estimators: Fix attach for LDA
+* bug-fix: Estimators: allow code_location to have no key prefix
 
 1.4.1
 =====

src/sagemaker/estimator.py

@@ -571,7 +571,7 @@ def _stage_user_code_in_s3(self):
             code_s3_prefix = '{}/source'.format(self._current_job_name)
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)
-            code_s3_prefix = '{}/{}/source'.format(key_prefix, self._current_job_name)
+            code_s3_prefix = '/'.join(filter(None, [key_prefix, self._current_job_name, 'source']))
 
         return tar_and_upload_dir(session=self.sagemaker_session.boto_session,
                                   bucket=code_bucket,

tests/unit/test_estimator.py

@@ -127,6 +127,24 @@ def test_custom_code_bucket(time, sagemaker_session):
     assert train_kwargs['hyperparameters']['sagemaker_submit_directory'] == json.dumps(expected_submit_dir)
 
 
+@patch('time.strftime', return_value=TIMESTAMP)
+def test_custom_code_bucket_without_prefix(time, sagemaker_session):
+    code_bucket = 'codebucket'
+    code_location = 's3://{}'.format(code_bucket)
+    t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
+                       train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
+                       code_location=code_location)
+    t.fit('s3://bucket/mydata')
+
+    expected_key = '{}/source/sourcedir.tar.gz'.format(JOB_NAME)
+    _, s3_args, _ = sagemaker_session.boto_session.resource('s3').Object.mock_calls[0]
+    assert s3_args == (code_bucket, expected_key)
+
+    expected_submit_dir = 's3://{}/{}'.format(code_bucket, expected_key)
+    _, _, train_kwargs = sagemaker_session.train.mock_calls[0]
+    assert train_kwargs['hyperparameters']['sagemaker_submit_directory'] == json.dumps(expected_submit_dir)
+
+
 def test_invalid_custom_code_bucket(sagemaker_session):
     code_location = 'thisllworkright?'
     t = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,