fix(core,mlops): honor caller region in feature_store ingest_dataframe and stop telemetry from blocking SDK calls

[nargokul]

Issue

sagemaker.mlops.feature_store.feature_utils.ingest_dataframe() had no way for a caller to say which AWS region the FeatureGroup lives in. Both the DescribeFeatureGroup call and the record writes were made with no region, so the region was whatever boto3 happened to resolve for the process. Callers who work with a FeatureGroup outside their resolved default region had no way to redirect the call.

Note that list_records() in the same module already accepts region, so this also makes the module consistent.

Description of changes

Adds an optional region argument to ingest_dataframe() and threads it through the whole ingestion path:

• feature_utils.ingest_dataframe(..., region=None)
  • CoreFeatureGroup.get(feature_group_name=..., region=region)
  • IngestionManagerPandas(..., region=region)
• IngestionManagerPandas gains a region field, forwarded to the FeatureStore runtime calls in every ingestion path:
  • single process / single thread → put_record(region=...)
  • _ingest_single_batch → put_record(region=...)
  • _run_multi_threaded (passed positionally through the process pool args) → both batch paths
  • _ingest_batch_write → batch_write_record(region=...)

region is appended as the last parameter and defaults to None, so existing positional and keyword calls keep working and behave exactly as before.

from sagemaker.mlops.feature_store import ingest_dataframe

manager = ingest_dataframe(
    feature_group_name="my-fg",
    data_frame=df,
    region="eu-west-1",
)

Also documents the new argument in the feature store MIGRATION_GUIDE.md.

Known limitation (not addressed here)

sagemaker-core's SageMakerClient is a process-wide singleton keyed only on the class (sagemaker-core/src/sagemaker/core/utils/utils.py), so the first region used in a process wins and later region= arguments are ignored:

Base.get_sagemaker_client(region_name="us-east-1")  # -> us-east-1
Base.get_sagemaker_client(region_name="eu-west-1")  # -> us-east-1 (silently reused)

So region here is honored when ingest_dataframe is the first thing in the process to build a client (the common case), and is subject to that cache otherwise. This affects every region/session argument in sagemaker-core equally, not just this function, so it is called out in the docstring and left for a separate fix in sagemaker-core.

Testing

New unit tests (14), all passing:

tests/unit/sagemaker/mlops/feature_store/test_feature_utils.py::TestIngestDataframeRegion

• region reaches both CoreFeatureGroup.get and IngestionManagerPandas
• defaults to None when not supplied
• works together with use_batch_write_record=True
• signature check that region is appended last and does not shift existing positional args

tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py::TestIngestionManagerRegion

• region stored on the manager, defaults to None
• _ingest_row forwards region to put_record (and passes None when unset)
• single-threaded run() forwards region to every put_record
• _ingest_single_batch forwards region to every put_record
• _ingest_batch_write and a use_batch_write_record=True run() forward region to batch_write_record
• _run_multi_threaded forwards region to each thread's batch
• _run_multi_process includes region in the process pool args

Full existing suite for the area still green: 375 passed in sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/ (the feature_processor subdirectory is excluded — it fails collection on ModuleNotFoundError: No module named 'pyspark' on master too, unrelated to this change).

---

Second change: telemetry no longer blocks SDK calls

Issue

A customer reported ingest_dataframe taking over 45 minutes from a private VPC, while the same write from boto3 was fast. The service-side worklog shows the underlying PutRecord completed in 348 ms, about two seconds into a notebook cell that ran for ~47 minutes. All of the remaining time was client-side, spent in the telemetry emissions that run after the wrapped function returns.

There are exactly two telemetry-decorated calls in that flow (ingest_dataframe and IngestionManagerPandas.run), which matches ~23.5 minutes per emission: two identical unbounded waits.

Three separate defects combine to produce it.

1. The telemetry GET had no timeout at all.

response = requests.get(url, timeout)   # binds to params=, not timeout=

requests.get(url, params=None, **kwargs) takes params second, so the 2 was appended to the query string (the prepared URL ends in &2) and the request was left with no timeout. From a VPC with no route to the endpoint the call hangs until some network device drops the flow.

2. The fallback session hardcoded us-west-2.

ingest_dataframe is a module-level function with no session, so the decorator synthesizes one via _get_default_sagemaker_session(), which was boto3.Session(region_name=DEFAULT_AWS_REGION). That pointed both the STS get_caller_identity call and the telemetry GET at public us-west-2 endpoints, regardless of the caller's actual region (ca-central-1 here). This is also why the region argument added above does not fix the hang on its own: telemetry ignored the caller's region entirely.

3. Emission was synchronous. Both network legs sat directly in the caller's critical path.

