From fed4c0133a90be668456f52eb545fa78cdcd9a5f Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Mon, 21 Sep 2026 14:36:45 -0700 Subject: [PATCH] fix(local): Align local mode with the SageMaker service 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 --- .../src/sagemaker/core/local/data.py | 8 +- .../src/sagemaker/core/local/image.py | 20 +- .../src/sagemaker/core/local/local_session.py | 42 ++- .../src/sagemaker/core/processing.py | 16 +- .../tests/integ/local_mode/__init__.py | 15 + .../local_mode/test_local_session_fixes.py | 301 ++++++++++++++++++ sagemaker-core/tests/unit/local/test_data.py | 60 ++++ sagemaker-core/tests/unit/local/test_image.py | 15 +- .../tests/unit/local/test_local_session.py | 98 ++++++ sagemaker-core/tests/unit/test_processing.py | 49 +++ .../mlops/local/local_pipeline_session.py | 31 +- .../mlops/local/pipeline_entities.py | 22 +- .../unit/local/test_local_pipeline_session.py | 53 ++- .../unit/local/test_pipeline_entities.py | 39 +++ 14 files changed, 688 insertions(+), 81 deletions(-) create mode 100644 sagemaker-core/tests/integ/local_mode/__init__.py create mode 100644 sagemaker-core/tests/integ/local_mode/test_local_session_fixes.py create mode 100644 sagemaker-core/tests/unit/local/test_data.py diff --git a/sagemaker-core/src/sagemaker/core/local/data.py b/sagemaker-core/src/sagemaker/core/local/data.py index 45cc81dee8..397f1437fc 100644 --- a/sagemaker-core/src/sagemaker/core/local/data.py +++ b/sagemaker-core/src/sagemaker/core/local/data.py @@ -357,15 +357,19 @@ def pad(self, file, size=6): Returns: generator of records """ - buffer = "" + buffer = None for element in self.splitter.split(file): + if buffer is None: + # Match the buffer type to the record type so binary inputs + # concatenate correctly instead of raising a TypeError. + buffer = b"" if isinstance(element, bytes) else "" if _payload_size_within_limit(buffer + element, size): buffer += element else: tmp = buffer buffer = element yield tmp - if _validate_payload_size(buffer, size): + if buffer is not None and _validate_payload_size(buffer, size): yield buffer diff --git a/sagemaker-core/src/sagemaker/core/local/image.py b/sagemaker-core/src/sagemaker/core/local/image.py index 1f42a06bd0..9e8e732302 100644 --- a/sagemaker-core/src/sagemaker/core/local/image.py +++ b/sagemaker-core/src/sagemaker/core/local/image.py @@ -39,6 +39,7 @@ from sagemaker.core.config.config_schema import CONTAINER_CONFIG, LOCAL import sagemaker.core from sagemaker.core.common_utils import custom_extractall_tarfile +from sagemaker.core.constants import DIR_PARAM_NAME, SAGEMAKER_OUTPUT_LOCATION CONTAINER_PREFIX = "algo" STUDIO_HOST_NAME = "sagemaker-local" @@ -277,9 +278,7 @@ def train(self, input_data_config, output_data_config, hyperparameters, environm data_dir, input_data_config, output_data_config, hyperparameters ) # If local, source directory needs to be updated to mounted /opt/ml/code path - hyperparameters = self._update_local_src_path( - hyperparameters, key=sagemaker.serve.model_builder.DIR_PARAM_NAME - ) + hyperparameters = self._update_local_src_path(hyperparameters, key=DIR_PARAM_NAME) # Create the configuration files for each container that we will create # Each container will map the additional local volumes (if any). @@ -344,15 +343,15 @@ def serve(self, model_dir, environment): volumes = self._prepare_serving_volumes(model_dir) # If the user script was passed as a file:// mount it to the container. - if sagemaker.serve.model_builder.DIR_PARAM_NAME.upper() in environment: - script_dir = environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()] + if DIR_PARAM_NAME.upper() in environment: + script_dir = environment[DIR_PARAM_NAME.upper()] parsed_uri = urlparse(script_dir) if parsed_uri.scheme == "file": host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path) volumes.append(_Volume(host_dir, "/opt/ml/code")) # Update path to mount location environment = environment.copy() - environment[sagemaker.serve.model_builder.DIR_PARAM_NAME.upper()] = "/opt/ml/code" + environment[DIR_PARAM_NAME.upper()] = "/opt/ml/code" if _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image): _pull_image(self.image) @@ -583,8 +582,8 @@ def _prepare_training_volumes( # If there is a training script directory and it is a local directory, # mount it to the container. - if sagemaker.serve.model_builder.DIR_PARAM_NAME in hyperparameters: - training_dir = json.loads(hyperparameters[sagemaker.serve.model_builder.DIR_PARAM_NAME]) + if DIR_PARAM_NAME in hyperparameters: + training_dir = json.loads(hyperparameters[DIR_PARAM_NAME]) parsed_uri = urlparse(training_dir) if parsed_uri.scheme == "file": host_dir = os.path.abspath(parsed_uri.netloc + parsed_uri.path) @@ -593,10 +592,7 @@ def _prepare_training_volumes( volumes.append(_Volume(shared_dir, "/opt/ml/shared")) parsed_uri = urlparse(output_data_config["S3OutputPath"]) - if ( - parsed_uri.scheme == "file" - and sagemaker.serve.model_builder.SAGEMAKER_OUTPUT_LOCATION in hyperparameters - ): + if parsed_uri.scheme == "file" and SAGEMAKER_OUTPUT_LOCATION in hyperparameters: dir_path = os.path.abspath(parsed_uri.netloc + parsed_uri.path) intermediate_dir = os.path.join(dir_path, "output", "intermediate") if not os.path.exists(intermediate_dir): diff --git a/sagemaker-core/src/sagemaker/core/local/local_session.py b/sagemaker-core/src/sagemaker/core/local/local_session.py index 663b7b4a09..9d6ab7b199 100644 --- a/sagemaker-core/src/sagemaker/core/local/local_session.py +++ b/sagemaker-core/src/sagemaker/core/local/local_session.py @@ -86,6 +86,25 @@ def __init__(self, sagemaker_session=None): """ self.sagemaker_session = sagemaker_session or LocalSession() + def describe_user_profile(self, DomainId, UserProfileName, **kwargs): + """Passes ``describe_user_profile`` through to the real SageMaker client. + + Local mode does not model user profiles, but Studio role resolution calls + this during session setup. Delegating to a real boto client keeps that path + working instead of raising ``AttributeError``. + + Args: + DomainId (str): The domain ID the user profile belongs to. + UserProfileName (str): The name of the user profile to describe. + **kwargs: Additional keyword arguments forwarded to the boto client. + + Returns: (dict) DescribeUserProfile response. + """ + boto_client = self.sagemaker_session.boto_session.client("sagemaker") + return boto_client.describe_user_profile( + DomainId=DomainId, UserProfileName=UserProfileName, **kwargs + ) + @_telemetry_emitter(Feature.LOCAL_MODE, "local_session.create_processing_job") def create_processing_job( self, @@ -533,7 +552,28 @@ def invoke_endpoint( Body = Body.encode("utf-8") r = self.http.request("POST", url, body=Body, preload_content=False, headers=headers) - return {"Body": r, "ContentType": Accept} + # Mirror the response shape of the real SageMaker runtime client so the same + # code works in local mode and against a hosted endpoint. HTTP header names + # are case-insensitive and containers send e.g. ``Content-Type`` or + # ``content-type``, so look them up without regard to case. + try: + response_headers = dict(r.headers) + except (TypeError, ValueError): + response_headers = {} + lowered = {str(k).lower(): v for k, v in response_headers.items()} + request_id = lowered.get("x-amzn-requestid", "local-request-id") + status_code = getattr(r, "status", None) + return { + "Body": r, + "ContentType": lowered.get("content-type", Accept), + "InvokedProductionVariant": TargetVariant or "AllTraffic", + "ResponseMetadata": { + "RequestId": request_id, + "HTTPStatusCode": status_code, + "HTTPHeaders": response_headers, + "RetryAttempts": 0, + }, + } class LocalSession(Session): diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 9dec7060bc..cf9338dc45 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -278,11 +278,12 @@ def __init__( self.role = resolve_value_from_config( role, PROCESSING_JOB_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session ) - if not self.role: + if not self.role and self.instance_type not in ("local", "local_gpu"): # Originally IAM role was a required parameter. # Now we marked that as Optional because we can fetch it from SageMakerConfig # Because of marking that parameter as optional, we should validate if it is None, even - # after fetching the config. + # after fetching the config. In Local Mode the role is never used + # (LocalSagemakerClient.create_processing_job discards it), so it is not required. raise ValueError("An AWS IAM role is required to create a Processing job.") self.env = resolve_value_from_config( @@ -750,11 +751,12 @@ def _get_process_args(self, inputs, outputs, experiment_config): process_request_args["network_config"] = self.network_config._to_request_dict() else: process_request_args["network_config"] = None - process_request_args["role_arn"] = ( - self.role - if is_pipeline_variable(self.role) - else self.sagemaker_session.expand_role(self.role) - ) + if self.role is None or is_pipeline_variable(self.role): + # No role in Local Mode (see __init__): the local client discards RoleArn, + # and expand_role(None) would raise. + process_request_args["role_arn"] = self.role + else: + process_request_args["role_arn"] = self.sagemaker_session.expand_role(self.role) process_request_args["tags"] = self.tags return process_request_args diff --git a/sagemaker-core/tests/integ/local_mode/__init__.py b/sagemaker-core/tests/integ/local_mode/__init__.py new file mode 100644 index 0000000000..99b2735713 --- /dev/null +++ b/sagemaker-core/tests/integ/local_mode/__init__.py @@ -0,0 +1,15 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integ tests for sagemaker-core Local Mode session/runtime fixes.""" + +from __future__ import absolute_import diff --git a/sagemaker-core/tests/integ/local_mode/test_local_session_fixes.py b/sagemaker-core/tests/integ/local_mode/test_local_session_fixes.py new file mode 100644 index 0000000000..c76ba0f3e7 --- /dev/null +++ b/sagemaker-core/tests/integ/local_mode/test_local_session_fixes.py @@ -0,0 +1,301 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""End-to-end Local Mode tests for the sagemaker-core session/runtime fixes. + +Each test locks in one bug fixed in this PR and states what a pass proves: + +- #5562 -- ``Processor`` no longer requires an IAM role when ``instance_type`` is + ``local``/``local_gpu``. A pass proves a real local processing job runs through + Docker with ``role=None`` and reaches ``Completed`` with no role ``ValueError`` + and no ``expand_role(None)`` failure while building the request. +- #3348 + #4996 -- ``LocalSagemakerRuntimeClient.invoke_endpoint`` returns the + full SageMaker runtime response shape (``Body``, ``ContentType`` taken from the + container's response header, ``InvokedProductionVariant`` and + ``ResponseMetadata``) and binary payloads round-trip through the local invoke + path unchanged. A pass drives a real container endpoint end-to-end through + Docker and asserts the response contract and the bytes round-trip. + +#4417 (``describe_user_profile`` passthrough) is covered by unit tests only: an +integ test would just prove that a boto client can call the service. + +These tests need Docker (local mode); #5562 pulls a public SageMaker ECR image. +Where a local blocker prevents a step (ECR pull denied) the test skips with a +precise reason rather than weakening the assertion. Marked ``local_mode`` and +``serial`` so local-mode tests do not run concurrently. +""" + +from __future__ import absolute_import + +import fcntl +import os +import shutil +import tempfile +import textwrap +import time +import uuid +from contextlib import contextmanager + +import boto3 +import botocore.exceptions +import pytest + +from sagemaker.core import image_uris +from sagemaker.core.local.local_session import ( + LocalSagemakerClient, + LocalSagemakerRuntimeClient, + LocalSession, +) +from sagemaker.core.processing import Processor + +REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-west-2")) +LOCK_PATH = os.path.join(tempfile.gettempdir(), "sagemaker_test_local_mode_lock") + +pytestmark = [pytest.mark.local_mode, pytest.mark.serial] + + +@contextmanager +def _local_mode_lock(path=LOCK_PATH): + """Serialize local-mode tests that share Docker and the fixed 8080 port.""" + f = open(path, "w") + try: + fcntl.lockf(f.fileno(), fcntl.LOCK_EX) + yield + finally: + time.sleep(5) + fcntl.lockf(f.fileno(), fcntl.LOCK_UN) + f.close() + + +def _boto_session(): + return boto3.Session(region_name=REGION) + + +# --------------------------------------------------------------------------- # +# #5562 -- local processing job without a role +# --------------------------------------------------------------------------- # +def test_local_processor_runs_without_role_5562(): + """A local processing job must run with ``role=None``. + + Before the fix, ``Processor.__init__`` raised ``ValueError('An AWS IAM role + is required...')`` even for ``instance_type='local'`` where the role is never + used. A pass proves the constructor accepts ``role=None`` in local mode and + the job runs through Docker to ``Completed``. + """ + with _local_mode_lock(): + image_uri = image_uris.retrieve( + "sklearn", REGION, version="1.2-1", py_version="py3", instance_type="ml.m5.large" + ) + + # Constructing with role=None must NOT raise -- this is the fix itself. + processor = Processor( + image_uri=image_uri, + instance_type="local", + instance_count=1, + role=None, + sagemaker_session=LocalSession(boto_session=_boto_session()), + entrypoint=["python3", "-c", 'print("ok")'], + ) + assert processor.role is None + + job_name = "local-proc-no-role-%s" % uuid.uuid4().hex[:8] + try: + # LocalSagemakerClient.create_processing_job runs the container synchronously, + # so the job is finished when run() returns. wait=False because the V3 + # ProcessingJob resource waiter polls the real service client, not local mode. + processor.run(wait=False, logs=False, job_name=job_name) + except botocore.exceptions.ClientError as err: + code = err.response.get("Error", {}).get("Code", "") + if code in ("AccessDenied", "AccessDeniedException", "UnrecognizedClientException"): + pytest.skip(f"ECR pull for the sklearn image denied here: {code}") + raise + except Exception as err: # pylint: disable=broad-except + msg = str(err).lower() + if "pull" in msg and ("denied" in msg or "unauthorized" in msg or "403" in msg): + pytest.skip(f"ECR image pull denied locally: {err}") + if "toomanyrequests" in msg or "rate limit" in msg: + pytest.skip(f"Registry pull rate-limited locally: {err}") + raise + + described = processor.sagemaker_session.sagemaker_client.describe_processing_job( + ProcessingJobName=job_name + ) + assert described["ProcessingJobStatus"] == "Completed", described + + +# --------------------------------------------------------------------------- # +# #3348 + #4996 -- local endpoint invoke response shape and binary payload +# --------------------------------------------------------------------------- # +_DOCKERFILE = """\ +FROM python:3.10-slim +COPY serve.py /serve.py +ENTRYPOINT ["python3", "/serve.py"] +""" + +_SERVE_PY = textwrap.dedent('''\ + """Tiny SageMaker-style serving container for local endpoint tests. + + GET /ping -> 200. POST /invocations -> echoes the request body back with + Content-Type: application/octet-stream and an x-amzn-RequestId header, so the + test can prove ContentType comes from the container response (not the echoed + Accept) and that a binary payload round-trips unchanged. + """ + from http.server import BaseHTTPRequestHandler, HTTPServer + + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/ping": + self.send_response(200) + self.end_headers() + else: + self.send_response(404) + self.end_headers() + + def do_POST(self): + if self.path == "/invocations": + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("x-amzn-RequestId", "local-echo-req-id") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404) + self.end_headers() + + def log_message(self, *args): + pass + + + if __name__ == "__main__": + HTTPServer(("0.0.0.0", 8080), Handler).serve_forever() + ''') + + +def _docker_available(): + import subprocess + + try: + return subprocess.run(["docker", "info"], capture_output=True).returncode == 0 + except (OSError, ValueError): + return False + + +def _build_serving_image(tag): + import subprocess + + build_dir = tempfile.mkdtemp(prefix="local-serve-") + with open(os.path.join(build_dir, "Dockerfile"), "w") as f: + f.write(_DOCKERFILE) + with open(os.path.join(build_dir, "serve.py"), "w") as f: + f.write(_SERVE_PY) + result = subprocess.run( + ["docker", "build", "-t", tag, build_dir], capture_output=True, text=True + ) + return result + + +def test_local_endpoint_invoke_response_shape_and_binary_3348_4996(): + """Local endpoint invoke returns the full runtime shape and round-trips bytes. + + Before the fix, ``invoke_endpoint`` returned only ``{"Body", "ContentType"}`` + (ContentType echoing the request Accept), so callers written against the real + runtime client's ``ResponseMetadata``/``InvokedProductionVariant`` shape broke + in local mode (#3348), and the response ContentType did not reflect what the + container actually returned. A pass drives a real container endpoint through + Docker and asserts: + + - ``Body`` is present and the bytes payload round-trips unchanged (locks the + binary path shared with the #4996 MultiRecordStrategy fix -- see note below); + - ``ContentType`` comes from the container's response header + (``application/octet-stream``), not the echoed ``Accept``; + - ``InvokedProductionVariant`` is present; + - ``ResponseMetadata`` carries ``RequestId`` (from the container's + ``x-amzn-RequestId``), ``HTTPStatusCode``, ``HTTPHeaders`` and + ``RetryAttempts``. + """ + if not _docker_available(): + pytest.skip("Docker is not available for local endpoint test.") + + with _local_mode_lock(): + tag = "sagemaker-local-echo:%s" % uuid.uuid4().hex[:8] + build = _build_serving_image(tag) + if build.returncode != 0: + stderr = (build.stderr or "").lower() + if "toomanyrequests" in stderr or "rate limit" in stderr: + pytest.skip("Docker Hub pull rate-limited building the serving image.") + if "pull access denied" in stderr or "unauthorized" in stderr: + pytest.skip("python:3.10-slim base image pull denied locally.") + raise AssertionError("docker build failed:\n%s" % build.stderr) + + session = LocalSession(boto_session=_boto_session()) + client = LocalSagemakerClient(session) + runtime = LocalSagemakerRuntimeClient(session.config) + + suffix = uuid.uuid4().hex[:8] + model_name = "local-echo-model-%s" % suffix + config_name = "local-echo-config-%s" % suffix + endpoint_name = "local-echo-endpoint-%s" % suffix + + # _LocalEndpoint mounts ModelDataUrl at /opt/ml/model; an empty local dir is enough. + model_dir = tempfile.mkdtemp(prefix="sagemaker-local-echo-model-") + try: + client.create_model( + ModelName=model_name, + PrimaryContainer={ + "Image": tag, + "ModelDataUrl": "file://" + model_dir, + "Environment": {}, + }, + ) + client.create_endpoint_config( + EndpointConfigName=config_name, + ProductionVariants=[ + { + "VariantName": "AllTraffic", + "ModelName": model_name, + "InitialInstanceCount": 1, + "InstanceType": "local", + } + ], + ) + client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name) + + payload = bytes(range(256)) # non-UTF-8 bytes -> proves binary round-trip + response = runtime.invoke_endpoint( + Body=payload, + EndpointName=endpoint_name, + ContentType="application/octet-stream", + Accept="application/json", + ) + + # Response shape (#3348) + assert "Body" in response + assert response["ContentType"] == "application/octet-stream", response["ContentType"] + assert "InvokedProductionVariant" in response + metadata = response["ResponseMetadata"] + assert metadata["RequestId"] == "local-echo-req-id", metadata + assert metadata["HTTPStatusCode"] == 200, metadata + assert isinstance(metadata["HTTPHeaders"], dict) + assert metadata["RetryAttempts"] == 0 + + # Binary round-trip (path shared with the #4996 buffer-type fix) + echoed = response["Body"].read() + assert echoed == payload, (len(echoed), len(payload)) + finally: + client.delete_endpoint(endpoint_name) + client.delete_endpoint_config(config_name) + client.delete_model(model_name) + shutil.rmtree(model_dir, ignore_errors=True) diff --git a/sagemaker-core/tests/unit/local/test_data.py b/sagemaker-core/tests/unit/local/test_data.py new file mode 100644 index 0000000000..9764d34c5f --- /dev/null +++ b/sagemaker-core/tests/unit/local/test_data.py @@ -0,0 +1,60 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. + +"""Unit tests for sagemaker.core.local.data batch strategies""" + +from __future__ import absolute_import + +from sagemaker.core.local.data import MultiRecordStrategy + + +class _ListSplitter: + """A minimal Splitter stand-in that yields pre-set records.""" + + def __init__(self, records): + self._records = records + + def split(self, file): # pylint: disable=unused-argument + for record in self._records: + yield record + + +def test_multi_record_strategy_pad_bytes(): + """Binary records must not raise TypeError from str + bytes concatenation.""" + records = [b"abc", b"def", b"ghi"] + strategy = MultiRecordStrategy(_ListSplitter(records)) + + result = list(strategy.pad("dummy", size=6)) + + assert result == [b"abcdefghi"] + assert all(isinstance(chunk, bytes) for chunk in result) + + +def test_multi_record_strategy_pad_str(): + """Text records continue to be grouped as strings.""" + records = ["abc", "def", "ghi"] + strategy = MultiRecordStrategy(_ListSplitter(records)) + + result = list(strategy.pad("dummy", size=6)) + + assert result == ["abcdefghi"] + assert all(isinstance(chunk, str) for chunk in result) + + +def test_multi_record_strategy_pad_empty(): + """An empty split yields nothing instead of raising.""" + strategy = MultiRecordStrategy(_ListSplitter([])) + + result = list(strategy.pad("dummy", size=6)) + + assert result == [] diff --git a/sagemaker-core/tests/unit/local/test_image.py b/sagemaker-core/tests/unit/local/test_image.py index d7a42d5da1..79940405b4 100644 --- a/sagemaker-core/tests/unit/local/test_image.py +++ b/sagemaker-core/tests/unit/local/test_image.py @@ -621,9 +621,6 @@ def test_process_with_multiple_inputs(self, mock_session): "test-job", ) - @pytest.mark.skip( - reason="Requires sagemaker-serve module which is not installed in sagemaker-core tests" - ) def test_train_with_multiple_channels(self, mock_session): """Test train method with multiple input channels""" with patch( @@ -666,7 +663,7 @@ def test_train_with_multiple_channels(self, mock_session): mock_data_source.return_value = mock_source with patch("os.path.isdir", return_value=False): with patch( - "sagemaker.serve.model_builder.DIR_PARAM_NAME", "sagemaker_program" + "sagemaker.core.local.image.DIR_PARAM_NAME", "sagemaker_program" ): with patch.object( container, @@ -712,9 +709,6 @@ def test_train_with_multiple_channels(self, mock_session): == "/tmp/model.tar.gz" ) - @pytest.mark.skip( - reason="Requires sagemaker-serve module which is not installed in sagemaker-core tests" - ) def test_serve_with_environment_variables(self, mock_session): """Test serve method with environment variables""" with patch( @@ -741,7 +735,7 @@ def test_serve_with_environment_variables(self, mock_session): mock_data_source.return_value = mock_source with patch("os.path.isdir", return_value=False): with patch( - "sagemaker.serve.model_builder.DIR_PARAM_NAME", "sagemaker_program" + "sagemaker.core.local.image.DIR_PARAM_NAME", "sagemaker_program" ): with patch( "sagemaker.core.local.image._ecr_login_if_needed", @@ -873,9 +867,6 @@ def test_write_config_files(self, mock_session): assert mock_write.call_count == 3 # hyperparameters, resourceconfig, inputdataconfig - @pytest.mark.skip( - reason="Requires sagemaker-serve module which is not installed in sagemaker-core tests" - ) def test_prepare_training_volumes_with_local_code(self, mock_session): """Test _prepare_training_volumes with local code directory""" with patch( @@ -899,7 +890,7 @@ def test_prepare_training_volumes_with_local_code(self, mock_session): with patch("os.path.isdir", return_value=False): with patch("os.mkdir"): with patch( - "sagemaker.serve.model_builder.DIR_PARAM_NAME", "sagemaker_program" + "sagemaker.core.local.image.DIR_PARAM_NAME", "sagemaker_program" ): with patch( "sagemaker.core.local.data.get_data_source_instance" diff --git a/sagemaker-core/tests/unit/local/test_local_session.py b/sagemaker-core/tests/unit/local/test_local_session.py index faf7556bf7..91f4e4feaa 100644 --- a/sagemaker-core/tests/unit/local/test_local_session.py +++ b/sagemaker-core/tests/unit/local/test_local_session.py @@ -372,6 +372,104 @@ def test_invoke_endpoint_with_string_body(self, mock_pool_manager_class, mock_ge # String should be encoded to bytes assert isinstance(body, bytes) + @patch("sagemaker.core.local.local_session.get_docker_host") + @patch("urllib3.PoolManager") + def test_invoke_endpoint_response_shape(self, mock_pool_manager_class, mock_get_host): + """Response mirrors the real runtime client (ResponseMetadata + variant).""" + mock_get_host.return_value = "localhost" + + mock_pool = Mock() + mock_response = Mock() + mock_response.status = 200 + mock_response.headers = { + "Content-type": "application/json", + "x-amzn-RequestId": "req-abc-123", + } + mock_pool.request.return_value = mock_response + mock_pool_manager_class.return_value = mock_pool + + client = LocalSagemakerRuntimeClient() + + response = client.invoke_endpoint( + Body=b"test data", + EndpointName="test-endpoint", + Accept="application/json", + TargetVariant="variant1", + ) + + assert response["Body"] == mock_response + assert response["ContentType"] == "application/json" + assert response["InvokedProductionVariant"] == "variant1" + metadata = response["ResponseMetadata"] + assert metadata["HTTPStatusCode"] == 200 + assert metadata["RequestId"] == "req-abc-123" + assert metadata["HTTPHeaders"]["Content-type"] == "application/json" + assert metadata["RetryAttempts"] == 0 + + @patch("sagemaker.core.local.local_session.get_docker_host") + @patch("urllib3.PoolManager") + def test_invoke_endpoint_header_lookup_is_case_insensitive( + self, mock_pool_manager_class, mock_get_host + ): + """Real containers send ``Content-Type``; the lookup must not depend on casing.""" + mock_get_host.return_value = "localhost" + + mock_pool = Mock() + mock_response = Mock() + mock_response.status = 200 + mock_response.headers = { + "Content-Type": "application/octet-stream", + "X-Amzn-Requestid": "req-mixed-case", + } + mock_pool.request.return_value = mock_response + mock_pool_manager_class.return_value = mock_pool + + client = LocalSagemakerRuntimeClient() + response = client.invoke_endpoint( + Body=b"test data", EndpointName="test-endpoint", Accept="application/json" + ) + + assert response["ContentType"] == "application/octet-stream" + assert response["ResponseMetadata"]["RequestId"] == "req-mixed-case" + + @patch("sagemaker.core.local.local_session.get_docker_host") + @patch("urllib3.PoolManager") + def test_invoke_endpoint_default_variant(self, mock_pool_manager_class, mock_get_host): + """Without TargetVariant the response reports the default variant name.""" + mock_get_host.return_value = "localhost" + + mock_pool = Mock() + mock_response = Mock() + mock_response.status = 200 + mock_response.headers = {} + mock_pool.request.return_value = mock_response + mock_pool_manager_class.return_value = mock_pool + + client = LocalSagemakerRuntimeClient() + + response = client.invoke_endpoint(Body=b"data", EndpointName="ep") + + assert response["InvokedProductionVariant"] == "AllTraffic" + assert response["ResponseMetadata"]["RequestId"] == "local-request-id" + + def test_describe_user_profile_passthrough(self): + """describe_user_profile delegates to the real boto SageMaker client.""" + mock_session = Mock() + mock_boto_client = Mock() + mock_boto_client.describe_user_profile.return_value = { + "UserSettings": {"ExecutionRole": "arn:aws:iam::123456789012:role/Studio"} + } + mock_session.boto_session.client.return_value = mock_boto_client + + client = LocalSagemakerClient(sagemaker_session=mock_session) + result = client.describe_user_profile(DomainId="d-123", UserProfileName="user-a") + + mock_session.boto_session.client.assert_called_once_with("sagemaker") + mock_boto_client.describe_user_profile.assert_called_once_with( + DomainId="d-123", UserProfileName="user-a" + ) + assert result["UserSettings"]["ExecutionRole"] == ("arn:aws:iam::123456789012:role/Studio") + class TestLocalSession: """Test cases for LocalSession""" diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index 559ac4179c..2129dfb014 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -2292,3 +2292,52 @@ def test_sparkjar_processor_forwards_instance_preferences(self, mock_session): sagemaker_session=mock_session, ) assert processor.instance_preferences == self._PREFS + + +class TestProcessorLocalModeRole: + def test_role_not_required_for_local_instance(self, mock_session): + """A role is optional when instance_type is local (role is unused locally).""" + processor = Processor( + image_uri="test-image:latest", + instance_count=1, + instance_type="local", + sagemaker_session=mock_session, + ) + assert processor.role is None + assert processor.instance_type == "local" + + def test_role_not_required_for_local_gpu_instance(self, mock_session): + """A role is optional when instance_type is local_gpu.""" + processor = Processor( + image_uri="test-image:latest", + instance_count=1, + instance_type="local_gpu", + sagemaker_session=mock_session, + ) + assert processor.role is None + + def test_role_still_required_for_managed_instance(self, mock_session): + """A missing role still raises for non-local instance types.""" + with pytest.raises( + ValueError, match="An AWS IAM role is required to create a Processing job." + ): + Processor( + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + ) + + def test_request_args_skip_role_expansion_when_no_role(self, mock_session): + """Building the create request must not call expand_role(None) in local mode.""" + processor = Processor( + image_uri="test-image:latest", + instance_count=1, + instance_type="local", + sagemaker_session=mock_session, + ) + processor._current_job_name = "local-job" + processor._normalize_args = lambda *a, **k: None + request = processor._get_process_args(inputs=[], outputs=[], experiment_config=None) + assert request["role_arn"] is None + mock_session.expand_role.assert_not_called() diff --git a/sagemaker-mlops/src/sagemaker/mlops/local/local_pipeline_session.py b/sagemaker-mlops/src/sagemaker/mlops/local/local_pipeline_session.py index 300548700f..2dd15bd14d 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/local/local_pipeline_session.py +++ b/sagemaker-mlops/src/sagemaker/mlops/local/local_pipeline_session.py @@ -47,9 +47,10 @@ def __init__(self, *args, **kwargs): Accepts the same arguments as LocalSession. """ super().__init__(*args, **kwargs) - # Add pipeline storage to the sagemaker_client - if not hasattr(self.sagemaker_client, "_pipelines"): - self.sagemaker_client._pipelines = {} + # Own the local pipeline registry on the session rather than mutating the + # sagemaker_client. Attaching state to the client leaks across sessions that + # share a client and risks colliding with real client attributes. + self._local_pipelines = {} @_telemetry_emitter(Feature.LOCAL_MODE, "local_pipeline_session.create_pipeline") def create_pipeline( @@ -69,7 +70,7 @@ def create_pipeline( pipeline_description=pipeline_description, local_session=self, ) - self.sagemaker_client._pipelines[pipeline.name] = local_pipeline + self._local_pipelines[pipeline.name] = local_pipeline return {"PipelineArn": pipeline.name} def update_pipeline( @@ -84,7 +85,7 @@ def update_pipeline( Returns: Pipeline metadata (PipelineArn) """ - if pipeline.name not in self.sagemaker_client._pipelines: + if pipeline.name not in self._local_pipelines: error_response = { "Error": { "Code": "ResourceNotFound", @@ -92,11 +93,9 @@ def update_pipeline( } } raise ClientError(error_response, "update_pipeline") - self.sagemaker_client._pipelines[pipeline.name].pipeline_description = pipeline_description - self.sagemaker_client._pipelines[pipeline.name].pipeline = pipeline - self.sagemaker_client._pipelines[pipeline.name].last_modified_time = ( - datetime.now().timestamp() - ) + self._local_pipelines[pipeline.name].pipeline_description = pipeline_description + self._local_pipelines[pipeline.name].pipeline = pipeline + self._local_pipelines[pipeline.name].last_modified_time = datetime.now().timestamp() return {"PipelineArn": pipeline.name} def describe_pipeline(self, PipelineName): @@ -108,7 +107,7 @@ def describe_pipeline(self, PipelineName): Returns: Pipeline metadata (PipelineArn, PipelineDefinition, LastModifiedTime, etc) """ - if PipelineName not in self.sagemaker_client._pipelines: + if PipelineName not in self._local_pipelines: error_response = { "Error": { "Code": "ResourceNotFound", @@ -116,7 +115,7 @@ def describe_pipeline(self, PipelineName): } } raise ClientError(error_response, "describe_pipeline") - return self.sagemaker_client._pipelines[PipelineName].describe() + return self._local_pipelines[PipelineName].describe() def delete_pipeline(self, PipelineName): """Delete the local pipeline. @@ -127,8 +126,8 @@ def delete_pipeline(self, PipelineName): Returns: Pipeline metadata (PipelineArn) """ - if PipelineName in self.sagemaker_client._pipelines: - del self.sagemaker_client._pipelines[PipelineName] + if PipelineName in self._local_pipelines: + del self._local_pipelines[PipelineName] return {"PipelineArn": PipelineName} def start_pipeline_execution(self, PipelineName, **kwargs): @@ -144,7 +143,7 @@ def start_pipeline_execution(self, PipelineName, **kwargs): logger.warning("Parallelism configuration is not supported in local mode.") if "SelectiveExecutionConfig" in kwargs: raise ValueError("SelectiveExecutionConfig is not supported in local mode.") - if PipelineName not in self.sagemaker_client._pipelines: + if PipelineName not in self._local_pipelines: error_response = { "Error": { "Code": "ResourceNotFound", @@ -152,4 +151,4 @@ def start_pipeline_execution(self, PipelineName, **kwargs): } } raise ClientError(error_response, "start_pipeline_execution") - return self.sagemaker_client._pipelines[PipelineName].start(**kwargs) + return self._local_pipelines[PipelineName].start(**kwargs) diff --git a/sagemaker-mlops/src/sagemaker/mlops/local/pipeline_entities.py b/sagemaker-mlops/src/sagemaker/mlops/local/pipeline_entities.py index 652c83127d..1f0b5b1140 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/local/pipeline_entities.py +++ b/sagemaker-mlops/src/sagemaker/mlops/local/pipeline_entities.py @@ -17,7 +17,8 @@ import enum import datetime import logging -from uuid import uuid4 +import random +import string from copy import deepcopy from botocore.exceptions import ClientError @@ -25,12 +26,24 @@ logger = logging.getLogger(__name__) +_EXECUTION_ID_LENGTH = 12 +_EXECUTION_ID_ALPHABET = string.ascii_uppercase + string.digits + + +def _generate_execution_id(): + """Generate a service-like local execution id. + + The SageMaker service returns short uppercase alphanumeric execution ids + (for example ``2DRR2511NGO3``). Local mode previously used a 36-char UUID, + which overflowed downstream name-length limits. This mirrors the service + format so local and remote executions behave the same. + """ + return "".join(random.choices(_EXECUTION_ID_ALPHABET, k=_EXECUTION_ID_LENGTH)) + class _LocalPipeline(object): """Class representing a local SageMaker Pipeline""" - _executions = {} - def __init__( self, pipeline, @@ -39,6 +52,7 @@ def __init__( ): from sagemaker.core.local import LocalSession + self._executions = {} self.local_session = local_session or LocalSession() self.pipeline = pipeline self.pipeline_description = pipeline_description @@ -63,7 +77,7 @@ def start(self, **kwargs): """Start a pipeline execution. Returns a _LocalPipelineExecution object.""" from sagemaker.mlops.local.pipeline import LocalPipelineExecutor - execution_id = str(uuid4()) + execution_id = _generate_execution_id() execution = _LocalPipelineExecution( execution_id=execution_id, pipeline=self.pipeline, diff --git a/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py b/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py index 2f11b4f108..44a646f1bd 100644 --- a/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py +++ b/sagemaker-mlops/tests/unit/local/test_local_pipeline_session.py @@ -32,7 +32,7 @@ def mock_pipeline(): def local_session(): def mock_init(self, *args, **kwargs): self.sagemaker_client = Mock() - self.sagemaker_client._pipelines = {} + self._local_pipelines = {} with patch.object(LocalPipelineSession, "__init__", mock_init): session = LocalPipelineSession() @@ -40,31 +40,33 @@ def mock_init(self, *args, **kwargs): def test_local_pipeline_session_init(): - """Test LocalPipelineSession initialization.""" + """Test LocalPipelineSession initialization owns the registry on the session.""" def mock_parent_init(self, *args, **kwargs): - self.sagemaker_client = Mock(spec=[]) # Empty spec means no attributes initially + self.sagemaker_client = Mock(spec=[]) # Empty spec means no attributes with patch("sagemaker.core.local.LocalSession.__init__", mock_parent_init): session = LocalPipelineSession() - # Verify _pipelines attribute is created as a dict - assert hasattr(session.sagemaker_client, "_pipelines") - assert session.sagemaker_client._pipelines == {} + # Registry lives on the session, not on the sagemaker_client. + assert session._local_pipelines == {} + assert not hasattr(session.sagemaker_client, "_pipelines") -def test_local_pipeline_session_init_with_existing_pipelines(): - """Test LocalPipelineSession initialization when _pipelines already exists.""" +def test_local_pipeline_session_registry_isolated_per_session(): + """Two sessions sharing a client do not share pipeline state (issue #5604).""" def mock_parent_init(self, *args, **kwargs): - self.sagemaker_client = Mock() - self.sagemaker_client._pipelines = {"existing": "pipeline"} + self.sagemaker_client = Mock(spec=[]) with patch("sagemaker.core.local.LocalSession.__init__", mock_parent_init): - session = LocalPipelineSession() + session1 = LocalPipelineSession() + session2 = LocalPipelineSession() - # Should not overwrite existing _pipelines - assert session.sagemaker_client._pipelines == {"existing": "pipeline"} + session1._local_pipelines["pipelineA"] = "A" + + assert session2._local_pipelines == {} + assert session1._local_pipelines is not session2._local_pipelines def test_create_pipeline(local_session, mock_pipeline): @@ -81,11 +83,8 @@ def test_create_pipeline(local_session, mock_pipeline): ) assert result == {"PipelineArn": "test-pipeline"} - assert "test-pipeline" in local_session.sagemaker_client._pipelines - assert ( - local_session.sagemaker_client._pipelines["test-pipeline"] - == mock_local_pipeline_instance - ) + assert "test-pipeline" in local_session._local_pipelines + assert local_session._local_pipelines["test-pipeline"] == mock_local_pipeline_instance mock_local_pipeline.assert_called_once_with( pipeline=mock_pipeline, @@ -117,7 +116,7 @@ def test_update_pipeline(local_session, mock_pipeline): mock_local_pipeline.pipeline = Mock() mock_local_pipeline.last_modified_time = 1000.0 - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline new_pipeline = Mock() new_pipeline.name = "test-pipeline" @@ -143,7 +142,7 @@ def test_update_pipeline_not_found(local_session, mock_pipeline): def test_update_pipeline_with_kwargs(local_session, mock_pipeline): """Test update_pipeline ignores extra kwargs.""" mock_local_pipeline = Mock() - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.update_pipeline( local_session, mock_pipeline, "Description", extra_param="ignored" @@ -163,7 +162,7 @@ def test_describe_pipeline(local_session): } ) - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.describe_pipeline(local_session, "test-pipeline") @@ -185,12 +184,12 @@ def test_describe_pipeline_not_found(local_session): def test_delete_pipeline(local_session): """Test delete_pipeline removes pipeline.""" mock_local_pipeline = Mock() - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.delete_pipeline(local_session, "test-pipeline") assert result == {"PipelineArn": "test-pipeline"} - assert "test-pipeline" not in local_session.sagemaker_client._pipelines + assert "test-pipeline" not in local_session._local_pipelines def test_delete_pipeline_not_found(local_session): @@ -206,7 +205,7 @@ def test_start_pipeline_execution(local_session): mock_execution = Mock() mock_local_pipeline.start = Mock(return_value=mock_execution) - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.start_pipeline_execution(local_session, "test-pipeline") @@ -220,7 +219,7 @@ def test_start_pipeline_execution_with_kwargs(local_session): mock_execution = Mock() mock_local_pipeline.start = Mock(return_value=mock_execution) - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.start_pipeline_execution( local_session, @@ -242,7 +241,7 @@ def test_start_pipeline_execution_with_parallelism_config(local_session, caplog) mock_execution = Mock() mock_local_pipeline.start = Mock(return_value=mock_execution) - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline result = LocalPipelineSession.start_pipeline_execution( local_session, "test-pipeline", ParallelismConfiguration={"MaxParallelExecutionSteps": 5} @@ -255,7 +254,7 @@ def test_start_pipeline_execution_with_parallelism_config(local_session, caplog) def test_start_pipeline_execution_with_selective_execution_config(local_session): """Test start_pipeline_execution raises error for selective execution config.""" mock_local_pipeline = Mock() - local_session.sagemaker_client._pipelines["test-pipeline"] = mock_local_pipeline + local_session._local_pipelines["test-pipeline"] = mock_local_pipeline with pytest.raises(ValueError) as exc_info: LocalPipelineSession.start_pipeline_execution( diff --git a/sagemaker-mlops/tests/unit/local/test_pipeline_entities.py b/sagemaker-mlops/tests/unit/local/test_pipeline_entities.py index 44f56bb697..f247f95f4a 100644 --- a/sagemaker-mlops/tests/unit/local/test_pipeline_entities.py +++ b/sagemaker-mlops/tests/unit/local/test_pipeline_entities.py @@ -116,6 +116,45 @@ def test_start_creates_execution(self, mock_pipeline, mock_local_session): mock_executor.assert_called_once() mock_executor_instance.execute.assert_called_once() + def test_executions_isolated_per_instance(self, mock_pipeline, mock_local_session): + """Each _LocalPipeline owns its own execution registry (issue #5572).""" + pipeline1 = _LocalPipeline(pipeline=mock_pipeline, local_session=mock_local_session) + pipeline2 = _LocalPipeline(pipeline=mock_pipeline, local_session=mock_local_session) + + pipeline1._executions["exec1"] = "execution1" + pipeline2._executions["exec2"] = "execution2" + + assert pipeline1._executions == {"exec1": "execution1"} + assert pipeline2._executions == {"exec2": "execution2"} + assert pipeline1._executions is not pipeline2._executions + + def test_start_uses_service_like_execution_id(self, mock_pipeline, mock_local_session): + """Execution ids are 12-char uppercase alphanumerics (issue #5269).""" + import re + + mock_pipeline.steps = [] + mock_pipeline.parameters = [] + + captured = {} + + def fake_execution(execution_id, **kwargs): + captured["execution_id"] = execution_id + return Mock() + + with patch("sagemaker.mlops.local.pipeline.LocalPipelineExecutor") as mock_executor: + mock_executor.return_value.execute = Mock(return_value=Mock()) + with patch( + "sagemaker.mlops.local.pipeline_entities._LocalPipelineExecution", + side_effect=fake_execution, + ): + local_pipeline = _LocalPipeline( + pipeline=mock_pipeline, local_session=mock_local_session + ) + local_pipeline.start() + + execution_id = captured["execution_id"] + assert re.fullmatch(r"[A-Z0-9]{12}", execution_id), execution_id + class TestLocalPipelineExecution: """Tests for _LocalPipelineExecution class."""