change: Allow telemetry only in supported regions (#5009)

* change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions * change: Allow telemetry only in supported regions --------- Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>
aws · Jan 31, 2025 · 90e9c9f · 90e9c9f
1 parent c753da0
commit 90e9c9f
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 2 deletions.
diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py
@@ -42,3 +42,40 @@ class Status(Enum):
     def __str__(self):  # pylint: disable=E0307
         """Return the status name."""
         return self.name
+
+
+class Region(str, Enum):
+    """Telemetry: List of all supported AWS regions."""
+
+    # Classic
+    US_EAST_1 = "us-east-1"  # IAD
+    US_EAST_2 = "us-east-2"  # CMH
+    US_WEST_1 = "us-west-1"  # SFO
+    US_WEST_2 = "us-west-2"  # PDX
+    AP_NORTHEAST_1 = "ap-northeast-1"  # NRT
+    AP_NORTHEAST_2 = "ap-northeast-2"  # ICN
+    AP_NORTHEAST_3 = "ap-northeast-3"  # KIX
+    AP_SOUTH_1 = "ap-south-1"  # BOM
+    AP_SOUTHEAST_1 = "ap-southeast-1"  # SIN
+    AP_SOUTHEAST_2 = "ap-southeast-2"  # SYD
+    CA_CENTRAL_1 = "ca-central-1"  # YUL
+    EU_CENTRAL_1 = "eu-central-1"  # FRA
+    EU_NORTH_1 = "eu-north-1"  # ARN
+    EU_WEST_1 = "eu-west-1"  # DUB
+    EU_WEST_2 = "eu-west-2"  # LHR
+    EU_WEST_3 = "eu-west-3"  # CDG
+    SA_EAST_1 = "sa-east-1"  # GRU
+    # Opt-in
+    AP_EAST_1 = "ap-east-1"  # HKG
+    AP_SOUTHEAST_3 = "ap-southeast-3"  # CGK
+    AF_SOUTH_1 = "af-south-1"  # CPT
+    EU_SOUTH_1 = "eu-south-1"  # MXP
+    ME_SOUTH_1 = "me-south-1"  # BAH
+    MX_CENTRAL_1 = "mx-central-1"  # QRO
+    AP_SOUTHEAST_7 = "ap-southeast-7"  # BKK
+    AP_SOUTH_2 = "ap-south-2"  # HYD
+    AP_SOUTHEAST_4 = "ap-southeast-4"  # MEL
+    EU_CENTRAL_2 = "eu-central-2"  # ZRH
+    EU_SOUTH_2 = "eu-south-2"  # ZAZ
+    IL_CENTRAL_1 = "il-central-1"  # TLV
+    ME_CENTRAL_1 = "me-central-1"  # DXB
diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py
@@ -27,6 +27,7 @@
 from sagemaker.telemetry.constants import (
     Feature,
     Status,
+    Region,
     DEFAULT_AWS_REGION,
 )
 from sagemaker.user_agent import SDK_VERSION, process_studio_metadata_file
@@ -189,8 +190,16 @@ def _send_telemetry_request(
     """Make GET request to an empty object in S3 bucket"""
     try:
         accountId = _get_accountId(session) if session else "NotAvailable"
-        # telemetry will be sent to us-west-2 if no session availale
-        region = _get_region_or_default(session) if session else DEFAULT_AWS_REGION
+        region = _get_region_or_default(session)
+
+        try:
+            Region(region)  # Validate the region
+        except ValueError:
+            logger.warning(
+                "Region not found in supported regions. Telemetry request will not be emitted."
+            )
+            return
+
         url = _construct_url(
             accountId,
             region,
@@ -268,6 +277,7 @@ def _get_region_or_default(session):
 
 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)
 

diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py
@@ -300,3 +300,39 @@ def test_get_default_sagemaker_session_with_no_region(self):
         assert "Must setup local AWS configuration with a region supported by SageMaker." in str(
             context.exception
         )
+
+    @patch("sagemaker.telemetry.telemetry_logging._get_accountId")
+    @patch("sagemaker.telemetry.telemetry_logging._get_region_or_default")
+    def test_send_telemetry_request_valid_region(self, mock_get_region, mock_get_accountId):
+        """Test to verify telemetry request is sent when region is valid"""
+        mock_get_accountId.return_value = "testAccountId"
+        mock_session = MagicMock()
+
+        # Test with valid region
+        mock_get_region.return_value = "us-east-1"
+        with patch(
+            "sagemaker.telemetry.telemetry_logging._requests_helper"
+        ) as mock_requests_helper:
+            _send_telemetry_request(1, [1, 2], mock_session)
+            # Assert telemetry request was sent
+            mock_requests_helper.assert_called_once_with(
+                "https://sm-pysdk-t-us-east-1.s3.us-east-1.amazonaws.com/telemetry?"
+                "x-accountId=testAccountId&x-status=1&x-feature=1,2",
+                2,
+            )
+
+    @patch("sagemaker.telemetry.telemetry_logging._get_accountId")
+    @patch("sagemaker.telemetry.telemetry_logging._get_region_or_default")
+    def test_send_telemetry_request_invalid_region(self, mock_get_region, mock_get_accountId):
+        """Test to verify telemetry request is not sent when region is invalid"""
+        mock_get_accountId.return_value = "testAccountId"
+        mock_session = MagicMock()
+
+        # Test with invalid region
+        mock_get_region.return_value = "invalid-region"
+        with patch(
+            "sagemaker.telemetry.telemetry_logging._requests_helper"
+        ) as mock_requests_helper:
+            _send_telemetry_request(1, [1, 2], mock_session)
+            # Assert telemetry request was not sent
+            mock_requests_helper.assert_not_called()