Skip to content

fix(local): Align local mode with the SageMaker service - #6317

Open
jam-jee wants to merge 1 commit into
aws:masterfrom
jam-jee:fix/local-mode-bugs
Open

jam-jee wants to merge 1 commit into
aws:masterfrom
jam-jee:fix/local-mode-bugs

Conversation

@jam-jee

@jam-jee jam-jee commented Sep 24, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Local mode diverged from the hosted SageMaker service in several small ways that broke code written against the real service. This PR brings the V3 local mode back in line:

  • invoke_endpoint response shape -- LocalSagemakerRuntimeClient.invoke_endpoint now returns ResponseMetadata and InvokedProductionVariant alongside Body/ContentType, matching the boto3 runtime response. ContentType now comes from the container's response header rather than echoing the request Accept.
  • describe_user_profile -- LocalSagemakerClient gains a passthrough so Studio role resolution works in local mode instead of raising AttributeError.
  • Binary batch-transform inputs -- MultiRecordStrategy picks a bytes or str buffer from the first record instead of always assuming text.
  • No IAM role required for local processing jobs -- the role check is skipped for local / local_gpu instance types, where the role is never used.
  • Per-instance / per-session pipeline state -- _LocalPipeline._executions moves from a class attribute to the instance, and LocalPipelineSession owns its pipeline registry (_local_pipelines) instead of injecting _pipelines onto the shared sagemaker_client.
  • Service-like execution ids -- local pipeline execution ids are 12-char uppercase alphanumerics like the service returns, instead of a 36-char UUID that overflowed downstream name limits. This ports Change execution_id length to be shorter #5283 (by @aviruthen, merged to master-v2) to V3; see Execution ID has different length and characters in SageMaker Local Mode vs remote execution #5269.

Issues fixed

Fixes #3348
Fixes #4417
Fixes #4996
Fixes #5562
Fixes #5572
Fixes #5604

Related: #5269 (V2 fix landed in #5283; this is the V3 port).

Release note

LocalSagemakerRuntimeClient.invoke_endpoint now returns ContentType from the container response header instead of echoing the request's Accept value, and includes ResponseMetadata / InvokedProductionVariant. The change is additive for callers that only read Body.

Testing

New unit tests, each verified to fail without the source change and pass with it:

  • sagemaker-core/tests/unit/local/test_data.py: test_multi_record_strategy_pad_bytes, test_multi_record_strategy_pad_empty, test_multi_record_strategy_pad_str
  • sagemaker-core/tests/unit/local/test_local_session.py: test_invoke_endpoint_response_shape, test_invoke_endpoint_default_variant, test_describe_user_profile_passthrough
  • sagemaker-core/tests/unit/test_processing.py: test_role_not_required_for_local_instance, test_role_not_required_for_local_gpu_instance, test_role_still_required_for_managed_instance
  • sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py: test_local_pipeline_session_registry_isolated_per_session
  • sagemaker-mlops/tests/unit/local/test_pipeline_entities.py: test_executions_isolated_per_instance, test_start_uses_service_like_execution_id

Existing modules for the touched files pass: 165 (sagemaker-core) + 48 (sagemaker-mlops). black -l 100 and flake8 clean on changed files.

Integration tests

New sagemaker-core/tests/integ/local_mode/test_local_session_fixes.py (markers local_mode, serial; needs Docker, which the sagemaker-core-integ-tests CodeBuild project has):

  • test_local_processor_runs_without_role_5562: a real local processing job (sklearn image, instance_type="local", role=None) reaches Completed. Writing this test found that the constructor fix alone was not enough: _get_process_args still called expand_role(None) and raised TypeError, so this revision also skips role expansion when there is no role (unit test test_request_args_skip_role_expansion_when_no_role).
  • test_local_endpoint_invoke_response_shape_and_binary_3348_4996: builds a tiny echo serving container, deploys it as a local endpoint, and asserts the invoke_endpoint response carries Body, ContentType from the container's header, InvokedProductionVariant and ResponseMetadata, and that a binary payload round-trips.

