Make code_location to be S3 URI instead of bucket in training_config()

CHANGELOG.rst

@@ -7,6 +7,7 @@ CHANGELOG
 
 * feature: Estimators: dependencies attribute allows export of additional libraries into the container
 * feature: Add APIs to export Airflow transform and deploy config
+* bug-fix: Allow code_location argument to be S3 URI in training_config API
 
 1.15.0
 ======

src/sagemaker/estimator.py

@@ -670,8 +670,10 @@ def __init__(self, entry_point, source_dir=None, hyperparameters=None, enable_cl
                 training jobs. This will be ignored for now and removed in a further release.
             container_log_level (int): Log level to use within the container (default: logging.INFO).
                 Valid values are defined in the Python logging module.
-            code_location (str): Name of the S3 bucket where custom code is uploaded (default: None).
-                If not specified, default bucket created by ``sagemaker.session.Session`` is used.
+            code_location (str): The S3 prefix URI where custom code will be uploaded (default: None).
+                The code file uploaded in S3 is 'code_location/source/sourcedir.tar.gz'.
+                If not specified, the default code location is s3://default_bucket/job-name/. And code file
+                uploaded to S3 is s3://default_bucket/job-name/source/sourcedir.tar.gz
             image_name (str): An alternate image name to use instead of the official Sagemaker image
                 for the framework. This is useful to run one of the Sagemaker supported frameworks
                 with an image containing custom dependencies.

src/sagemaker/workflow/airflow.py

@@ -27,8 +27,12 @@ def prepare_framework(estimator, s3_operations):
         estimator (sagemaker.estimator.Estimator): The framework estimator to get information from and update.
         s3_operations (dict): The dict to specify s3 operations (upload `source_dir`).
     """
-    bucket = estimator.code_location if estimator.code_location else estimator.sagemaker_session._default_bucket
-    key = '{}/source/sourcedir.tar.gz'.format(estimator._current_job_name)
+    if estimator.code_location is not None:
+        bucket, key = fw_utils.parse_s3_url(estimator.code_location)
+        key = os.path.join(key, 'source', 'sourcedir.tar.gz')
+    else:
+        bucket = estimator.sagemaker_session._default_bucket
+        key = os.path.join(estimator._current_job_name, 'source', 'sourcedir.tar.gz')
     script = os.path.basename(estimator.entry_point)
     if estimator.source_dir and estimator.source_dir.lower().startswith('s3://'):
         code_dir = estimator.source_dir

tests/unit/test_airflow.py

@@ -244,7 +244,7 @@ def test_framework_training_config_all_args(sagemaker_session):
         source_dir="{{ source_dir }}",
         enable_cloudwatch_metrics=False,
         container_log_level="{{ log_level }}",
-        code_location="{{ bucket_name }}",
+        code_location="s3://{{ bucket_name }}/{{ prefix }}",
         training_steps=1000,
         evaluation_steps=100,
         checkpoint_path="{{ checkpoint_path }}",
@@ -304,9 +304,7 @@ def test_framework_training_config_all_args(sagemaker_session):
             'SecurityGroupIds': ['{{ security_group_ids }}']
         },
         'HyperParameters': {
-            'sagemaker_submit_directory': '"s3://{{ bucket_name }}/{{ base_job_name }}-'
-                                          '{{ execution_date.strftime(\'%Y-%m-%d-%H-%M-%S\') }}'
-                                          '/source/sourcedir.tar.gz"',
+            'sagemaker_submit_directory': '"s3://{{ bucket_name }}/{{ prefix }}/source/sourcedir.tar.gz"',
             'sagemaker_program': '"{{ entry_point }}"',
             'sagemaker_enable_cloudwatch_metrics': 'false',
             'sagemaker_container_log_level': '"{{ log_level }}"',
@@ -322,8 +320,7 @@ def test_framework_training_config_all_args(sagemaker_session):
             'S3Upload': [{
                 'Path': '{{ source_dir }}',
                 'Bucket': '{{ bucket_name }}',
-                'Key': "{{ base_job_name }}-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}"
-                       "/source/sourcedir.tar.gz",
+                'Key': "{{ prefix }}/source/sourcedir.tar.gz",
                 'Tar': True}]
         }
     }