From 25f16ef757e7874b925b426ecf5d1130206d3c2c Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 28 Jan 2025 23:57:10 -0800 Subject: [PATCH 1/5] change: Allow telemetry only in supported regions --- src/sagemaker/telemetry/constants.py | 35 ++++++++++++++++++ src/sagemaker/telemetry/telemetry_logging.py | 16 +++++++-- .../telemetry/test_telemetry_logging.py | 36 +++++++++++++++++++ 3 files changed, 85 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index 2108ff9fd6..28bb758b05 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -42,3 +42,38 @@ class Status(Enum): def __str__(self): # pylint: disable=E0307 """Return the status name.""" return self.name + + +class Region(str, Enum): + # 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 index b45550b2c2..55c0e205d9 100644 --- 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,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.debug( + "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 url = _construct_url( accountId, region, diff --git a/tests/unit/sagemaker/telemetry/test_telemetry_logging.py b/tests/unit/sagemaker/telemetry/test_telemetry_logging.py index 9107256b5b..bd8db82a16 100644 --- 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() From 0ed85d67e6cfd8bfe39529a12aaa5dfe43785835 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 28 Jan 2025 23:57:10 -0800 Subject: [PATCH 2/5] change: Allow telemetry only in supported regions --- src/sagemaker/telemetry/constants.py | 37 ++++++++++++++++++++ src/sagemaker/telemetry/telemetry_logging.py | 2 +- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index 28bb758b05..a18a4a4a0f 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -44,6 +44,43 @@ def __str__(self): # pylint: disable=E0307 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 + + class Region(str, Enum): # Classic US_EAST_1 = "us-east-1" # IAD diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 55c0e205d9..887a574ca1 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -197,7 +197,7 @@ def _send_telemetry_request( try: Region(region) except ValueError: - logger.debug( + logger.warning( "Region not found in supported regions. Telemetry request will not be emitted." ) return From b69ffcb952948626166c35ec4264d6fff8a2ce17 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 28 Jan 2025 23:57:10 -0800 Subject: [PATCH 3/5] change: Allow telemetry only in supported regions --- src/sagemaker/telemetry/constants.py | 34 ---------------------------- 1 file changed, 34 deletions(-) diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index a18a4a4a0f..d6f19dc3d2 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -80,37 +80,3 @@ class Region(str, Enum): IL_CENTRAL_1 = "il-central-1" # TLV ME_CENTRAL_1 = "me-central-1" # DXB - -class Region(str, Enum): - # 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 From 8d7f4a8e3c1645b4d892479cbcbd723951e77081 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Wed, 29 Jan 2025 13:03:28 -0800 Subject: [PATCH 4/5] change: Allow telemetry only in supported regions --- src/sagemaker/telemetry/constants.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/sagemaker/telemetry/constants.py b/src/sagemaker/telemetry/constants.py index d6f19dc3d2..cb83a78279 100644 --- a/src/sagemaker/telemetry/constants.py +++ b/src/sagemaker/telemetry/constants.py @@ -79,4 +79,3 @@ class Region(str, Enum): EU_SOUTH_2 = "eu-south-2" # ZAZ IL_CENTRAL_1 = "il-central-1" # TLV ME_CENTRAL_1 = "me-central-1" # DXB - From dadbb220d42205b4658dde8e09861a4b72389507 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Thu, 30 Jan 2025 12:18:59 -0800 Subject: [PATCH 5/5] change: Allow telemetry only in supported regions --- src/sagemaker/telemetry/telemetry_logging.py | 22 +++++++++----------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/sagemaker/telemetry/telemetry_logging.py b/src/sagemaker/telemetry/telemetry_logging.py index 887a574ca1..b0ecedee4c 100644 --- a/src/sagemaker/telemetry/telemetry_logging.py +++ b/src/sagemaker/telemetry/telemetry_logging.py @@ -190,19 +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" + 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 - # 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 url = _construct_url( accountId, region, @@ -280,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)