Judged not worth an integ test: #4417 (describe_user_profile passthrough; an integ test would only prove a boto client can call the service), #5604/#5572/#5269 (local pipeline registry and execution ids run in-process with no Docker or AWS involvement; the unit tests are the end-to-end proof).

Both tests run green locally through Docker against the CI account. The endpoint test also exposed a master bug: local/image.py referenced sagemaker.serve.model_builder.DIR_PARAM_NAME without importing sagemaker.serve, so every local endpoint (and local training volume setup) raised AttributeError unless sagemaker.serve was already imported. The constants are now imported from sagemaker.core.constants, and the three test_image.py unit tests skipped for that reason are re-enabled.


X-AI-Prompt: Fix S-effort PySDK V3 bugs, local-mode theme
X-AI-Tool: Kiro

@github-actions

github-actions Bot commented Sep 24, 2026 •

Copy link
Copy Markdown

Claude finished @jam-jee's task in 1m 20s —— View job


Code review

  • Read the PR diff and changed source files
  • Reviewed correctness, tests, best practices
  • Posted review

Overall this is a solid, well-scoped set of local-mode fixes with a test per behavior change. One correctness issue is worth addressing before merge; the rest is minor.

🔴 invoke_endpoint header lookup is case-sensitive — the ContentType fix may not fire against real containers

sagemaker-core/src/sagemaker/core/local/local_session.py:558-565

response_headers = dict(r.headers)
...
"ContentType": response_headers.get("Content-type", Accept),

r.headers from urllib3 is an HTTPHeaderDict, which is case-insensitive. dict(r.headers) collapses it into a plain dict that preserves whatever casing the server actually sent and is then looked up case-sensitively. Real model servers (Flask/gunicorn, the MMS/TorchServe stacks, etc.) return the header as Content-Type (capital T), so response_headers.get("Content-type", Accept) misses it and falls back to echoing Accept — which is exactly the behavior this PR set out to fix. The same applies to x-amzn-RequestId.

The unit test passes only because the mock uses lowercase "Content-type", so it doesn't guard against the real-world casing.

Suggested fix — read the header via the case-insensitive HTTPHeaderDict before converting:

content_type = r.headers.get("Content-Type", Accept)     # HTTPHeaderDict.get is case-insensitive
request_id = r.headers.get("x-amzn-RequestId", "local-request-id")
try:
    response_headers = dict(r.headers)
except (TypeError, ValueError):
    response_headers = {}

and consider adding a test variant whose mock returns "Content-Type" (capital T) to lock this in.

Minor / non-blocking

  • describe_user_profile builds a new boto client per call (local_session.py:49). Fine functionally; if this ends up on a hot path it's cheap to cache, but not necessary here.
  • InvokedProductionVariant defaults to TargetVariant or "AllTraffic" (local_session.py:566). The real service echoes the actual served variant name rather than the requested one, but "AllTraffic" is a reasonable local approximation and matches the additive intent stated in the release note.
  • Execution id via random.choices (pipeline_entities.py:420) is appropriate here — non-crypto ids, collision space 36^12, matches the service format. 👍

