feature: Allow setting S3 endpoint URL for Local Session (#1359)

Author: larsll

src/sagemaker/estimator.py

@@ -1713,6 +1713,7 @@ def _stage_user_code_in_s3(self):
             directory=self.source_dir,
             dependencies=self.dependencies,
             kms_key=kms_key,
+            s3_resource=self.sagemaker_session.s3_resource,
         )
 
     def _model_source_dir(self):

src/sagemaker/fw_utils.py

@@ -307,7 +307,14 @@ def validate_source_dir(script, directory):
 
 
 def tar_and_upload_dir(
-    session, bucket, s3_key_prefix, script, directory=None, dependencies=None, kms_key=None
+    session,
+    bucket,
+    s3_key_prefix,
+    script,
+    directory=None,
+    dependencies=None,
+    kms_key=None,
+    s3_resource=None,
 ):
     """Package source files and upload a compress tar file to S3. The S3
     location will be ``s3://<bucket>/s3_key_prefix/sourcedir.tar.gz``.
@@ -331,6 +338,9 @@ def tar_and_upload_dir(
             copied into /opt/ml/lib
         kms_key (str): Optional. KMS key ID used to upload objects to the bucket
             (default: None).
+        s3_resource (boto3.resource("s3")): Optional. Pre-instantiated Boto3 Resource
+            for S3 connections, can be used to customize the configuration,
+            e.g. set the endpoint URL (default: None).
     Returns:
         sagemaker.fw_utils.UserCode: An object with the S3 bucket and key (S3 prefix) and
             script name.
@@ -354,7 +364,12 @@ def tar_and_upload_dir(
         else:
             extra_args = None
 
-        session.resource("s3").Object(bucket, key).upload_file(tar_file, ExtraArgs=extra_args)
+        if s3_resource is None:
+            s3_resource = session.resource("s3", region_name=session.region_name)
+        else:
+            print("Using provided s3_resource")
+
+        s3_resource.Object(bucket, key).upload_file(tar_file, ExtraArgs=extra_args)
     finally:
         shutil.rmtree(tmp)
 

src/sagemaker/local/image.py

@@ -46,6 +46,7 @@
 # Environment variables to be set during training
 REGION_ENV_NAME = "AWS_REGION"
 TRAINING_JOB_NAME_ENV_NAME = "TRAINING_JOB_NAME"
+S3_ENDPOINT_URL_ENV_NAME = "S3_ENDPOINT_URL"
 
 logger = logging.getLogger(__name__)
 
@@ -139,6 +140,11 @@ def train(self, input_data_config, output_data_config, hyperparameters, job_name
             REGION_ENV_NAME: self.sagemaker_session.boto_region_name,
             TRAINING_JOB_NAME_ENV_NAME: job_name,
         }
+        if self.sagemaker_session.s3_resource is not None:
+            training_env_vars[
+                S3_ENDPOINT_URL_ENV_NAME
+            ] = self.sagemaker_session.s3_resource.meta.client._endpoint.host
+
         compose_data = self._generate_compose_file(
             "train", additional_volumes=volumes, additional_env_vars=training_env_vars
         )
@@ -206,6 +212,7 @@ def serve(self, model_dir, environment):
             "serve", additional_env_vars=environment, additional_volumes=volumes
         )
         compose_command = self._compose()
+
         self.container = _HostingContainer(compose_command)
         self.container.start()
 

src/sagemaker/local/local_session.py

@@ -98,6 +98,7 @@ def create_training_job(
         )
         training_job = _LocalTrainingJob(container)
         hyperparameters = kwargs["HyperParameters"] if "HyperParameters" in kwargs else {}
+        logger.info("Starting training job")
         training_job.start(InputDataConfig, OutputDataConfig, hyperparameters, TrainingJobName)
 
         LocalSagemakerClient._training_jobs[TrainingJobName] = training_job
@@ -377,7 +378,9 @@ def invoke_endpoint(
 class LocalSession(Session):
     """Placeholder docstring"""
 
-    def __init__(self, boto_session=None):
+    def __init__(self, boto_session=None, s3_endpoint_url=None):
+        self.s3_endpoint_url = s3_endpoint_url
+
         super(LocalSession, self).__init__(boto_session)
 
         if platform.system() == "Windows":
@@ -407,6 +410,10 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
         self.local_mode = True
 
+        if self.s3_endpoint_url is not None:
+            self.s3_resource = boto_session.resource("s3", endpoint_url=self.s3_endpoint_url)
+            self.s3_client = boto_session.client("s3", endpoint_url=self.s3_endpoint_url)
+
         sagemaker_config_file = os.path.join(os.path.expanduser("~"), ".sagemaker", "config.yaml")
         if os.path.exists(sagemaker_config_file):
             try:

src/sagemaker/multidatamodel.py

@@ -92,7 +92,13 @@ def __init__(
         self.model = model
         self.container_mode = MULTI_MODEL_CONTAINER_MODE
         self.sagemaker_session = sagemaker_session or Session()
-        self.s3_client = self.sagemaker_session.boto_session.client("s3")
+
+        if self.sagemaker_session.s3_client is None:
+            self.s3_client = self.sagemaker_session.boto_session.client(
+                "s3", region_name=self.sagemaker_session.boto_session.region_name
+            )
+        else:
+            self.s3_client = self.sagemaker_session.s3_client
 
         # Set the ``Model`` parameters if the model parameter is not specified
         if not self.model:

src/sagemaker/session.py

@@ -107,8 +107,8 @@ def __init__(
         """
         self._default_bucket = None
         self._default_bucket_name_override = default_bucket
