Skip to content
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

change: Allow telemetry only in supported regions #5009

Merged
merged 7 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from 4 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
16 changes: 14 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,19 @@ 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

# Validate region if session exists
if session:
region = _get_region_or_default(session)
try:
Region(region)
except ValueError:
logger.warning(
"Region not found in supported regions. Telemetry request will not be emitted."
)
return
else: # telemetry will be sent to us-west-2 if no session available
region = DEFAULT_AWS_REGION
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isnt this redundant? This condition is already covered in _get_region_or_default

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so, _get_region_or_default covers the condition where session is available but the region value is not allocated.
The else part takes care of the condition when session details are not available.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems like we are assigning the same thing in 2 different places which is something we can simplify.

By removing the if session condition here in line195 and handling everything within _get_region_or_default

If session is None , it will automatically raise an Exception and will get assigned the default value in L#278.

url = _construct_url(
accountId,
region,
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()
Loading