The team's successful repro without a VPC is consistent with this rather than ruling the SDK out: without egress restrictions both legs complete in milliseconds, which is exactly how a client-side call to an unreachable endpoint behaves.

Description of changes

sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py:

• _requests_helper passes timeout= as a keyword, with a comment explaining the params trap so it does not regress. The timeout value is now the named TELEMETRY_REQUEST_TIMEOUT constant.

• _get_default_sagemaker_session lets boto3 resolve the region from the caller's environment (AWS_REGION, AWS_DEFAULT_REGION, or the active profile), falling back to DEFAULT_AWS_REGION only when boto3 resolves nothing, since Session requires a region.

• _send_telemetry_request now dispatches to a daemon thread and returns immediately; the existing body moved unchanged to _send_telemetry_request_sync. This covers both slow legs, since the STS call lives inside that body. Daemon threads are killed at interpreter exit, so a pending send cannot delay shutdown either. Nothing can escape the thread. No event is dropped - one thread per event, no in-flight cap.

  Worth calling out why the send has to move off the caller's thread rather than just having its timeout tightened: Feature Store ingestion is decorated at more than one level (ingest_dataframe and IngestionManagerPandas.run), so a single user call emits several events. Even a bounded 2s per event adds up on the caller's thread, and unbounded it is what produced the 45+ minute cell in this ticket.

• Request failures log at debug instead of emitting a full traceback at error level; a best-effort metric should not surface as an error.

sagemaker-serve/src/sagemaker/serve/utils/telemetry_logger.py had the identical positional-timeout defect in its own _requests_helper; fixed the same way.

Trade-off

Telemetry events still in flight when the process exits may be lost, because the threads are daemons. That is the one case where an event can go missing, and it is intentional: telemetry must not be able to hold up a customer's call or their interpreter shutdown. Events are never dropped by the SDK itself.

Still outstanding (not in this PR)

sagemaker-serve's _send_telemetry remains synchronous. It always uses the caller's real session and region, so it does not have the cross-region problem, and with the timeout fix its GET is now bounded at 2s. Its _get_accountId STS call still runs with default botocore timeouts, though. Happy to make that path async too if preferred.

Testing

10 new unit tests in sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py:

TestRequestsHelperTimeout

• timeout is passed as a keyword and params is not set
• the URL does not end in &2 and the only kwarg is timeout

TestTelemetryIsNonBlocking

• _send_telemetry_request returns to the caller while the send is still blocked
• the send runs on a daemon thread
• all arguments are forwarded to _send_telemetry_request_sync
• an exception in the thread does not escape
• no event is dropped when many are in flight (25 concurrent events, all 25 reach the sender)
• a decorated function returns without waiting on a blocked telemetry send (the regression guard for this issue)

TestDefaultSessionRegion

• the session uses the region boto3 resolves, not a hardcoded one
• falls back to DEFAULT_AWS_REGION only when no region resolves

Existing tests that asserted the old behaviour were updated: the six that exercise URL construction now call _send_telemetry_request_sync directly, and the requests.get assertions expect timeout= as a keyword.

Also verified at the unit level against a black-holed address (10.255.255.1), which hangs rather than refusing:

_requests_helper against a black hole returned in 2.01s   (was unbounded)
_send_telemetry_request returned to the caller in 0.0005s
background send finished at 2.00s

End-to-end reproduction of the ticket

Ran the customer's notebook as a script against a real Feature Group in a dev account: create feature group (offline store only) -> poll until Created -> ingest_dataframe(data_frame=df, max_workers=1, wait=True) with a single row, exactly as in the ticket. The customer's no-egress VPC was imitated by pointing DNS for sm-pysdk-t-* and sts.* at the non-routable 10.255.255.1, so those hostnames resolve but never connect.