-
-        # currently is used for local_code in local mode
+        self.s3_resource = None
+        self.s3_client = None
         self.config = None
 
         self._initialize(
@@ -199,7 +199,10 @@ def upload_data(self, path, bucket=None, key_prefix="data", extra_args=None):
             key_suffix = name
 
         bucket = bucket or self.default_bucket()
-        s3 = self.boto_session.resource("s3")
+        if self.s3_resource is None:
+            s3 = self.boto_session.resource("s3", region_name=self.boto_region_name)
+        else:
+            s3 = self.s3_resource
 
         for local_path, s3_key in files:
             s3.Object(bucket, s3_key).upload_file(local_path, ExtraArgs=extra_args)
@@ -227,7 +230,11 @@ def upload_string_as_file_body(self, body, bucket, key, kms_key=None):
             str: The S3 URI of the uploaded file.
                 The URI format is: ``s3://{bucket name}/{key}``.
         """
-        s3 = self.boto_session.resource("s3")
+        if self.s3_resource is None:
+            s3 = self.boto_session.resource("s3", region_name=self.boto_region_name)
+        else:
+            s3 = self.s3_resource
+
         s3_object = s3.Object(bucket_name=bucket, key=key)
 
         if kms_key is not None:
@@ -254,7 +261,10 @@ def download_data(self, path, bucket, key_prefix="", extra_args=None):
 
         """
         # Initialize the S3 client.
-        s3 = self.boto_session.client("s3")
+        if self.s3_client is None:
+            s3 = self.boto_session.client("s3", region_name=self.boto_region_name)
+        else:
+            s3 = self.s3_client
 
         # Initialize the variables used to loop through the contents of the S3 bucket.
         keys = []
@@ -299,7 +309,10 @@ def read_s3_file(self, bucket, key_prefix):
             str: The body of the s3 file as a string.
 
         """
-        s3 = self.boto_session.client("s3")
+        if self.s3_client is None:
+            s3 = self.boto_session.client("s3", region_name=self.boto_region_name)
+        else:
+            s3 = self.s3_client
 
         # Explicitly passing a None kms_key to boto3 throws a validation error.
         s3_object = s3.get_object(Bucket=bucket, Key=key_prefix)
@@ -317,7 +330,10 @@ def list_s3_files(self, bucket, key_prefix):
             [str]: The list of files at the S3 path.
 
         """
-        s3 = self.boto_session.resource("s3")
+        if self.s3_resource is None:
+            s3 = self.boto_session.resource("s3", region_name=self.boto_region_name)
+        else:
+            s3 = self.s3_resource
 
         s3_bucket = s3.Bucket(name=bucket)
         s3_objects = s3_bucket.objects.filter(Prefix=key_prefix).all()
@@ -330,6 +346,7 @@ def default_bucket(self):
             str: The name of the default bucket, which is of the form:
                 ``sagemaker-{region}-{AWS account ID}``.
         """
+
         if self._default_bucket:
             return self._default_bucket
 
@@ -364,10 +381,14 @@ def _create_s3_bucket_if_it_does_not_exist(self, bucket_name, region):
                 already being created, no exception is raised.
 
         """
-        bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
+        if self.s3_resource is None:
+            s3 = self.boto_session.resource("s3", region_name=region)
+        else:
+            s3 = self.s3_resource
+
+        bucket = s3.Bucket(name=bucket_name)
         if bucket.creation_date is None:
             try:
-                s3 = self.boto_session.resource("s3", region_name=region)
                 if region == "us-east-1":
                     # 'us-east-1' cannot be specified because it is the default region:
                     # https://github.com/boto/boto3/issues/125

tests/unit/test_airflow.py

@@ -36,6 +36,8 @@ def sagemaker_session():
         boto_region_name=REGION,
         config=None,
         local_mode=False,
+        s3_resource=None,
+        s3_client=None,
     )
     session.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
     session._default_bucket = BUCKET_NAME

tests/unit/test_chainer.py

@@ -62,6 +62,8 @@ def sagemaker_session():
         boto_region_name=REGION,
         config=None,
         local_mode=False,
+        s3_resource=None,
+        s3_client=None,
     )
 
     describe = {"ModelArtifacts": {"S3ModelArtifacts": "s3://m/m.tar.gz"}}

tests/unit/test_estimator.py

@@ -174,6 +174,8 @@ def sagemaker_session():
         boto_region_name=REGION,
         config=None,
         local_mode=False,
+        s3_client=None,
+        s3_resource=None,
     )
     sms.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(

tests/unit/test_fm.py

@@ -55,6 +55,8 @@ def sagemaker_session():
         region_name=REGION,
         config=None,
         local_mode=False,
+        s3_client=False,
+        s3_resource=False,
     )
     sms.boto_region_name = REGION
     sms.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)