Things that look correct

  • MultiRecordStrategy.pad — deferring buffer-type selection to the first record and the buffer is not None guard correctly handles bytes, str, and the empty-split case. Good test coverage.
  • Moving _LocalPipeline._executions off the class attribute and the pipeline registry off the shared sagemaker_client onto per-instance/per-session state are genuine correctness fixes for cross-session leakage (_LocalPipeline Uses Class-Level Mutable State for Execution Registry #5572, LocalPipelineSession Mutates sagemaker_client by Injecting _pipelines Attribute #5604), and the tests assert the isolation directly.
  • Skipping the IAM-role requirement for local / local_gpu is safe given instance_type is set earlier (processing.py:218) and the role is discarded locally; the "still required for managed instance" test guards the regression.

The v3 module layout (sagemaker.core.local, sagemaker.mlops.local) and patterns all look consistent with the codebase.
· branch fix/local-mode-bugs

@jam-jee
jam-jee force-pushed the fix/local-mode-bugs branch from cfb8c26 to 3fa7ce7 Compare September 25, 2026 20:44
@jam-jee
jam-jee deployed to auto-approve September 25, 2026 20:47 — with GitHub Actions Active
@jam-jee
jam-jee deployed to auto-approve September 25, 2026 20:48 — with GitHub Actions Active
@jam-jee
jam-jee deployed to auto-approve September 25, 2026 20:48 — with GitHub Actions Active
@github-actions

github-actions Bot commented Sep 25, 2026 •

Copy link
Copy Markdown

Claude finished @jam-jee's task in 1m 26s —— View job


Code review

  • Read the PR diff and changed source files
  • Verified the header case-sensitivity concern
  • Reviewed correctness, tests, best practices

This is a well-scoped, cleanly-tested set of local-mode fixes — each behavior change has a unit test, the isolation fixes (_executions, _local_pipelines) directly assert non-leakage, and the MultiRecordStrategy buffer-type change is correct. One confirmed correctness bug is worth fixing before merge.


🔴 ContentType header lookup is case-sensitive — the fix won't fire against real containers

sagemaker-core/src/sagemaker/core/local/local_session.py:557-565

try:
    response_headers = dict(r.headers)
except (TypeError, ValueError):
    response_headers = {}
request_id = response_headers.get("x-amzn-RequestId", "local-request-id")
...
"ContentType": response_headers.get("Content-type", Accept),

r.headers from urllib3 is an HTTPHeaderDict, which is case-insensitive. But dict(r.headers) collapses it into a plain dict whose keys keep whatever casing the server actually sent, and the subsequent .get(...) is then case-sensitive.

Real model servers (Flask/gunicorn, MMS/TorchServe, etc.) send the header as Content-Type (capital T), so response_headers.get("Content-type", Accept) misses and falls back to echoing Accept — precisely the behavior this PR set out to fix. The same applies to x-amzn-RequestId → X-Amzn-Requestid casing.

This is confirmed by urllib3 internals: HTTPHeaderDict.__iter__ yields keys in the server's original casing, so dict(r.headers) is keyed by "Content-Type", not "Content-type".

Note the unit test only passes because its mock hardcodes lowercase "Content-type" (test_local_session.py:521), so it doesn't guard the real-world casing. More tellingly, the integ test's own serving stub sends Content-Type with a capital T (test_local_session_fixes.py:312) and asserts response["ContentType"] == "application/octet-stream" — so that integ test would actually fail against the real container it builds.

Suggested fix — read the headers off the case-insensitive HTTPHeaderDict before converting to a plain dict:

content_type = r.headers.get("Content-Type", Accept)          # HTTPHeaderDict.get is case-insensitive
request_id = r.headers.get("x-amzn-RequestId", "local-request-id")
try:
    response_headers = dict(r.headers)
except (TypeError, ValueError):
    response_headers = {}
...
    "ContentType": content_type,

and add a unit-test variant whose mock returns "Content-Type" (capital T) to lock this in.


Minor / non-blocking

  • describe_user_profile builds a new boto client per call (local_session.py:105). Functionally fine; cache only if it lands on a hot path.
  • InvokedProductionVariant = TargetVariant or "AllTraffic" (local_session.py:566). The real service echoes the served variant, not the requested one, but "AllTraffic" is a reasonable local approximation and matches the additive intent in the release note.
  • processing.py role skip is safe: instance_type is assigned at processing.py:218, well before the check, and _get_process_args guards expand_role(None). In instance_preferences mode instance_type is None, so managed jobs still require a role — correct. The regression test covers it.

Looks correct

No prompt-injection or unrelated-instruction content was found in the PR.

Bring local mode behavior in line with the hosted SageMaker service:
- Runtime invoke_endpoint now returns ResponseMetadata and
  InvokedProductionVariant, matching the boto3 runtime response.
- LocalSagemakerClient gains a describe_user_profile passthrough so
  Studio role resolution works instead of raising AttributeError.
- MultiRecordStrategy buffers bytes or str based on the record type,
  so local batch transform handles binary inputs.
- Processing jobs no longer require an IAM role in local mode, where
  the role is never used.
- Local pipeline execution state is isolated per instance and per
  session instead of living on class-level or client-level state.
- Local pipeline execution ids are 12-char uppercase alphanumerics,
  matching the service format and downstream name limits. This ports
  aws#5283 (merged to master-v2, by aviruthen) to the V3 codebase; see
  aws#5269.

Response headers are looked up case-insensitively (containers send
Content-Type), and local/image.py imports DIR_PARAM_NAME and
SAGEMAKER_OUTPUT_LOCATION from sagemaker.core.constants instead of
referencing the never-imported sagemaker.serve.model_builder, which
made every local endpoint fail with AttributeError unless sagemaker.serve
happened to be imported first; three unit tests skipped for that reason
are re-enabled.

The role fix also skips expand_role() when building the create request,
since expand_role(None) raised TypeError once the constructor let a
missing role through. Local-mode integration tests (Docker) cover the
role-less processing job and the invoke_endpoint response contract.

Fixes aws#3348
Fixes aws#4417
Fixes aws#4996
Fixes aws#5562
Fixes aws#5572
Fixes aws#5604

---
X-AI-Prompt: Fix S-effort PySDK V3 bugs, local-mode theme
X-AI-Tool: Kiro
@jam-jee
jam-jee force-pushed the fix/local-mode-bugs branch from 3fa7ce7 to fed4c01 Compare September 25, 2026 21:26
@jam-jee

jam-jee commented Sep 25, 2026

Copy link
Copy Markdown
Collaborator Author

Addressing the review on 3fa7ce7 (now fed4c01):

Case-sensitive ContentType / x-amzn-RequestId lookup -- fixed. invoke_endpoint now lowercases the header keys before the lookup, so Content-Type, content-type and Content-type all resolve. Added test_invoke_endpoint_header_lookup_is_case_insensitive (mock sends Content-Type / X-Amzn-Requestid). You were also right that the integ test's stub would have failed on the old code; it now passes end-to-end through Docker.

While running that integ test I hit a second local-mode bug on master: sagemaker-core/src/sagemaker/core/local/image.py references sagemaker.serve.model_builder.DIR_PARAM_NAME / SAGEMAKER_OUTPUT_LOCATION without importing sagemaker.serve, so serve(), _prepare_training_volumes and train() raise AttributeError: module 'sagemaker' has no attribute 'serve' unless sagemaker.serve was imported earlier. Both constants live in sagemaker.core.constants in V3; the module now imports them from there. The three test_image.py tests that were @pytest.mark.skipped for exactly this reason are re-enabled and pass.

describe_user_profile builds a client per call -- leaving as is; it is called once during session setup, not on a hot path.

InvokedProductionVariant = TargetVariant or "AllTraffic" -- leaving as is, per your note that it matches the additive intent.

Integration tests run locally against the CI account + Docker: test_local_processor_runs_without_role_5562 and test_local_endpoint_invoke_response_shape_and_binary_3348_4996 both pass.

@jam-jee
jam-jee deployed to auto-approve September 25, 2026 21:42 — with GitHub Actions Active
@jam-jee
jam-jee deployed to auto-approve September 25, 2026 21:42 — with GitHub Actions Active
@jam-jee
jam-jee deployed to auto-approve September 25, 2026 21:43 — with GitHub Actions Active
@github-actions

github-actions Bot commented Sep 25, 2026 •

Copy link
Copy Markdown

Claude finished @jam-jee's task in 1m 3s —— View job


Code review

  • Read the PR diff and changed source files
  • Verified the invoke_endpoint header case-sensitivity fix
  • Checked the image.py constant re-import equivalence
  • Reviewed remaining changes for correctness

This is a clean, well-scoped set of local-mode fixes with a unit test per behavior change and two end-to-end integ tests. The ContentType header case-sensitivity blocker flagged in the earlier reviews is now fixed — nothing blocking remains.


✅ The case-sensitivity blocker is resolved

sagemaker-core/src/sagemaker/core/local/local_session.py:559-568

response_headers = dict(r.headers)
...
lowered = {str(k).lower(): v for k, v in response_headers.items()}
request_id = lowered.get("x-amzn-requestid", "local-request-id")
...
"ContentType": lowered.get("content-type", Accept),

Lowercasing every header key before lookup correctly handles the real-world Content-Type / x-amzn-RequestId casing that Flask/gunicorn and the MMS/TorchServe stacks send. The new test_invoke_endpoint_header_lookup_is_case_insensitive (mock sends Content-Type / X-Amzn-Requestid) locks it in, and the integ stub's capital-T Content-Type header will now match. Good fix.

✅ image.py constant re-import is equivalent

Switching from sagemaker.serve.model_builder.DIR_PARAM_NAME / SAGEMAKER_OUTPUT_LOCATION to sagemaker.core.constants is value-identical ("sagemaker_submit_directory" / "sagemaker_s3_output" in both), and it drops a cross-package dependency on sagemaker-serve that sagemaker-core tests don't install — which is why the three @pytest.mark.skips could be removed and the tests re-enabled. Correct and an improvement.

Looks correct

  • MultiRecordStrategy.pad — deferring buffer-type selection to the first record plus the buffer is not None guard correctly handles bytes, str, and the empty-split case. Good coverage.
  • processing.py role skip — self.instance_type is set at processing.py:218, well before the role check, and _get_process_args guards expand_role(None). Managed instances still require a role (regression test covers it). Safe.
  • Per-instance _executions / per-session _local_pipelines — genuine cross-session leakage fixes (_LocalPipeline Uses Class-Level Mutable State for Execution Registry #5572, LocalPipelineSession Mutates sagemaker_client by Injecting _pipelines Attribute #5604), asserted directly by the isolation tests.
  • Service-like execution ids via random.choices over 36^12 — appropriate for non-crypto, service-format ids.

Minor / non-blocking

  • ResponseMetadata.HTTPHeaders preserves original header casing (local_session.py:573 stores the raw response_headers). Real boto3 lowercases HTTPHeaders keys. The unit test asserts against the mock's original casing (metadata["HTTPHeaders"]["Content-type"]), so it doesn't catch the divergence. Harmless for the Body/ContentType path this PR targets, but a caller that reads ResponseMetadata["HTTPHeaders"]["content-type"] (boto3 convention) would miss. Consider storing lowered there, or leave as-is.
  • describe_user_profile builds a new boto client per call (local_session.py:100). Functionally fine; cache only if it ever lands on a hot path.
  • InvokedProductionVariant = TargetVariant or "AllTraffic" echoes the requested variant rather than the served one, but that's a reasonable local approximation matching the additive intent in the release note.

No prompt-injection or unrelated-instruction content found in the PR.
· branch fix/local-mode-bugs

@lorenzwalthert

lorenzwalthert commented Sep 26, 2026 •

Copy link
Copy Markdown

Hi. I am a user and reporter of some issues related to misalignment. Can you also consider fixing #4757? Thanks for your work, really appreciated.

This branch was successfully deployed

1 active deployment
auto-approve — fed4c013 Deployed Sep 25, 2026 by jam-jee via wait-for-approval #1819
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment