Skip to content

Commit cd53577

Browse files
authored
fix(local): Align local mode with the SageMaker service (#6317)
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 #5283 (merged to master-v2, by aviruthen) to the V3 codebase; see #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 #3348 Fixes #4417 Fixes #4996 Fixes #5562 Fixes #5572 Fixes #5604 --- X-AI-Prompt: Fix S-effort PySDK V3 bugs, local-mode theme X-AI-Tool: Kiro
1 parent f3b8581 commit cd53577

14 files changed

Lines changed: 688 additions & 81 deletions

File tree

‎sagemaker-core/src/sagemaker/core/local/data.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -357,15 +357,19 @@ def pad(self, file, size=6):
357357
Returns:
358358
generator of records
359359
"""
360-
buffer = ""
360+
buffer = None
361361
for element in self.splitter.split(file):
362+
if buffer is None:
363+
# Match the buffer type to the record type so binary inputs
364+
# concatenate correctly instead of raising a TypeError.
365+
buffer = b"" if isinstance(element, bytes) else ""
362366
if _payload_size_within_limit(buffer + element, size):
363367
buffer += element
364368
else:
365369
tmp = buffer
366370
buffer = element
367371
yield tmp
368-
if _validate_payload_size(buffer, size):
372+
if buffer is not None and _validate_payload_size(buffer, size):
369373
yield buffer
370374

371375

‎sagemaker-core/src/sagemaker/core/local/image.py‎

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
from sagemaker.core.config.config_schema import CONTAINER_CONFIG, LOCAL
4040
import sagemaker.core
4141
from sagemaker.core.common_utils import custom_extractall_tarfile
42+
from sagemaker.core.constants import DIR_PARAM_NAME, SAGEMAKER_OUTPUT_LOCATION
4243

4344
CONTAINER_PREFIX = "algo"
4445
STUDIO_HOST_NAME = "sagemaker-local"
@@ -277,9 +278,7 @@ def train(self, input_data_config, output_data_config, hyperparameters, environm
277278
data_dir, input_data_config, output_data_config, hyperparameters
278279
)
279280
# If local, source directory needs to be updated to mounted /opt/ml/code path
280-
hyperparameters = self._update_local_src_path(
281-
hyperparameters, key=sagemaker.serve.model_builder.DIR_PARAM_NAME
282-
)
281+
hyperparameters = self._update_local_src_path(hyperparameters, key=DIR_PARAM_NAME)
283282

284283
# Create the configuration files for each container that we will create
285284
# Each container will map the additional local volumes (if any).
@@ -344,15 +343,15 @@ def serve(self, model_dir, environment):
344343
volumes = self._prepare_serving_volumes(model_dir)
345344

346345
# If the user script was passed as a file:// mount it to the container.
347-
if sagemaker.serve.model_builder.DIR_PARAM_NAME.upper() in environment:
348-
script_dir = environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()]
346+
if DIR_PARAM_NAME.upper() in environment:
347+
script_dir = environment[DIR_PARAM_NAME.upper()]
349348
parsed_uri = urlparse(script_dir)
350349
if parsed_uri.scheme == "file":
351350
host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
352351
volumes.append(_Volume(host_dir, "/opt/ml/code"))
353352
# Update path to mount location
354353
environment = environment.copy()
355-
environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()] = "/opt/ml/code"
354+
environment[DIR_PARAM_NAME.upper()] = "/opt/ml/code"
356355

357356
if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image):
358357
_pull_image(self.image)
@@ -583,8 +582,8 @@ def _prepare_training_volumes(
583582

584583
# If there is a training script directory and it is a local directory,
585584
# mount it to the container.
586-
if sagemaker.serve.model_builder.DIR_PARAM_NAME in hyperparameters:
587-
training_dir = json.loads(hyperparameters[sagemaker.serve.model_builder.DIR_PARAM_NAME])
585+
if DIR_PARAM_NAME in hyperparameters:
586+
training_dir = json.loads(hyperparameters[DIR_PARAM_NAME])
588587
parsed_uri = urlparse(training_dir)
589588
if parsed_uri.scheme == "file":
590589
host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
@@ -593,10 +592,7 @@ def _prepare_training_volumes(
593592
volumes.append(_Volume(shared_dir, "/opt/ml/shared"))
594593

595594
parsed_uri = urlparse(output_data_config["S3OutputPath"])
596-
if (
597-
parsed_uri.scheme == "file"
598-
and sagemaker.serve.model_builder.SAGEMAKER_OUTPUT_LOCATION in hyperparameters
599-
):
595+
if parsed_uri.scheme == "file" and SAGEMAKER_OUTPUT_LOCATION in hyperparameters:
600596
dir_path = os.path.abspath(parsed_uri.netloc + parsed_uri.path)
601597
intermediate_dir = os.path.join(dir_path, "output", "intermediate")
602598
if not os.path.exists(intermediate_dir):

‎sagemaker-core/src/sagemaker/core/local/local_session.py‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,25 @@ def __init__(self, sagemaker_session=None):
8686
"""
8787
self.sagemaker_session = sagemaker_session or LocalSession()
8888

89+
def describe_user_profile(self, DomainId, UserProfileName, **kwargs):
90+
"""Passes ``describe_user_profile`` through to the real SageMaker client.
91+
92+
Local mode does not model user profiles, but Studio role resolution calls
93+
this during session setup. Delegating to a real boto client keeps that path
94+
working instead of raising ``AttributeError``.
95+
96+
Args:
97+
DomainId (str): The domain ID the user profile belongs to.
98+
UserProfileName (str): The name of the user profile to describe.
99+
**kwargs: Additional keyword arguments forwarded to the boto client.
100+
101+
Returns: (dict) DescribeUserProfile response.
102+
"""
103+
boto_client = self.sagemaker_session.boto_session.client("sagemaker")
104+
return boto_client.describe_user_profile(
105+
DomainId=DomainId, UserProfileName=UserProfileName, **kwargs
106+
)
107+
89108
@_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_processing_job")
90109
def create_processing_job(
91110
self,
@@ -533,7 +552,28 @@ def invoke_endpoint(
533552
Body = Body.encode("utf-8")
534553
r = self.http.request("POST", url, body=Body, preload_content=False, headers=headers)
535554

536-
return {"Body": r, "ContentType": Accept}
555+
# Mirror the response shape of the real SageMaker runtime client so the same
556+
# code works in local mode and against a hosted endpoint. HTTP header names
557+
# are case-insensitive and containers send e.g. ``Content-Type`` or
558+
# ``content-type``, so look them up without regard to case.
559+
try:
560+
response_headers = dict(r.headers)
561+
except (TypeError, ValueError):
562+
response_headers = {}
563+
lowered = {str(k).lower(): v for k, v in response_headers.items()}
564+
request_id = lowered.get("x-amzn-requestid", "local-request-id")
565+
status_code = getattr(r, "status", None)
566+
return {
567+
"Body": r,
568+
"ContentType": lowered.get("content-type", Accept),
569+
"InvokedProductionVariant": TargetVariant or "AllTraffic",
570+
"ResponseMetadata": {
571+
"RequestId": request_id,
572+
"HTTPStatusCode": status_code,
573+
"HTTPHeaders": response_headers,
574+
"RetryAttempts": 0,
575+
},
576+
}
537577

538578

539579
class LocalSession(Session):

‎sagemaker-core/src/sagemaker/core/processing.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -278,11 +278,12 @@ def __init__(
278278
self.role = resolve_value_from_config(
279279
role, PROCESSING_JOB_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
280280
)
281-
if not self.role:
281+
if not self.role and self.instance_type not in ("local", "local_gpu"):
282282
# Originally IAM role was a required parameter.
283283
# Now we marked that as Optional because we can fetch it from SageMakerConfig
284284
# Because of marking that parameter as optional, we should validate if it is None, even
285-
# after fetching the config.
285+
# after fetching the config. In Local Mode the role is never used
286+
# (LocalSagemakerClient.create_processing_job discards it), so it is not required.
286287
raise ValueError("An AWS IAM role is required to create a Processing job.")
287288

288289
self.env = resolve_value_from_config(
@@ -760,11 +761,12 @@ def _get_process_args(self, inputs, outputs, experiment_config):
760761
process_request_args["network_config"] = self.network_config._to_request_dict()
761762
else:
762763
process_request_args["network_config"] = None
763-
process_request_args["role_arn"] = (
764-
self.role
765-
if is_pipeline_variable(self.role)
766-
else self.sagemaker_session.expand_role(self.role)
767-
)
764+
if self.role is None or is_pipeline_variable(self.role):
765+
# No role in Local Mode (see __init__): the local client discards RoleArn,
766+
# and expand_role(None) would raise.
767+
process_request_args["role_arn"] = self.role
768+
else:
769+
process_request_args["role_arn"] = self.sagemaker_session.expand_role(self.role)
768770
process_request_args["tags"] = self.tags
769771
return process_request_args
770772

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integ tests for sagemaker-core Local Mode session/runtime fixes."""
14+
15+
from __future__ import absolute_import

0 commit comments

Comments
 (0)