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,9 @@ 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).
+            code_location (str): The S3 URI where custom code is uploaded (default: None).
                 If not specified, default bucket created by ``sagemaker.session.Session`` is used.
+                The default S3 path is default_bucket/job-name/source/.
             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}]
         }
     }