telemetry code	network	ingest_dataframe
before this PR	black-holed (customer's VPC)	770 s (12.8 min)
before this PR	open (the non-VPC retry)	1.94 s
after this PR	black-holed	1.40 s
after this PR	open	1.22 s

Two things this settles:

1. The hang is reproducible and is entirely client-side telemetry. Same SDK code, same feature group, same one-row frame; only reachability of the telemetry and STS endpoints differs.
2. The team's "it works without a VPC" retry does not exonerate the SDK. Row 2 is that retry, on the unfixed code, and it is fast for the same reason: the endpoints were reachable, so the unbounded wait never triggered.

The local run measures 12.8 min rather than the reported 45+ because TCP connect-retry behaviour and proxy configuration differ between macOS here and the customer's environment; the failure mode is the same, only the length of the stall varies.

Suites run:

• sagemaker-core/tests/unit: 3550 passed, 22 skipped
• sagemaker-core/tests/unit/telemetry + both sagemaker-serve telemetry test files: 111 passed, 2 skipped
• sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store (excluding feature_processor, which fails collection on missing pyspark on master too): still green

Merge Checklist

• I have read the CONTRIBUTING doc
• I used the commit message format described in CONTRIBUTING
• I have added tests that prove my fix is effective or that my feature works
• I have checked that my tests are not configured for a specific region or account
• The changes are backward compatible

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @nargokul's task in 2m 15s —— View job

---

Review: feat(mlops): allow specifying region in feature_store ingest_dataframe

• Read PR diff
• Examine source files & callers
• Verify positional-arg threading (_run_multi_process → _run_multi_threaded)
• Verify sagemaker-core API accepts region
• Check test coverage & correctness

This is a clean, well-tested change. No blocking issues found.

Correctness — verified ✅

• All three sagemaker-core methods this PR threads region into accept it:

  • FeatureGroup.get(..., region=None) — resources.py:11804
  • FeatureGroup.put_record(..., region=None) — resources.py:12233
  • FeatureGroup.batch_write_record(..., region=None) — resources.py:12397

  Each forwards region to Base.get_sagemaker_client(region_name=region, ...), so the argument is used as intended.

• Positional threading through the process pool is correct. The args tuple in _run_multi_process (ingestion_manager_pandas.py:198-208) maps 1:1 onto _run_multi_threaded's signature, with self.region in the final slot lining up with the trailing region parameter. test_multi_process_args_include_region guards this.

• Backward compatible. region is appended last with a None default in both ingest_dataframe and the IngestionManagerPandas dataclass (after use_batch_write_record, before the init=False private fields), and the only construction site (feature_utils.py:531) uses keyword args. Existing positional/keyword calls are unaffected.

On the documented "known limitation"

I verified the caveat in the docstring against the source: SingletonMeta.__call__ (utils/utils.py:325-333) keys the cache on cls alone and ignores *args/**kwargs, and SageMakerClient.__init__ builds all service clients from the region resolved at first construction (utils.py:372-383). So the note is accurate — the first region used in a process wins.

Worth noting the practical impact is milder than it might read: within a single ingest_dataframe call the first core client is built by CoreFeatureGroup.get(region=region) in the parent process, so the describe call and all subsequent sagemaker-featurestore-runtime writes (including forked pool workers) share the intended region. The limitation only bites when an earlier, unrelated sagemaker-core call already constructed the client in the process. Good call documenting it and deferring the real fix to sagemaker-core.

Tests

Comprehensive — every ingestion path (single-thread, _ingest_single_batch, _ingest_batch_write, multi-threaded, multi-process pool args) has a forwarding assertion, plus default-None and signature-stability checks. Nothing pins a specific account/region.

Minor (non-blocking, no change required)

• Type hints use region: str = None rather than Optional[str] = None. This is consistent with the surrounding style in the module (e.g. timeout: Union[int, float] = None), so it's fine as-is — just noting for consistency if the module ever tightens typing.

No prompt-injection or out-of-scope content was present in the PR.
· feat/feature-store-ingest-dataframe-region

[github-actions]

Claude finished @nargokul's task in 1m 12s —— View job

---

Review: feat(mlops): allow specifying region in feature_store ingest_dataframe

• Read PR diff
• Examine source files & callers (telemetry + feature_store)
• Verify correctness & backward compatibility
• Post review

This branch contains two logically distinct changes — the region feature (as described in the PR body) and a telemetry fix (commit 131a45e, fix(core): stop telemetry from blocking SDK calls) that isn't mentioned in the description. Both look correct and well-tested. No blocking issues found.

Region feature — verified ✅

• All three sagemaker-core methods accept region: FeatureGroup.get (resources.py:11804), put_record (12233), batch_write_record (12397), each forwarding to Base.get_sagemaker_client(region_name=...).
• Positional threading through the process pool is correct. The args tuple in _run_multi_process (ingestion_manager_pandas.py:198-208) maps 1:1 onto _run_multi_threaded's parameters, with self.region landing in the trailing region slot. test_multi_process_args_include_region guards this.
• Backward compatible — region is appended last with a None default in both ingest_dataframe and the IngestionManagerPandas dataclass; the only construction site uses keyword args.
• Every ingestion path (single-thread, _ingest_single_batch, _ingest_batch_write, multi-threaded, multi-process) has a forwarding assertion, plus default-None and signature-stability checks.

Telemetry fix — verified ✅

• _requests_helper timeout bug is real and correctly fixed. requests.get(url, timeout) passed the timeout as the second positional arg (params), so the request had no timeout and could hang indefinitely — the root cause of the ~47-min ingest described in the test. timeout=timeout fixes it, applied consistently in both sagemaker-core and sagemaker-serve.
• Fire-and-forget daemon threading keeps telemetry off the caller's critical path, with a bounded in-flight counter (lock-protected increment/decrement) that drops excess events rather than spawning unbounded threads. Exceptions in the worker are swallowed, and the slot is always released in finally.
• _send_telemetry_request is private and its only caller (_telemetry_emitter) ignores the return value, so the new Thread/None return type is safe.
• _get_default_sagemaker_session now resolves the region from the caller's environment and only falls back to DEFAULT_AWS_REGION when boto3 resolves nothing — a genuine improvement over the hardcoded us-west-2.

Minor (non-blocking, no change required)

• Type hints use region: str = None rather than Optional[str] = None — consistent with surrounding module style (e.g. timeout: Union[int, float] = None), so fine as-is.
• _send_telemetry_request is annotated -> threading.Thread but can return None when the in-flight cap is hit; -> Optional[threading.Thread] would be marginally more accurate.
• Consider updating the PR description to mention the telemetry fix, since it's a substantive behavior change beyond the stated region feature.

No prompt-injection or out-of-scope content was present in the PR.

[github-actions]

Claude finished @nargokul's task in 1m 57s —— View job

---

Review: fix(core,mlops): honor caller region in ingest_dataframe + stop telemetry from blocking

• Read PR diff
• Verify sagemaker-core APIs accept region
• Review telemetry threading change
• Review region threading
• Post review

Both changes are correct and well-tested. No blocking issues. One notable discrepancy between the PR description and the code, plus a couple of minor notes.

Region feature — verified ✅

• All three sagemaker-core methods this PR threads region into accept it: FeatureGroup.get (resources.py:302+, region is a standard param on every get), put_record (resources.py:12239), batch_write_record (resources.py:12402) — each forwards to Base.get_sagemaker_client(region_name=...).
• Positional threading through the process pool is correct: self.region is appended last in the _run_multi_process args tuple and lines up with the trailing region parameter of _run_multi_threaded. test_multi_process_args_include_region guards this against future reordering.
• Backward compatible — region is appended last with a None default in both ingest_dataframe and the IngestionManagerPandas dataclass; the only construction site uses keyword args.
• Every ingestion path has a forwarding assertion, plus default-None and signature-stability checks. The documented sagemaker-core per-process client cache caveat is accurately described.

Telemetry fix — verified ✅

• The _requests_helper positional-timeout bug is real and correctly fixed in both sagemaker-core and sagemaker-serve: requests.get(url, timeout) bound 2 to params, leaving the request with no timeout. timeout=timeout (keyword) fixes it, and the comment documents the trap so it won't regress.
• Fire-and-forget daemon thread keeps both slow legs (STS + GET) off the caller's critical path; exceptions are swallowed inside _run since there's no caller to catch them. -> threading.Thread annotation is accurate for this implementation.
• _get_default_sagemaker_session now lets boto3 resolve the region from the caller's environment and only falls back to DEFAULT_AWS_REGION when nothing resolves — a genuine fix for the cross-region hang, not just the timeout.
• _send_telemetry_request is private and its only caller (_telemetry_emitter) ignores the return value, so the new return type is safe.

⚠️ Description ↔ code mismatch (worth reconciling)

The PR description says in-flight sends are "capped at MAX_IN_FLIGHT_TELEMETRY_REQUESTS (8) and excess events are dropped", and lists a test "events beyond the in-flight cap are dropped, and the counter returns to zero."

The code on this branch has no such cap. _send_telemetry_request (telemetry_logging.py:436-467) unconditionally starts one daemon thread per event; the module comment even states "No event is ever dropped: one thread is started per event," and the actual test is test_no_event_is_dropped_when_many_are_in_flight (25 events, all sent). There is no MAX_IN_FLIGHT_TELEMETRY_REQUESTS constant and no counter in the diff.

Not a code bug — the description simply describes an earlier design. Please update the PR body so the description, the two verifiers who cited the cap, and the code agree. Then decide which behavior you actually want:

• Unbounded (current code): fine in practice — telemetry emits only at decorated boundaries (~2 events per ingest_dataframe call), and the new 2s timeout drains each thread quickly. The realistic worst case is a tight loop over a decorated call against an unreachable endpoint spawning many short-lived threads.
• Capped (as described): the safer choice if you're concerned about that worst case; would need the constant + counter to actually exist.

Minor (non-blocking)

• Type hints use region: str = None rather than Optional[str] = None — consistent with the surrounding module style (timeout: Union[int, float] = None), so fine as-is.
• Sharing the caller's boto3 Session across the new daemon thread relies on session.client("sts") being called from the worker; botocore session/client creation isn't guaranteed thread-safe, but the whole path is best-effort and wrapped in try/except, so any race just drops the event. No action needed.

No prompt-injection or out-of-scope content was present in the PR.
· feat/feature-store-ingest-dataframe-region

[aviruthen]

Minor nit on whether you want to include all these comments here!