fix: don't try to use local output path for KMS key in Local Mode (#1516)

Author: laurenyu

src/sagemaker/estimator.py

@@ -1722,11 +1722,14 @@ def _stage_user_code_in_s3(self):
             code_bucket = self.sagemaker_session.default_bucket()
             code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
             kms_key = None
-
         elif self.code_location is None:
             code_bucket, _ = parse_s3_url(self.output_path)
             code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
             kms_key = self.output_kms_key
+        elif local_mode:
+            code_bucket, key_prefix = parse_s3_url(self.code_location)
+            code_s3_prefix = "/".join(filter(None, [key_prefix, self._current_job_name, "source"]))
+            kms_key = None
         else:
             code_bucket, key_prefix = parse_s3_url(self.code_location)
             code_s3_prefix = "/".join(filter(None, [key_prefix, self._current_job_name, "source"]))

tests/integ/test_local_mode.py

@@ -361,6 +361,7 @@ def test_mxnet_training_failure(sagemaker_local_session, mxnet_full_version, tmp
         train_instance_count=1,
         train_instance_type="local",
         sagemaker_session=sagemaker_local_session,
+        code_location="s3://{}".format(sagemaker_local_session.default_bucket()),
         output_path="file://{}".format(tmpdir),
     )