Skip to content

change: Allow telemetry only in supported regions #5009

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions src/sagemaker/telemetry/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
14 changes: 12 additions & 2 deletions src/sagemaker/telemetry/telemetry_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)

Expand Down
36 changes: 36 additions & 0 deletions tests/unit/sagemaker/telemetry/test_telemetry_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()