From 681d6b819e0ebb000c3ea66f2293f89d64457f07 Mon Sep 17 00:00:00 2001 From: sagemaker-bot Date: Tue, 29 Sep 2026 23:16:37 +0000 Subject: [PATCH] fix(#5828): ModelBuilder.register does not generate a repack step with a pipeline session in V3 but does in V2 --- .../src/sagemaker/serve/model_builder.py | 10 +- ..._model_builder_register_pipeline_repack.py | 119 ++++++++++++++++++ 2 files changed, 127 insertions(+), 2 deletions(-) create mode 100644 sagemaker-serve/tests/unit/test_model_builder_register_pipeline_repack.py diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 0d1e6e4746..94e5ec6177 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2652,8 +2652,14 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None: if isinstance(self.s3_model_data_url, dict): logging.warning("ModelDataSource currently doesn't support model repacking") return - if is_pipeline_variable(self.s3_model_data_url): - # model is not yet there, defer repacking to later during pipeline execution + if is_pipeline_variable(self.s3_model_data_url) or isinstance( + self.sagemaker_session, PipelineSession + ): + # Under a PipelineSession the repack must be deferred to a runtime + # `_RepackModelStep` inserted by `ModelStep`, whether the model data is a + # Pipeline variable (not yet materialized) or a static S3 URI. Repacking + # eagerly here would bypass that step and leave the registered model without + # the inference code bundled in (V2 parity, see issue #5828). if not isinstance(self.sagemaker_session, PipelineSession): logging.warning( "The model_data is a Pipeline variable of type %s, " diff --git a/sagemaker-serve/tests/unit/test_model_builder_register_pipeline_repack.py b/sagemaker-serve/tests/unit/test_model_builder_register_pipeline_repack.py new file mode 100644 index 0000000000..7d4f169029 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_register_pipeline_repack.py @@ -0,0 +1,119 @@ +# 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/ +"""Reproduction for issue #5828. + +ModelBuilder.register() under a PipelineSession, with custom inference +source_code (entry_script + source_dir) and an S3 model artifact, must flag a +runtime repack so that ``ModelStep`` inserts a ``_RepackModelStep`` (as V2 did). +The repack step bundles the inference code into the model artifact; without it +the registered model package cannot be repacked and is unusable. + +The pipeline-context return value of register() is a ``_ModelStepArguments`` +whose ``need_runtime_repack`` set is consumed by ``ModelStep`` to decide whether +to append the repack step. In V3 this set is left empty, so no repack step is +generated -- this test asserts it is populated (parity with V2). +""" +from __future__ import absolute_import + +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core.training.configs import SourceCode +from sagemaker.core.workflow.pipeline_context import PipelineSession + + +class _FakePipelineSession(PipelineSession): + """A PipelineSession usable without any AWS calls. + + Bypasses the base ``Session.__init__`` (which reaches out to boto) while + keeping the real ``context`` / ``init_model_step_arguments`` / + ``_intercept_create_request`` behavior that ``runnable_by_pipeline`` relies + on. It IS an instance of ``PipelineSession`` so the pipeline code path is + taken. + """ + + def __init__(self): + self._context = None + self._region_name = "us-west-2" + self.config = {} + self.sagemaker_config = {} + self.local_mode = False + self.default_bucket_prefix = "prefix" + self.settings = MagicMock() + self.boto_session = MagicMock() + self.boto_session.region_name = "us-west-2" + self.sagemaker_client = MagicMock() + + def default_bucket(self): + return "test-bucket" + + +@pytest.fixture(autouse=True) +def _ambient_execution_role(): + with patch( + "sagemaker.serve.model_builder.resolve_and_validate_role", + return_value="arn:aws:iam::123456789012:role/ambient-execution-role", + ): + yield + + +def test_register_generates_repack_step_with_pipeline_session(): + from sagemaker.serve import ModelBuilder + + pipeline_session = _FakePipelineSession() + + model_builder = ModelBuilder( + image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/img:latest", + s3_model_data_url="s3://bucket/model.tar.gz", + source_code=SourceCode( + entry_script="infer.py", + source_dir="s3://bucket/source/sourcedir.tar.gz", + ), + role_arn="arn:aws:iam::123456789012:role/ambient-execution-role", + sagemaker_session=pipeline_session, + env_vars={"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "application/json"}, + ) + model_builder.framework = None + model_builder.framework_version = None + + with patch("sagemaker.core.s3.determine_bucket_and_prefix", return_value=("bucket", "prefix")), \ + patch("sagemaker.serve.model_builder.repack_model") as mock_repack, \ + patch( + "sagemaker.serve.model_builder.get_model_package_args", + return_value={}, + ), \ + patch( + "sagemaker.serve.model_builder.create_model_package_from_containers", + return_value={ + "ModelPackageArn": "arn:aws:sagemaker:us-west-2:1:model-package/g/1" + }, + ), \ + patch( + "sagemaker.serve.model_builder.update_container_with_inference_params", + side_effect=lambda **kw: kw.get("container_def"), + ): + step_args = model_builder.register( + content_types=["application/jsonlines"], + response_types=["application/json"], + inference_instances=["ml.m5.large"], + transform_instances=["ml.m5.large"], + model_package_group_name="my-group", + approval_status="Approved", + ) + + # register() under a PipelineSession returns the _ModelStepArguments context. + assert step_args is not None + assert hasattr(step_args, "need_runtime_repack") + + # A repack step MUST be flagged (V2 parity). The bug leaves this empty and + # instead repacks eagerly, so ModelStep never inserts a _RepackModelStep. + assert step_args.need_runtime_repack, ( + "register() under a PipelineSession with source_code did not flag a " + "runtime repack; ModelStep will not generate a repack step (issue #5828)." + )