diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index b8f5205e91..d55a390056 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -16,6 +16,7 @@ import os import platform import sys +import threading from time import perf_counter from typing import List import functools @@ -66,6 +67,23 @@ ) _telemetry_msg_shown = False +# Seconds to wait for the telemetry endpoint before giving up on an event. +TELEMETRY_REQUEST_TIMEOUT = 2 + +# Telemetry must never sit in the caller's critical path, so every event is sent +# from a daemon thread. A slow or unreachable telemetry endpoint (for example from +# inside a VPC with no route to it) can no longer stall the SDK call the customer +# actually made, and because the threads are daemons a pending send cannot delay +# interpreter shutdown either. No event is ever dropped: one thread is started per +# event. +# +# Special case worth calling out: Feature Store ingestion is decorated at more than +# one level (``ingest_dataframe`` and ``IngestionManagerPandas.run``), so a single +# user call can emit several events. Sending them serially on the caller's thread is +# what turned an ingest that the service finished in under a second into a +# multi-minute wait, which is why the send is moved off that thread here rather than +# only having its timeout tightened. + FEATURE_TO_CODE = { str(Feature.SDK_DEFAULTS): 11, str(Feature.LOCAL_MODE): 12, @@ -422,6 +440,41 @@ def _send_telemetry_request( failure_reason: str = None, failure_type: str = None, extra_info: str = None, +) -> threading.Thread: + """Schedule a telemetry event to be sent on a background daemon thread. + + Every event is still sent; this only moves the send off the caller's thread so + that telemetry never adds latency to the SDK call that triggered it. + + Returns: + threading.Thread: The thread doing the send. Callers generally ignore this; + tests can join on it. + """ + + def _run(): + try: + _send_telemetry_request_sync( + status, feature_list, session, failure_reason, failure_type, extra_info + ) + except Exception: # pylint: disable=W0703 + # Nothing can be raised out of a fire-and-forget thread: there is no + # caller to catch it, and even the logging call inside + # _send_telemetry_request_sync can fail once the interpreter starts + # tearing down. Telemetry is best-effort, so drop the event silently. + pass + + thread = threading.Thread(target=_run, name="sagemaker-telemetry", daemon=True) + thread.start() + return thread + + +def _send_telemetry_request_sync( + status: int, + feature_list: List[int], + session: Session, + failure_reason: str = None, + failure_type: str = None, + extra_info: str = None, ) -> None: """Make GET request to an empty object in S3 bucket""" try: @@ -449,7 +502,7 @@ def _send_telemetry_request( ) # Send the telemetry request logger.debug("Sending telemetry request to [%s]", url) - _requests_helper(url, 2) + _requests_helper(url, TELEMETRY_REQUEST_TIMEOUT) logger.debug("SageMaker Python SDK telemetry successfully emitted.") except Exception: # pylint: disable=W0703 logger.debug("SageMaker Python SDK telemetry not emitted!") @@ -482,12 +535,17 @@ def _construct_url( def _requests_helper(url, timeout): - """Make a GET request to the given URL""" + """Make a GET request to the given URL + + ``timeout`` must be passed by keyword. ``requests.get`` takes ``params`` as + its second positional argument, so passing it positionally would append the + value to the query string and leave the request with no timeout at all. + """ response = None try: - response = requests.get(url, timeout) + response = requests.get(url, timeout=timeout) except requests.exceptions.RequestException as e: - logger.exception("Request exception: %s", str(e)) + logger.debug("Request exception: %s", str(e)) return response @@ -511,9 +569,24 @@ def _get_region_or_default(session): def _get_default_sagemaker_session(): - """Return the default sagemaker session""" + """Return the default sagemaker session + + The region is resolved by boto3 from the caller's own environment + (``AWS_REGION``, ``AWS_DEFAULT_REGION``, or the active profile in + ``~/.aws/config``). ``DEFAULT_AWS_REGION`` is only used as a last resort, + because ``Session`` requires a region. Hardcoding the default meant that + callers with no session of their own (module-level functions such as + ``ingest_dataframe``) had their telemetry pointed at a region they may have + no network route to. + """ - boto_session = boto3.Session(region_name=DEFAULT_AWS_REGION) + boto_session = boto3.Session() + if not boto_session.region_name: + logger.debug( + "No region resolved from the local AWS configuration. Falling back to %s.", + DEFAULT_AWS_REGION, + ) + boto_session = boto3.Session(region_name=DEFAULT_AWS_REGION) sagemaker_session = Session(boto_session=boto_session) return sagemaker_session diff --git a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py index 6c7359ceb1..e1e55d4241 100644 --- a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py +++ b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py @@ -12,16 +12,19 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import import os +import threading import unittest +from time import perf_counter import pytest import requests from unittest.mock import Mock, patch, MagicMock import boto3 import sagemaker -from sagemaker.core.telemetry.constants import Feature +from sagemaker.core.telemetry.constants import Feature, DEFAULT_AWS_REGION from sagemaker.core.telemetry.attribution import _CREATED_BY_ENV_VAR from sagemaker.core.telemetry.telemetry_logging import ( _send_telemetry_request, + _send_telemetry_request_sync, _telemetry_emitter, _construct_url, _get_accountId, @@ -30,6 +33,7 @@ _get_default_sagemaker_session, OS_NAME_VERSION, PYTHON_VERSION, + TELEMETRY_REQUEST_TIMEOUT, ) from sagemaker.core.user_agent import SDK_VERSION, process_studio_metadata_file @@ -76,18 +80,18 @@ def test_log_sucessfully(self, mock_get_accountId, mock_request_helper): """Test to check if the telemetry logging is successful""" MOCK_SESSION.boto_session.region_name = "us-west-2" mock_get_accountId.return_value = "testAccountId" - _send_telemetry_request("someStatus", "1", MOCK_SESSION) + _send_telemetry_request_sync("someStatus", "1", MOCK_SESSION) mock_request_helper.assert_called_with( "https://sm-pysdk-t-us-west-2.s3.us-west-2.amazonaws.com/" "telemetry?x-accountId=testAccountId&x-status=someStatus&x-feature=1", - 2, + TELEMETRY_REQUEST_TIMEOUT, ) @patch("sagemaker.core.telemetry.telemetry_logging._get_accountId") def test_log_handle_exception(self, mock_get_accountId): """Test to check if the exception is handled while logging telemetry""" mock_get_accountId.side_effect = Exception("Internal error") - _send_telemetry_request("someStatus", "1", MOCK_SESSION) + _send_telemetry_request_sync("someStatus", "1", MOCK_SESSION) self.assertRaises(Exception) @patch("sagemaker.core.telemetry.telemetry_logging._get_accountId") @@ -101,11 +105,11 @@ def test_send_telemetry_request_success(self, mock_get_region, mock_get_accountI "sagemaker.core.telemetry.telemetry_logging._requests_helper" ) as mock_requests_helper: mock_requests_helper.return_value = None - _send_telemetry_request(1, [1, 2], MagicMock(), None, None, "extra_info") + _send_telemetry_request_sync(1, [1, 2], MagicMock(), None, None, "extra_info") mock_requests_helper.assert_called_with( "https://sm-pysdk-t-us-west-2.s3.us-west-2.amazonaws.com/" "telemetry?x-accountId=testAccountId&x-status=1&x-feature=1,2&x-extra=extra_info", - 2, + TELEMETRY_REQUEST_TIMEOUT, ) @patch("sagemaker.core.telemetry.telemetry_logging._get_accountId") @@ -119,14 +123,14 @@ def test_send_telemetry_request_failure(self, mock_get_region, mock_get_accountI "sagemaker.core.telemetry.telemetry_logging._requests_helper" ) as mock_requests_helper: mock_requests_helper.return_value = None - _send_telemetry_request( + _send_telemetry_request_sync( 0, [1, 2], MagicMock(), "failure_reason", "failure_type", "extra_info" ) mock_requests_helper.assert_called_with( "https://sm-pysdk-t-us-west-2.s3.us-west-2.amazonaws.com/" "telemetry?x-accountId=testAccountId&x-status=0&x-feature=1,2" "&x-failureReason=failure_reason&x-failureType=failure_type&x-extra=extra_info", - 2, + TELEMETRY_REQUEST_TIMEOUT, ) @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") @@ -249,7 +253,9 @@ def test_requests_helper_success(self, mock_requests_get): response = _requests_helper(url, timeout) - mock_requests_get.assert_called_once_with(url, timeout) + # timeout must be a keyword argument: positionally it becomes `params`, + # which leaves the request with no timeout at all. + mock_requests_get.assert_called_once_with(url, timeout=timeout) self.assertEqual(response, mock_response) @patch("sagemaker.core.telemetry.telemetry_logging.requests.get") @@ -261,7 +267,7 @@ def test_requests_helper_exception(self, mock_requests_get): response = _requests_helper(url, timeout) - mock_requests_get.assert_called_once_with(url, timeout) + mock_requests_get.assert_called_once_with(url, timeout=timeout) self.assertIsNone(response) def test_get_accountId_success(self): @@ -340,12 +346,12 @@ def test_send_telemetry_request_valid_region(self, mock_get_region, mock_get_acc with patch( "sagemaker.core.telemetry.telemetry_logging._requests_helper" ) as mock_requests_helper: - _send_telemetry_request(1, [1, 2], mock_session) + _send_telemetry_request_sync(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, + TELEMETRY_REQUEST_TIMEOUT, ) @patch("sagemaker.core.telemetry.telemetry_logging._get_accountId") @@ -360,7 +366,7 @@ def test_send_telemetry_request_invalid_region(self, mock_get_region, mock_get_a with patch( "sagemaker.core.telemetry.telemetry_logging._requests_helper" ) as mock_requests_helper: - _send_telemetry_request(1, [1, 2], mock_session) + _send_telemetry_request_sync(1, [1, 2], mock_session) # Assert telemetry request was not sent mock_requests_helper.assert_not_called() @@ -701,3 +707,181 @@ def test_telemetry_opt_out_message_not_shown_when_opted_out( # Reset the flag for other tests telemetry_module._telemetry_msg_shown = False + + +class TestRequestsHelperTimeout(unittest.TestCase): + """The telemetry GET must actually carry a timeout. + + `requests.get(url, params=None, **kwargs)` takes `params` second, so passing + the timeout positionally appended it to the query string and left the + request with no timeout, letting an unreachable telemetry endpoint block the + caller indefinitely. + """ + + @patch("sagemaker.core.telemetry.telemetry_logging.requests.get") + def test_timeout_passed_as_keyword_not_params(self, mock_requests_get): + _requests_helper("https://example.com/telemetry?x-status=1", 2) + + _, kwargs = mock_requests_get.call_args + self.assertEqual(kwargs["timeout"], 2) + self.assertNotIn("params", kwargs) + + def test_timeout_reaches_prepared_request_not_the_url(self): + """Guard against the regression at the layer where it was observable.""" + captured = {} + + def fake_get(url, **kwargs): + captured["url"] = url + captured["kwargs"] = kwargs + return None + + target = "sagemaker.core.telemetry.telemetry_logging.requests.get" + with patch(target, side_effect=fake_get): + _requests_helper("https://example.com/telemetry?x-status=1", 2) + + # The old code produced a URL ending in "&2" and no timeout kwarg. + self.assertFalse(captured["url"].endswith("&2")) + self.assertEqual(captured["kwargs"], {"timeout": 2}) + + +class TestTelemetryIsNonBlocking(unittest.TestCase): + """Telemetry must never add latency to the SDK call that triggered it. + + Regression guard: a Feature Store ingest returned in under a second + server-side but the notebook cell took ~47 minutes, because each telemetry + emission blocked on an endpoint the caller's VPC had no route to. + """ + + def setUp(self): + import sagemaker.core.telemetry.telemetry_logging as telemetry_module + + self.telemetry_module = telemetry_module + + def test_send_returns_before_the_request_completes(self): + release = threading.Event() + entered = threading.Event() + + def blocking_send(*args, **kwargs): + entered.set() + release.wait(timeout=10) + + with patch.object( + self.telemetry_module, "_send_telemetry_request_sync", side_effect=blocking_send + ): + start = perf_counter() + thread = _send_telemetry_request(1, [1], MagicMock()) + elapsed = perf_counter() - start + + try: + self.assertLess(elapsed, 1, "_send_telemetry_request blocked on the network call") + self.assertTrue(entered.wait(timeout=5)) + finally: + release.set() + thread.join(timeout=5) + + def test_send_runs_on_a_daemon_thread(self): + """Daemon threads are killed at exit, so a pending send cannot hang shutdown.""" + with patch.object(self.telemetry_module, "_send_telemetry_request_sync"): + thread = _send_telemetry_request(1, [1], MagicMock()) + self.assertTrue(thread.daemon) + thread.join(timeout=5) + + def test_send_forwards_all_arguments(self): + session = MagicMock() + with patch.object(self.telemetry_module, "_send_telemetry_request_sync") as mock_sync: + thread = _send_telemetry_request( + 0, [1, 2], session, "failure_reason", "failure_type", "extra_info" + ) + thread.join(timeout=5) + + mock_sync.assert_called_once_with( + 0, [1, 2], session, "failure_reason", "failure_type", "extra_info" + ) + + def test_thread_swallows_exceptions(self): + """An exception in the thread has no caller to catch it, so it must not escape.""" + with patch.object( + self.telemetry_module, + "_send_telemetry_request_sync", + side_effect=RuntimeError("boom"), + ): + thread = _send_telemetry_request(1, [1], MagicMock()) + thread.join(timeout=5) + + self.assertFalse(thread.is_alive()) + + def test_no_event_is_dropped_when_many_are_in_flight(self): + """Making the send async must not cost us events, however many are pending.""" + release = threading.Event() + started = threading.Semaphore(0) + event_count = 25 + + def blocking_send(*args, **kwargs): + started.release() + release.wait(timeout=10) + + with patch.object( + self.telemetry_module, "_send_telemetry_request_sync", side_effect=blocking_send + ): + threads = [_send_telemetry_request(1, [1], MagicMock()) for _ in range(event_count)] + try: + self.assertTrue(all(t is not None for t in threads)) + for _ in range(event_count): + self.assertTrue(started.acquire(timeout=5), "an event was never sent") + finally: + release.set() + for t in threads: + t.join(timeout=5) + + @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config") + def test_decorated_function_returns_without_waiting_for_telemetry(self, mock_resolve_config): + mock_resolve_config.return_value = False + release = threading.Event() + + def blocking_send(*args, **kwargs): + release.wait(timeout=10) + + with patch.object( + self.telemetry_module, "_send_telemetry_request_sync", side_effect=blocking_send + ): + try: + start = perf_counter() + LocalSagemakerClientMock().mock_create_model() + elapsed = perf_counter() - start + self.assertLess(elapsed, 1, "the decorated call waited on the telemetry request") + finally: + release.set() + + +class TestDefaultSessionRegion(unittest.TestCase): + """The synthesized fallback session must use the caller's own region. + + Module-level functions such as `ingest_dataframe` have no session, so the + decorator builds one. Hardcoding us-west-2 pointed telemetry at a region the + caller may have no network route to. + """ + + @patch("sagemaker.core.telemetry.telemetry_logging.Session") + @patch("sagemaker.core.telemetry.telemetry_logging.boto3.Session") + def test_uses_region_resolved_from_environment(self, mock_boto_session, mock_session): + mock_boto_session.return_value.region_name = "ca-central-1" + + _get_default_sagemaker_session() + + # Called with no region_name so boto3 resolves it from the environment + # or the active profile, rather than being pinned to us-west-2. + mock_boto_session.assert_called_once_with() + mock_session.assert_called_once_with(boto_session=mock_boto_session.return_value) + + @patch("sagemaker.core.telemetry.telemetry_logging.Session") + @patch("sagemaker.core.telemetry.telemetry_logging.boto3.Session") + def test_falls_back_to_default_region_when_none_resolved(self, mock_boto_session, mock_session): + mock_boto_session.return_value.region_name = None + + _get_default_sagemaker_session() + + # Session requires a region, so the default is still the last resort. + self.assertEqual( + mock_boto_session.call_args_list[-1], + unittest.mock.call(region_name=DEFAULT_AWS_REGION), + ) diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/MIGRATION_GUIDE.md b/sagemaker-mlops/src/sagemaker/mlops/feature_store/MIGRATION_GUIDE.md index 40942fa6f3..414e42e65f 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/MIGRATION_GUIDE.md +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/MIGRATION_GUIDE.md @@ -236,6 +236,19 @@ manager = ingest_dataframe( # Access failed rows: manager.failed_rows ``` +In V2 the region came from the `Session` you passed to the `FeatureGroup`. In V3 +`ingest_dataframe` takes an explicit `region`; when you omit it, boto3 resolves the +region from the environment (`AWS_DEFAULT_REGION`, `AWS_REGION`, or the active +profile in `~/.aws/config`): + +```python +manager = ingest_dataframe( + feature_group_name="my-fg", + data_frame=df, + region="eu-west-1", +) +``` + --- ## Athena Query diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py index 0f80a509d8..8c6d9b2615 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_utils.py @@ -479,6 +479,7 @@ def ingest_dataframe( wait: bool = True, timeout: Union[int, float] = None, use_batch_write_record: bool = False, + region: str = None, ): """Ingest a pandas DataFrame to a FeatureGroup. @@ -493,20 +494,29 @@ def ingest_dataframe( call) instead of PutRecord (1 record per call) for significantly better throughput. Requires both ``sagemaker:BatchWriteRecord`` AND ``sagemaker:PutRecord`` IAM permissions. Default: False. + region: AWS region name of the FeatureGroup, e.g. "eu-west-1". Used both to + describe the FeatureGroup and to write the records. If not specified, the + region is resolved by boto3 from the environment (``AWS_DEFAULT_REGION``, + ``AWS_REGION``, or the active profile in ``~/.aws/config``). Default: None. Returns: IngestionManagerPandas instance. Raises: ValueError: If max_workers or max_processes <= 0. + + Note: + sagemaker-core caches its boto clients per process, so the first region used in + a process wins. Use a single region per process, or the same ``region`` value on + every call. """ - + if max_processes <= 0: raise ValueError("max_processes must be greater than 0.") if max_workers <= 0: raise ValueError("max_workers must be greater than 0.") - fg = CoreFeatureGroup.get(feature_group_name=feature_group_name) + fg = CoreFeatureGroup.get(feature_group_name=feature_group_name, region=region) feature_definitions = {} for fd in fg.feature_definitions: collection_type = getattr(fd, "collection_type", None) @@ -524,6 +534,7 @@ def ingest_dataframe( max_workers=max_workers, max_processes=max_processes, use_batch_write_record=use_batch_write_record, + region=region, ) manager.run(data_frame=data_frame, wait=wait, timeout=timeout) return manager diff --git a/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py b/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py index 0762193e28..df49a63d5e 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py +++ b/sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py @@ -52,6 +52,12 @@ class IngestionManagerPandas: max_workers (int): number of threads to create. max_processes (int): number of processes to create. Each process spawns ``max_workers`` threads. + use_batch_write_record (bool): whether to use the BatchWriteRecord API + instead of PutRecord. + region (str): AWS region name to write the records to, e.g. "eu-west-1". + If None, the region is resolved by boto3 from the environment + (``AWS_DEFAULT_REGION``, ``AWS_REGION``, or the active profile in + ``~/.aws/config``). """ feature_group_name: str @@ -59,6 +65,7 @@ class IngestionManagerPandas: max_workers: int = 1 max_processes: int = 1 use_batch_write_record: bool = False + region: str = None _async_result: Any = field(default=None, init=False) _processing_pool: Pool = field(default=None, init=False) _failed_indices: List[int] = field(default_factory=list, init=False) @@ -151,6 +158,7 @@ def _run_single_process_single_thread( start_index=0, end_index=len(data_frame), target_stores=target_stores, + region=self.region, ) else: failed_rows = [] @@ -163,6 +171,7 @@ def _run_single_process_single_thread( feature_definitions=self.feature_definitions, failed_rows=failed_rows, target_stores=target_stores, + region=self.region, ) self._failed_indices = failed_rows @@ -195,6 +204,7 @@ def _run_multi_process( start_index, timeout, self.use_batch_write_record, + self.region, )) def init_worker(): @@ -220,6 +230,7 @@ def _run_multi_threaded( row_offset: int = 0, timeout: Union[int, float] = None, use_batch_write_record: bool = False, + region: str = None, ) -> List[int]: """Start multi-threaded ingestion within a single process.""" executor = ThreadPoolExecutor(max_workers=max_workers) @@ -238,6 +249,7 @@ def _run_multi_threaded( start_index=start_index, end_index=end_index, target_stores=target_stores, + region=region, ) else: future = executor.submit( @@ -248,6 +260,7 @@ def _run_multi_threaded( start_index=start_index, end_index=end_index, target_stores=target_stores, + region=region, ) futures[future] = (start_index + row_offset, end_index + row_offset) @@ -270,6 +283,7 @@ def _ingest_single_batch( start_index: int, end_index: int, target_stores: List[str] = None, + region: str = None, ) -> List[int]: """Ingest a single batch of DataFrame rows into FeatureStore.""" logger.info("Started ingesting index %d to %d", start_index, end_index) @@ -285,6 +299,7 @@ def _ingest_single_batch( feature_definitions=feature_definitions, failed_rows=failed_rows, target_stores=target_stores, + region=region, ) return failed_rows @@ -297,6 +312,7 @@ def _ingest_row( feature_definitions: Dict[str, Dict[Any, Any]], failed_rows: List[int], target_stores: List[str] = None, + region: str = None, ): """Ingest a single DataFrame row into FeatureStore using SageMaker Core.""" try: @@ -323,6 +339,7 @@ def _ingest_row( feature_group.put_record( record=record, target_stores=target_stores, + region=region, ) except Exception as e: @@ -408,6 +425,7 @@ def _ingest_batch_write( start_index: int, end_index: int, target_stores: List[str] = None, + region: str = None, ) -> List[int]: """Ingest records using BatchWriteRecord API (up to 25 per call). @@ -418,6 +436,8 @@ def _ingest_batch_write( start_index: Start index in the DataFrame slice. end_index: End index in the DataFrame slice. target_stores: Target stores for ingestion. + region: AWS region name to write the records to. If None, boto3 resolves + the region from the environment. Returns: List of row indices that failed to ingest. @@ -459,7 +479,7 @@ def _ingest_batch_write( try: fg = CoreFeatureGroup(feature_group_name=feature_group_name) - response = fg.batch_write_record(entries=entries) + response = fg.batch_write_record(entries=entries, region=region) # Handle partial failures from unprocessed entries if response.unprocessed_entries: diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py index 7fd55ceef6..93311b2284 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py @@ -251,6 +251,78 @@ def test_raises_on_invalid_max_processes(self): ingest_dataframe("my-fg", df, max_processes=-1) +class TestIngestDataframeRegion: + """``region`` must reach both the describe call and the ingestion manager.""" + + @pytest.fixture + def mock_feature_group(self): + mock_fg = MagicMock() + mock_fg.feature_definitions = [ + MagicMock(feature_name="id", feature_type="Integral"), + ] + return mock_fg + + @patch("sagemaker.mlops.feature_store.feature_utils.IngestionManagerPandas") + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_region_passed_to_describe_and_manager( + self, mock_fg_class, mock_manager_class, mock_feature_group + ): + mock_fg_class.get.return_value = mock_feature_group + + df = pd.DataFrame({"id": [1, 2, 3]}) + ingest_dataframe("my-fg", df, region="eu-west-1") + + mock_fg_class.get.assert_called_once_with( + feature_group_name="my-fg", region="eu-west-1" + ) + assert mock_manager_class.call_args[1]["region"] == "eu-west-1" + + @patch("sagemaker.mlops.feature_store.feature_utils.IngestionManagerPandas") + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_region_defaults_to_none( + self, mock_fg_class, mock_manager_class, mock_feature_group + ): + mock_fg_class.get.return_value = mock_feature_group + + df = pd.DataFrame({"id": [1, 2, 3]}) + ingest_dataframe("my-fg", df) + + mock_fg_class.get.assert_called_once_with(feature_group_name="my-fg", region=None) + assert mock_manager_class.call_args[1]["region"] is None + + @patch("sagemaker.mlops.feature_store.feature_utils.IngestionManagerPandas") + @patch("sagemaker.mlops.feature_store.feature_utils.CoreFeatureGroup") + def test_region_works_with_batch_write_record( + self, mock_fg_class, mock_manager_class, mock_feature_group + ): + mock_fg_class.get.return_value = mock_feature_group + + df = pd.DataFrame({"id": [1, 2, 3]}) + ingest_dataframe("my-fg", df, use_batch_write_record=True, region="ap-south-1") + + kwargs = mock_manager_class.call_args[1] + assert kwargs["region"] == "ap-south-1" + assert kwargs["use_batch_write_record"] is True + + def test_region_is_keyword_only_in_practice_and_does_not_shift_positionals(self): + """``region`` is appended last, so existing positional calls keep working.""" + import inspect + + from sagemaker.mlops.feature_store.feature_utils import ingest_dataframe as fn + + params = list(inspect.signature(fn).parameters) + assert params[-1] == "region" + assert params[:7] == [ + "feature_group_name", + "data_frame", + "max_workers", + "max_processes", + "wait", + "timeout", + "use_batch_write_record", + ] + + class TestGetSessionFromRole: @patch("sagemaker.mlops.feature_store.feature_utils.boto3") @patch("sagemaker.mlops.feature_store.feature_utils.Session") diff --git a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py index 46cdef158e..5e9c985ddf 100644 --- a/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py +++ b/sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py @@ -321,3 +321,182 @@ def test_async_with_parallelism_no_validation_error(self, mock_run, max_workers, # Verify it called the multi-process method (positive assertion) mock_run.assert_called_once() + + +class TestIngestionManagerRegion: + """``region`` must reach every FeatureStore runtime call.""" + + @pytest.fixture + def feature_definitions(self): + return {"id": {"FeatureType": "String", "CollectionType": None}} + + @pytest.fixture + def sample_dataframe(self): + return pd.DataFrame({"id": ["1", "2", "3"]}) + + def test_region_defaults_to_none(self, feature_definitions): + manager = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + ) + assert manager.region is None + + def test_region_stored_on_manager(self, feature_definitions): + manager = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + region="eu-west-1", + ) + assert manager.region == "eu-west-1" + + def test_ingest_row_passes_region_to_put_record(self, feature_definitions): + df = pd.DataFrame({"id": ["1"]}) + mock_fg = MagicMock() + failed_rows = [] + + for row in df.itertuples(): + IngestionManagerPandas._ingest_row( + data_frame=df, + row=row, + feature_group=mock_fg, + feature_definitions=feature_definitions, + failed_rows=failed_rows, + target_stores=None, + region="eu-west-1", + ) + + assert mock_fg.put_record.call_args[1]["region"] == "eu-west-1" + + def test_ingest_row_region_defaults_to_none(self, feature_definitions): + df = pd.DataFrame({"id": ["1"]}) + mock_fg = MagicMock() + failed_rows = [] + + for row in df.itertuples(): + IngestionManagerPandas._ingest_row( + data_frame=df, + row=row, + feature_group=mock_fg, + feature_definitions=feature_definitions, + failed_rows=failed_rows, + target_stores=None, + ) + + assert mock_fg.put_record.call_args[1]["region"] is None + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_single_thread_run_passes_region_to_put_record( + self, mock_fg_class, feature_definitions, sample_dataframe + ): + mock_fg = MagicMock() + mock_fg_class.return_value = mock_fg + + manager = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + region="eu-west-1", + ) + manager.run(data_frame=sample_dataframe, wait=True) + + assert mock_fg.put_record.call_count == 3 + for call in mock_fg.put_record.call_args_list: + assert call[1]["region"] == "eu-west-1" + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_single_batch_passes_region_to_put_record( + self, mock_fg_class, feature_definitions, sample_dataframe + ): + mock_fg = MagicMock() + mock_fg_class.return_value = mock_fg + + IngestionManagerPandas._ingest_single_batch( + data_frame=sample_dataframe, + feature_group_name="test-fg", + feature_definitions=feature_definitions, + start_index=0, + end_index=3, + region="eu-west-1", + ) + + assert mock_fg.put_record.call_count == 3 + for call in mock_fg.put_record.call_args_list: + assert call[1]["region"] == "eu-west-1" + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_passes_region( + self, mock_fg_class, feature_definitions, sample_dataframe + ): + mock_fg = MagicMock() + mock_fg.batch_write_record.return_value = MagicMock( + unprocessed_entries=[], errors=[] + ) + mock_fg_class.return_value = mock_fg + + IngestionManagerPandas._ingest_batch_write( + data_frame=sample_dataframe, + feature_group_name="test-fg", + feature_definitions=feature_definitions, + start_index=0, + end_index=3, + region="eu-west-1", + ) + + mock_fg.batch_write_record.assert_called_once() + assert mock_fg.batch_write_record.call_args[1]["region"] == "eu-west-1" + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") + def test_batch_write_run_passes_region( + self, mock_fg_class, feature_definitions, sample_dataframe + ): + mock_fg = MagicMock() + mock_fg.batch_write_record.return_value = MagicMock( + unprocessed_entries=[], errors=[] + ) + mock_fg_class.return_value = mock_fg + + manager = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + use_batch_write_record=True, + region="eu-west-1", + ) + manager.run(data_frame=sample_dataframe, wait=True) + + assert mock_fg.batch_write_record.call_args[1]["region"] == "eu-west-1" + + @patch.object(IngestionManagerPandas, "_ingest_single_batch", return_value=[]) + def test_multi_threaded_forwards_region( + self, mock_ingest, feature_definitions, sample_dataframe + ): + IngestionManagerPandas._run_multi_threaded( + max_workers=2, + feature_group_name="test-fg", + feature_definitions=feature_definitions, + data_frame=sample_dataframe, + region="eu-west-1", + ) + + assert mock_ingest.call_count == 2 + for call in mock_ingest.call_args_list: + assert call[1]["region"] == "eu-west-1" + + @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.Pool") + def test_multi_process_args_include_region( + self, mock_pool_class, feature_definitions, sample_dataframe + ): + mock_pool = MagicMock() + mock_pool_class.return_value = mock_pool + + manager = IngestionManagerPandas( + feature_group_name="test-fg", + feature_definitions=feature_definitions, + max_processes=2, + region="eu-west-1", + ) + manager.run(data_frame=sample_dataframe, wait=False) + + starmap_args = mock_pool.starmap_async.call_args[0][1] + assert len(starmap_args) == 2 + for process_args in starmap_args: + # region is the last positional argument passed to _run_multi_threaded + assert process_args[-1] == "eu-west-1" diff --git a/sagemaker-serve/src/sagemaker/serve/utils/telemetry_logger.py b/sagemaker-serve/src/sagemaker/serve/utils/telemetry_logger.py index 0ce68b153a..c2e7aee5c5 100644 --- a/sagemaker-serve/src/sagemaker/serve/utils/telemetry_logger.py +++ b/sagemaker-serve/src/sagemaker/serve/utils/telemetry_logger.py @@ -244,11 +244,16 @@ def _construct_url( def _requests_helper(url, timeout): - """Placeholder docstring""" + """Make a GET request to the given URL + + ``timeout`` must be passed by keyword. ``requests.get`` takes ``params`` as + its second positional argument, so passing it positionally would append the + value to the query string and leave the request with no timeout at all. + """ response = None try: - response = requests.get(url, timeout) + response = requests.get(url, timeout=timeout) except requests.exceptions.RequestException as e: logger.debug("Request exception: %s", str(e)) return response diff --git a/sagemaker-serve/tests/unit/test_telemetry_logger.py b/sagemaker-serve/tests/unit/test_telemetry_logger.py index ba7d487a8c..9059679448 100644 --- a/sagemaker-serve/tests/unit/test_telemetry_logger.py +++ b/sagemaker-serve/tests/unit/test_telemetry_logger.py @@ -135,7 +135,7 @@ def test_requests_helper_success(self, mock_get): result = _requests_helper("https://example.com", 2) self.assertEqual(result, mock_response) - mock_get.assert_called_once_with("https://example.com", 2) + mock_get.assert_called_once_with("https://example.com", timeout=2) @patch('sagemaker.serve.utils.telemetry_logger.requests.get') def test_requests_helper_exception(self, mock_get):