feature: allow use of short lived creds for local container (#3501)

* build: reset soft

* docs(docs): update docs with USE_SHORT_LIVED_CREDENTIALS bullet

* style(format): fix rst code format

Author: wcarpenter1-godaddy

doc/overview.rst

@@ -1578,6 +1578,7 @@ A few important notes:
 - If you are using S3 data as input, it is pulled from S3 to your local environment. Ensure you have sufficient space to store the data locally.
 - If you run into problems it often due to different Docker containers conflicting. Killing these containers and re-running often solves your problems.
 - Local Mode requires Docker Compose and `nvidia-docker2 <https://github.com/NVIDIA/nvidia-docker>`__ for ``local_gpu``.
+- Set ``USE_SHORT_LIVED_CREDENTIALS=1`` if running on EC2 and you would like to use the session credentials instead of EC2 Metadata Service credentials.
 
 .. warning::
 

src/sagemaker/local/image.py

@@ -1014,7 +1014,7 @@ def _aws_credentials(session):
                 "AWS_ACCESS_KEY_ID=%s" % (str(access_key)),
                 "AWS_SECRET_ACCESS_KEY=%s" % (str(secret_key)),
             ]
-        if not _aws_credentials_available_in_metadata_service():
+        if _use_short_lived_credentials() or not _aws_credentials_available_in_metadata_service():
             logger.warning(
                 "Using the short-lived AWS credentials found in session. They might expire while "
                 "running."
@@ -1052,6 +1052,11 @@ def _aws_credentials_available_in_metadata_service():
     return not instance_metadata_provider.load() is None
 
 
+def _use_short_lived_credentials():
+    """Use short-lived AWS credentials found in session."""
+    return os.environ.get("USE_SHORT_LIVED_CREDENTIALS") == "1"
+
+
 def _write_json_file(filename, content):
     """Write the contents dict as json to the file.
 

tests/unit/test_image.py renamed to tests/unit/sagemaker/local/test_local_image.py

@@ -860,5 +860,25 @@ def test__aws_credentials_with_short_lived_credentials_and_ec2_metadata_service_
     ]
 
 
+@patch("sagemaker.local.image._aws_credentials_available_in_metadata_service")
+def test__aws_credentials_with_short_lived_credentials_and_ec2_metadata_service_having_credentials_override(
+    mock,
+):
+    os.environ["USE_SHORT_LIVED_CREDENTIALS"] = "1"
+    credentials = Credentials(
+        access_key=_random_string(), secret_key=_random_string(), token=_random_string()
+    )
+    session = Mock()
+    session.get_credentials.return_value = credentials
+    mock.return_value = True
+    aws_credentials = _aws_credentials(session)
+
+    assert aws_credentials == [
+        "AWS_ACCESS_KEY_ID=%s" % credentials.access_key,
+        "AWS_SECRET_ACCESS_KEY=%s" % credentials.secret_key,
+        "AWS_SESSION_TOKEN=%s" % credentials.token,
+    ]
+
+
 def _random_string(size=6, chars=string.ascii_uppercase):
     return "".join(random.choice(chars) for x in range(size))