From b3e1695a344aaa3205db88383ecfe941cd066bd6 Mon Sep 17 00:00:00 2001 From: Gokul A Date: Thu, 20 Aug 2026 19:21:18 -0700 Subject: [PATCH 1/3] feat(mlops): allow specifying region in feature_store ingest_dataframe ingest_dataframe() had no way to say which AWS region the FeatureGroup lives in, so both the DescribeFeatureGroup call and the record writes fell back to whatever region boto3 resolved, with no caller control. Add an optional `region` argument and thread it through the whole ingestion path: - feature_utils.ingest_dataframe -> CoreFeatureGroup.get(region=...) and IngestionManagerPandas(region=...) - IngestionManagerPandas gains a `region` field, forwarded to put_record(region=...) and batch_write_record(region=...) in the single-thread, multi-thread, multi-process, and BatchWriteRecord paths `region` is appended last and defaults to None, so existing calls and their behavior are unchanged. --- .../mlops/feature_store/MIGRATION_GUIDE.md | 13 ++ .../mlops/feature_store/feature_utils.py | 15 +- .../feature_store/ingestion_manager_pandas.py | 22 ++- .../mlops/feature_store/test_feature_utils.py | 72 +++++++ .../test_ingestion_manager_pandas.py | 179 ++++++++++++++++++ 5 files changed, 298 insertions(+), 3 deletions(-) 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" From 131a45e0a385abac33d6824c1cdfcd01f8cd5f4f Mon Sep 17 00:00:00 2001 From: Gokul A Date: Fri, 21 Aug 2026 08:08:10 -0700 Subject: [PATCH 2/3] fix(core): stop telemetry from blocking SDK calls Telemetry could add unbounded latency to any decorated SDK call. A Feature Store ingest was reported taking ~47 minutes from a private VPC while the underlying PutRecord completed server-side in 348ms; all of the remaining time was spent in the two telemetry emissions that follow the call. Three defects, all in the emission path: 1. `_requests_helper` passed the timeout positionally. `requests.get` takes `params` as its second positional argument, so the value was appended to the query string and the request had no timeout at all. From a VPC with no route to the telemetry endpoint the GET hung until a network device dropped the flow. Now passed as `timeout=`, verified bounded at ~2s against a black-holed address. The same one-line defect existed in sagemaker-serve's telemetry_logger and is fixed there too. 2. `_get_default_sagemaker_session` hardcoded us-west-2. Module-level functions such as `ingest_dataframe` have no session of their own, so the decorator synthesizes one, which pointed both the STS `get_caller_identity` call and the telemetry GET at a region the caller may have no route to. The region is now resolved by boto3 from the caller's environment, with the default kept only as a last resort since `Session` requires a region. 3. Emission was synchronous. `_send_telemetry_request` now dispatches to a daemon thread and returns immediately, so neither the STS call nor the GET can sit in the caller's critical path. Daemon threads are killed at interpreter exit, so a pending send cannot delay shutdown either. In-flight sends are capped and excess events dropped rather than growing threads without bound, and nothing can escape the thread. The existing body moved to `_send_telemetry_request_sync`. Telemetry request failures now log at debug instead of logging a full traceback at error level; a best-effort metric should not look like an error. Note: events queued at process exit may now be lost. That is the intended trade-off for never blocking the caller. --- .../core/telemetry/telemetry_logging.py | 93 +++++++- .../unit/telemetry/test_telemetry_logging.py | 218 ++++++++++++++++-- .../sagemaker/serve/utils/telemetry_logger.py | 9 +- .../tests/unit/test_telemetry_logger.py | 2 +- 4 files changed, 300 insertions(+), 22 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index b8f5205e91..7e40a9da56 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,20 @@ ) _telemetry_msg_shown = False +# Seconds to wait for the telemetry endpoint before giving up on an event. +TELEMETRY_REQUEST_TIMEOUT = 2 + +# Telemetry is best-effort, so it must never sit in the caller's critical path. +# 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. If more than +# MAX_IN_FLIGHT_TELEMETRY_REQUESTS sends are already outstanding we drop the +# event rather than grow threads without bound. +MAX_IN_FLIGHT_TELEMETRY_REQUESTS = 8 +_in_flight_telemetry_requests = 0 +_in_flight_telemetry_lock = threading.Lock() + FEATURE_TO_CODE = { str(Feature.SDK_DEFAULTS): 11, str(Feature.LOCAL_MODE): 12, @@ -422,6 +437,52 @@ 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. + + This returns as soon as the thread is started so that telemetry never adds + latency to the SDK call that triggered it. The event is dropped (and None + returned) if too many sends are already outstanding. + + Returns: + threading.Thread: The thread doing the send, or None if the event was + dropped. Callers generally ignore this; tests can join on it. + """ + global _in_flight_telemetry_requests # pylint: disable=W0603 + with _in_flight_telemetry_lock: + if _in_flight_telemetry_requests >= MAX_IN_FLIGHT_TELEMETRY_REQUESTS: + logger.debug("Too many telemetry requests in flight. Dropping this one.") + return None + _in_flight_telemetry_requests += 1 + + def _run(): + global _in_flight_telemetry_requests # pylint: disable=W0603 + 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 + finally: + with _in_flight_telemetry_lock: + _in_flight_telemetry_requests -= 1 + + 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 +510,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 +543,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 +577,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..791617337d 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,189 @@ 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 + telemetry_module._in_flight_telemetry_requests = 0 + + def tearDown(self): + self.telemetry_module._in_flight_telemetry_requests = 0 + + 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_and_releases_its_slot(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()) + self.assertEqual(self.telemetry_module._in_flight_telemetry_requests, 0) + + def test_events_are_dropped_when_too_many_are_in_flight(self): + release = threading.Event() + started = [] + + def blocking_send(*args, **kwargs): + started.append(1) + release.wait(timeout=10) + + max_in_flight = self.telemetry_module.MAX_IN_FLIGHT_TELEMETRY_REQUESTS + with patch.object( + self.telemetry_module, "_send_telemetry_request_sync", side_effect=blocking_send + ): + threads = [_send_telemetry_request(1, [1], MagicMock()) for _ in range(max_in_flight)] + try: + self.assertTrue(all(t is not None for t in threads)) + # One more than the cap allows is dropped rather than spawning + # an unbounded number of threads. + self.assertIsNone(_send_telemetry_request(1, [1], MagicMock())) + finally: + release.set() + for t in threads: + t.join(timeout=5) + + # Counter is released once the sends finish, so later events go through. + self.assertEqual(self.telemetry_module._in_flight_telemetry_requests, 0) + + @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-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): From d1d06c41c6e64ae0eccf19a0133d362a17f13010 Mon Sep 17 00:00:00 2001 From: Gokul A Date: Fri, 21 Aug 2026 08:44:02 -0700 Subject: [PATCH 3/3] fix(core): never drop telemetry events when sending asynchronously Removes the in-flight cap added alongside the async send. Telemetry events are data we cannot silently lose, so every event now gets its own daemon thread rather than being dropped once eight sends are outstanding. Feature Store ingestion is the special case that motivated moving the send off the caller's thread at all: it is decorated at more than one level (ingest_dataframe and IngestionManagerPandas.run), so a single user call emits several events, and sending them serially turned an ingest the service finished in under a second into a multi-minute wait. --- .../core/telemetry/telemetry_logging.py | 42 ++++++++----------- .../unit/telemetry/test_telemetry_logging.py | 26 ++++-------- 2 files changed, 26 insertions(+), 42 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index 7e40a9da56..d55a390056 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -70,16 +70,19 @@ # Seconds to wait for the telemetry endpoint before giving up on an event. TELEMETRY_REQUEST_TIMEOUT = 2 -# Telemetry is best-effort, so it must never sit in the caller's critical path. -# 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. If more than -# MAX_IN_FLIGHT_TELEMETRY_REQUESTS sends are already outstanding we drop the -# event rather than grow threads without bound. -MAX_IN_FLIGHT_TELEMETRY_REQUESTS = 8 -_in_flight_telemetry_requests = 0 -_in_flight_telemetry_lock = threading.Lock() +# 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, @@ -440,23 +443,15 @@ def _send_telemetry_request( ) -> threading.Thread: """Schedule a telemetry event to be sent on a background daemon thread. - This returns as soon as the thread is started so that telemetry never adds - latency to the SDK call that triggered it. The event is dropped (and None - returned) if too many sends are already outstanding. + 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, or None if the event was - dropped. Callers generally ignore this; tests can join on it. + threading.Thread: The thread doing the send. Callers generally ignore this; + tests can join on it. """ - global _in_flight_telemetry_requests # pylint: disable=W0603 - with _in_flight_telemetry_lock: - if _in_flight_telemetry_requests >= MAX_IN_FLIGHT_TELEMETRY_REQUESTS: - logger.debug("Too many telemetry requests in flight. Dropping this one.") - return None - _in_flight_telemetry_requests += 1 def _run(): - global _in_flight_telemetry_requests # pylint: disable=W0603 try: _send_telemetry_request_sync( status, feature_list, session, failure_reason, failure_type, extra_info @@ -467,9 +462,6 @@ def _run(): # _send_telemetry_request_sync can fail once the interpreter starts # tearing down. Telemetry is best-effort, so drop the event silently. pass - finally: - with _in_flight_telemetry_lock: - _in_flight_telemetry_requests -= 1 thread = threading.Thread(target=_run, name="sagemaker-telemetry", daemon=True) thread.start() diff --git a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py index 791617337d..e1e55d4241 100644 --- a/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py +++ b/sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py @@ -756,10 +756,6 @@ def setUp(self): import sagemaker.core.telemetry.telemetry_logging as telemetry_module self.telemetry_module = telemetry_module - telemetry_module._in_flight_telemetry_requests = 0 - - def tearDown(self): - self.telemetry_module._in_flight_telemetry_requests = 0 def test_send_returns_before_the_request_completes(self): release = threading.Event() @@ -802,7 +798,7 @@ def test_send_forwards_all_arguments(self): 0, [1, 2], session, "failure_reason", "failure_type", "extra_info" ) - def test_thread_swallows_exceptions_and_releases_its_slot(self): + 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, @@ -813,34 +809,30 @@ def test_thread_swallows_exceptions_and_releases_its_slot(self): thread.join(timeout=5) self.assertFalse(thread.is_alive()) - self.assertEqual(self.telemetry_module._in_flight_telemetry_requests, 0) - def test_events_are_dropped_when_too_many_are_in_flight(self): + 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 = [] + started = threading.Semaphore(0) + event_count = 25 def blocking_send(*args, **kwargs): - started.append(1) + started.release() release.wait(timeout=10) - max_in_flight = self.telemetry_module.MAX_IN_FLIGHT_TELEMETRY_REQUESTS with patch.object( self.telemetry_module, "_send_telemetry_request_sync", side_effect=blocking_send ): - threads = [_send_telemetry_request(1, [1], MagicMock()) for _ in range(max_in_flight)] + threads = [_send_telemetry_request(1, [1], MagicMock()) for _ in range(event_count)] try: self.assertTrue(all(t is not None for t in threads)) - # One more than the cap allows is dropped rather than spawning - # an unbounded number of threads. - self.assertIsNone(_send_telemetry_request(1, [1], MagicMock())) + 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) - # Counter is released once the sends finish, so later events go through. - self.assertEqual(self.telemetry_module._in_flight_telemetry_requests, 0) - @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