Fix default session issue

Author: knikure

src/sagemaker/telemetry/telemetry_logging.py

@@ -20,6 +20,7 @@
 import functools
 import requests
 
+import boto3
 from sagemaker.session import Session
 from sagemaker.utils import resolve_value_from_config
 from sagemaker.config.config_schema import TELEMETRY_OPT_OUT_PATH
@@ -71,7 +72,9 @@ def wrapper(*args, **kwargs):
                 sagemaker_session = args[0].sagemaker_session
             elif feature == Feature.REMOTE_FUNCTION:
                 # Get the sagemaker_session from the function keyword arguments for remote function
-                sagemaker_session = kwargs.get("sagemaker_session", Session())
+                sagemaker_session = kwargs.get(
+                    "sagemaker_session", _get_default_sagemaker_session()
+                )
 
             if sagemaker_session:
                 logger.debug("sagemaker_session found, preparing to emit telemetry...")
@@ -243,3 +246,11 @@ def _get_region_or_default(session):
         return session.boto_session.region_name
     except Exception:  # pylint: disable=W0703
         return DEFAULT_AWS_REGION
+
+
+def _get_default_sagemaker_session():
+    """Return the default sagemaker session"""
+    boto_session = boto3.Session(region_name=DEFAULT_AWS_REGION)
+    sagemaker_session = Session(boto_session=boto_session)
+
+    return sagemaker_session

tests/unit/sagemaker/telemetry/test_telemetry_logging.py

@@ -15,6 +15,7 @@
 import pytest
 import requests
 from unittest.mock import Mock, patch, MagicMock
+import boto3
 import sagemaker
 from sagemaker.telemetry.constants import Feature
 from sagemaker.telemetry.telemetry_logging import (
@@ -24,6 +25,7 @@
     _get_accountId,
     _requests_helper,
     _get_region_or_default,
+    _get_default_sagemaker_session,
     OS_NAME_VERSION,
     PYTHON_VERSION,
 )
@@ -282,3 +284,19 @@ def test_get_region_or_default_exception(self):
             region = _get_region_or_default(mock_session)
             assert region == "us-west-2"
             assert "Error creating boto session" in str(exception)
+
+    @patch.object(boto3.Session, "region_name", "us-west-2")
+    def test_get_default_sagemaker_session(self):
+        sagemaker_session = _get_default_sagemaker_session()
+
+        assert sagemaker_session is sagemaker.Session
+        assert sagemaker_session.boto_session.region_name == "us-west-2"
+
+    @patch.object(boto3.Session, "region_name", None)
+    def test_get_default_sagemaker_session_with_no_region(self):
+        with self.assertRaises(ValueError) as context:
+            _get_default_sagemaker_session()
+
+        assert "Must setup local AWS configuration with a region supported by SageMaker." in str(
+            context.exception
+        )