From aaf749e8bd1e8828aae8d0113e42a68dbf29b6d4 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Thu, 15 Oct 2026 15:57:14 -0700 Subject: [PATCH] fix: keep source_code for a pipeline-variable model URI and stop repack mutating caller source Two coupled defects on the ModelBuilder pipeline repack path. #6450: ModelBuilder._build_for_passthrough() recognised a model artifact only when s3_model_data_url was a str starting with "s3://". When the artifact comes from an upstream pipeline step the value is a PipelineVariable, which failed that check, so the passthrough branch fell through and cleared source_dir, entry_point, dependencies and git_config. The CreateModel request carried the pipeline expression with an empty container Environment, and no repack happened. _upload_code(repack=True) already defers a PipelineVariable artifact to a runtime repack, so the fix is to treat a PipelineVariable s3_model_data_url as a model artifact and route through that existing machinery. A concrete "s3://" URI keeps its current eager repack behaviour. #6481: _RepackModelStep wrote _repack_model.py and _repack_script_launcher.sh into the caller-provided source_dir -- once in the constructor and again on every `arguments` access -- and for an s3:// source_dir it downloaded the archive, added the files and overwrote the original S3 object. The helper files now go into a copy the SDK owns: a temporary directory for a local source_dir, a new sibling S3 key for an s3:// one. Staging runs once, so repeated serialization neither re-stages nor re-uploads, and the caller's object is read but never overwritten. The two are sequenced together because #6481 is latent today only while #6450 blocks that repack path: fixing #6450 alone would ship code that mutates customer-owned source directories and overwrites their original S3 objects. Fixes #6450 Fixes #6481 --- .../src/sagemaker/mlops/workflow/_utils.py | 132 +++++++++----- .../tests/unit/workflow/test_utils.py | 162 +++++++++++++++++- .../src/sagemaker/serve/model_builder.py | 5 + ...sthrough_source_code_repack_integration.py | 67 ++++++++ .../tests/unit/test_model_builder_methods.py | 65 +++++++ 5 files changed, 384 insertions(+), 47 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py index f421eb851e..7a46ecd0ff 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py @@ -47,6 +47,9 @@ INSTANCE_TYPE = "ml.m5.large" REPACK_SCRIPT = "_repack_model.py" REPACK_SCRIPT_LAUNCHER = "_repack_script_launcher.sh" +# Suffix appended to an ``s3://`` source_dir object name to derive the SDK-owned +# staging key. The caller's original object is never overwritten (issue #6481). +STAGED_SOURCE_DIR_SUFFIX = "-sagemaker-repack.tar.gz" LAUNCH_REPACK_SCRIPT_CMD = """ #!/bin/bash @@ -179,8 +182,11 @@ def __init__( self._entry_point_basename = os.path.basename(self._entry_point) self._source_dir = source_dir self._requirements = requirements + # The SDK-owned copy of the source that the repack helper files are written + # into. ``None`` until the first ``_prepare_for_repacking()`` call. + self._staged_source_dir = None - # Prepare source directory with repack scripts + # Stage an SDK-owned copy of the source with the repack scripts in it self._prepare_for_repacking() # Handle requirements.txt like ModelTrainer @@ -194,7 +200,7 @@ def __init__( from sagemaker.core.training.configs import SourceCode, Compute, Networking source_code = SourceCode( - source_dir=self._source_dir, + source_dir=self._staged_source_dir, entry_script=REPACK_SCRIPT_LAUNCHER, requirements=requirements_file, ) @@ -251,11 +257,31 @@ def __init__( ) def _prepare_for_repacking(self): - """Prepares the source for the model trainer.""" + """Stages an SDK-owned copy of the source for the model trainer. + + The caller-provided ``source_dir`` is never modified: the repack helper + files are written into a copy that the SDK owns -- a temporary directory + for a local ``source_dir``, a new S3 object for an ``s3://`` one. Staging + runs at most once, so repeated ``arguments`` access neither re-stages nor + re-uploads (see issue #6481). + """ + if self._staged_source_dir is not None: + return + + if self._source_dir is not None and self._source_dir.lower().startswith("s3://"): + self._staged_source_dir = self._stage_s3_source_dir() + return + if self._source_dir is None: + # _establish_source_dir already builds a directory that the SDK owns, + # so the helper files can go straight into it. self._establish_source_dir() + self._staged_source_dir = self._source_dir + else: + self._staged_source_dir = tempfile.mkdtemp() + shutil.copytree(self._source_dir, self._staged_source_dir, dirs_exist_ok=True) - self._inject_repack_script_and_launcher() + self._inject_repack_script_and_launcher(self._staged_source_dir) def _establish_source_dir(self): """If the source_dir is None, creates it for the repacking job. @@ -272,10 +298,10 @@ def _establish_source_dir(self): shutil.copy2(self._entry_point, os.path.join(self._source_dir, self._entry_point_basename)) self._entry_point = self._entry_point_basename - def _inject_repack_script_and_launcher(self): + def _inject_repack_script_and_launcher(self, target_dir): """Injects the _repack_model.py script and _repack_script_launcher.sh - into S3 or local source directory. + into a local directory that the SDK owns. Note: The bash file is needed because if not supplied, the SKLearn training job will auto install all dependencies listed in requirements.txt. @@ -283,58 +309,82 @@ def _inject_repack_script_and_launcher(self): since it should just simply repack the model along with other supplied files, e.g. the requirements.txt. - If the source_dir is an S3 path: + It performs the following: + 1) copies the _repack_model.py script into ``target_dir`` + 2) creates the _repack_script_launcher.sh in ``target_dir`` + + Args: + target_dir (str): A local directory owned by the SDK. This is never + the caller-provided ``source_dir`` (see issue #6481). + """ + fname = os.path.join(os.path.dirname(__file__), REPACK_SCRIPT) + shutil.copy2(fname, os.path.join(target_dir, REPACK_SCRIPT)) + # newline="\n" forces LF endings so the bash launcher stays runnable in the + # Linux repack container even when written from a Windows host (see issue #3762). + with open( + os.path.join(target_dir, REPACK_SCRIPT_LAUNCHER), "w", newline="\n" + ) as launcher_file: + launcher_file.write(LAUNCH_REPACK_SCRIPT_CMD) + + def _stage_s3_source_dir(self): + """Stages an ``s3://`` source archive into a new, SDK-owned S3 object. + + It performs the following: 1) downloads the source_dir tar.gz 2) extracts it 3) copies the _repack_model.py script into the extracted directory 4) creates the _repack_script_launcher.sh in the extracted dir 5) rezips the directory - 6) overwrites the S3 source_dir with the new tar.gz + 6) uploads the new tar.gz under a new key, alongside the original - If the source_dir is a local path: - 1) copies the _repack_model.py script into the source dir - 2) creates the _repack_script_launcher.sh in the source dir + The caller's object is read but never overwritten (see issue #6481). The + staging key is derived from the original one, so it stays in the caller's + bucket and prefix and inherits the same access and encryption policies. + + Returns: + str: The S3 URI of the staged source archive. """ - fname = os.path.join(os.path.dirname(__file__), REPACK_SCRIPT) - if self._source_dir.lower().startswith("s3://"): - with tempfile.TemporaryDirectory() as tmp: - targz_contents_dir = os.path.join(tmp, "extracted") + staged_uri = self._staged_s3_source_dir_uri() + with tempfile.TemporaryDirectory() as tmp: + targz_contents_dir = os.path.join(tmp, "extracted") - old_targz_path = os.path.join(tmp, "old.tar.gz") - download_file_from_url(self._source_dir, old_targz_path, self.sagemaker_session) + old_targz_path = os.path.join(tmp, "old.tar.gz") + download_file_from_url(self._source_dir, old_targz_path, self.sagemaker_session) - with tarfile.open(name=old_targz_path, mode="r:gz") as t: - custom_extractall_tarfile(t, targz_contents_dir) + with tarfile.open(name=old_targz_path, mode="r:gz") as t: + custom_extractall_tarfile(t, targz_contents_dir) - shutil.copy2(fname, os.path.join(targz_contents_dir, REPACK_SCRIPT)) - # newline="\n" forces LF endings so the bash launcher stays runnable in the - # Linux repack container even when written from a Windows host (see issue #3762). - with open( - os.path.join(targz_contents_dir, REPACK_SCRIPT_LAUNCHER), "w", newline="\n" - ) as launcher_file: - launcher_file.write(LAUNCH_REPACK_SCRIPT_CMD) + self._inject_repack_script_and_launcher(targz_contents_dir) - new_targz_path = os.path.join(tmp, "new.tar.gz") - with tarfile.open(new_targz_path, mode="w:gz") as t: - t.add(targz_contents_dir, arcname=os.path.sep) + new_targz_path = os.path.join(tmp, "new.tar.gz") + with tarfile.open(new_targz_path, mode="w:gz") as t: + t.add(targz_contents_dir, arcname=os.path.sep) - _save_model(self._source_dir, new_targz_path, self.sagemaker_session, kms_key=None) - else: - shutil.copy2(fname, os.path.join(self._source_dir, REPACK_SCRIPT)) - # newline="\n" forces LF endings so the bash launcher stays runnable in the - # Linux repack container even when written from a Windows host (see issue #3762). - with open( - os.path.join(self._source_dir, REPACK_SCRIPT_LAUNCHER), "w", newline="\n" - ) as launcher_file: - launcher_file.write(LAUNCH_REPACK_SCRIPT_CMD) + _save_model(staged_uri, new_targz_path, self.sagemaker_session, kms_key=None) + + return staged_uri + + def _staged_s3_source_dir_uri(self): + """Derives the staging S3 URI for an ``s3://`` source_dir. + + Returns: + str: The original URI with its object name suffixed, keeping the + caller's bucket and prefix. + """ + prefix, _, object_name = self._source_dir.rpartition("/") + for extension in (".tar.gz", ".tgz"): + if object_name.lower().endswith(extension): + object_name = object_name[: -len(extension)] + break + return "/".join([prefix, object_name + STAGED_SOURCE_DIR_SUFFIX]) @property def arguments(self) -> RequestType: """The arguments dict that are used to call `create_training_job`. - This first prepares the source bundle for repackinglby placing artifacts - in locations which the training container will make available to the - repacking script and then gets the arguments for the training job. + The source bundle is staged once, by the constructor, so repeated access + here neither re-stages the source nor touches the caller's ``source_dir`` + (see issue #6481). """ self._prepare_for_repacking() return super(_RepackModelStep, self).arguments diff --git a/sagemaker-mlops/tests/unit/workflow/test_utils.py b/sagemaker-mlops/tests/unit/workflow/test_utils.py index efd4fa51a2..c5927bd2a3 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_utils.py +++ b/sagemaker-mlops/tests/unit/workflow/test_utils.py @@ -18,7 +18,7 @@ import os import tarfile import tempfile -from unittest.mock import Mock, patch +from unittest.mock import Mock, PropertyMock, patch from sagemaker.mlops.workflow._utils import ( FRAMEWORK_VERSION, @@ -26,6 +26,7 @@ REPACK_SCRIPT, REPACK_SCRIPT_LAUNCHER, LAUNCH_REPACK_SCRIPT_CMD, + STAGED_SOURCE_DIR_SUFFIX, _RepackModelStep, ) @@ -284,7 +285,11 @@ def test_establish_source_dir_creates_temp_dir(self, mock_session, temp_entry_po assert os.path.exists(step._source_dir) def test_inject_repack_script_local_source_dir(self, mock_session, temp_entry_point): - """Test _inject_repack_script_and_launcher with local source_dir.""" + """Test _inject_repack_script_and_launcher with local source_dir. + + The helper files land in the SDK-owned staging copy, not in the + caller-provided ``source_dir`` (see issue #6481). + """ with tempfile.TemporaryDirectory() as temp_dir: with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: with patch("sagemaker.train.ModelTrainer") as mock_trainer: @@ -297,7 +302,7 @@ def test_inject_repack_script_local_source_dir(self, mock_session, temp_entry_po mock_trainer.return_value = mock_trainer_instance mock_super.return_value = None - _RepackModelStep( + step = _RepackModelStep( name="repack-step", sagemaker_session=mock_session, role="arn:aws:iam::123456789012:role/SageMakerRole", @@ -306,13 +311,19 @@ def test_inject_repack_script_local_source_dir(self, mock_session, temp_entry_po source_dir=temp_dir, ) - # Verify repack script and launcher were created - repack_script_path = os.path.join(temp_dir, REPACK_SCRIPT) - launcher_path = os.path.join(temp_dir, REPACK_SCRIPT_LAUNCHER) + # Verify repack script and launcher were created in the staged copy + staged_dir = step._staged_source_dir + assert staged_dir != temp_dir + repack_script_path = os.path.join(staged_dir, REPACK_SCRIPT) + launcher_path = os.path.join(staged_dir, REPACK_SCRIPT_LAUNCHER) assert os.path.exists(repack_script_path) assert os.path.exists(launcher_path) + # ... and that the caller's directory is untouched + assert not os.path.exists(os.path.join(temp_dir, REPACK_SCRIPT)) + assert not os.path.exists(os.path.join(temp_dir, REPACK_SCRIPT_LAUNCHER)) + # Verify launcher content with open(launcher_path, "r") as f: content = f.read() @@ -427,6 +438,145 @@ def fake_download(source_dir, dest_path, session): f"got {launcher_open_calls}" ) + def test_caller_source_dir_not_modified(self, mock_session, temp_entry_point): + """Regression for #6481: the caller's source_dir must never be written to. + + The constructor used to drop ``_repack_model.py`` and + ``_repack_script_launcher.sh`` into the caller-provided directory, and + every ``arguments`` access did it again. + """ + with tempfile.TemporaryDirectory() as temp_dir: + with open(os.path.join(temp_dir, "inference.py"), "w") as fh: + fh.write("def model_fn(model_dir):\n return model_dir\n") + original_files = set(os.listdir(temp_dir)) + + with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: + with patch("sagemaker.train.ModelTrainer") as mock_trainer: + with patch( + "sagemaker.mlops.workflow._utils.TrainingStep.__init__" + ) as mock_super: + mock_retrieve.return_value = "sklearn-image:latest" + mock_trainer_instance = Mock() + mock_trainer_instance.train = Mock(return_value=Mock()) + mock_trainer.return_value = mock_trainer_instance + mock_super.return_value = None + + step = _RepackModelStep( + name="repack-step", + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/SageMakerRole", + model_data="s3://bucket/model.tar.gz", + entry_point=temp_entry_point, + source_dir=temp_dir, + ) + + assert set(os.listdir(temp_dir)) == original_files + + # Serializing the step, repeatedly, must not touch it either + with patch( + "sagemaker.mlops.workflow._utils.TrainingStep.arguments", + new_callable=PropertyMock, + return_value={}, + ): + step.arguments + step.arguments + + assert set(os.listdir(temp_dir)) == original_files + # the staged copy carries both helper files and the caller's code + assert set(os.listdir(step._staged_source_dir)) == original_files | { + REPACK_SCRIPT, + REPACK_SCRIPT_LAUNCHER, + } + + def test_staging_happens_once(self, mock_session, temp_entry_point): + """Regression for #6481: repeated ``arguments`` access must not re-stage.""" + with tempfile.TemporaryDirectory() as temp_dir: + with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: + with patch("sagemaker.train.ModelTrainer") as mock_trainer: + with patch( + "sagemaker.mlops.workflow._utils.TrainingStep.__init__" + ) as mock_super: + mock_retrieve.return_value = "sklearn-image:latest" + mock_trainer_instance = Mock() + mock_trainer_instance.train = Mock(return_value=Mock()) + mock_trainer.return_value = mock_trainer_instance + mock_super.return_value = None + + step = _RepackModelStep( + name="repack-step", + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/SageMakerRole", + model_data="s3://bucket/model.tar.gz", + entry_point=temp_entry_point, + source_dir=temp_dir, + ) + staged_dir = step._staged_source_dir + + with patch( + "sagemaker.mlops.workflow._utils.TrainingStep.arguments", + new_callable=PropertyMock, + return_value={}, + ): + with patch.object( + step, "_inject_repack_script_and_launcher" + ) as mock_inject: + step.arguments + step.arguments + + mock_inject.assert_not_called() + assert step._staged_source_dir == staged_dir + + def test_s3_source_dir_staged_to_new_key(self, mock_session, temp_entry_point): + """Regression for #6481: an s3:// source_dir must not be overwritten. + + The original archive is read, and the archive carrying the repack helper + files is uploaded under a new key next to it. + """ + source_dir_uri = "s3://caller-bucket/caller/prefix/sourcedir.tar.gz" + + def fake_download(source_dir, dest_path, session): + with tempfile.TemporaryDirectory() as seed: + seed_file = os.path.join(seed, "inference.py") + with open(seed_file, "w") as fh: + fh.write("def model_fn(model_dir):\n return model_dir\n") + with tarfile.open(dest_path, "w:gz") as t: + t.add(seed_file, arcname="inference.py") + + with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: + with patch("sagemaker.train.ModelTrainer") as mock_trainer: + with patch("sagemaker.mlops.workflow._utils.TrainingStep.__init__") as mock_super: + with patch( + "sagemaker.mlops.workflow._utils.download_file_from_url", + side_effect=fake_download, + ): + with patch( + "sagemaker.mlops.workflow._utils._save_model" + ) as mock_save_model: + mock_retrieve.return_value = "sklearn-image:latest" + mock_trainer_instance = Mock() + mock_trainer_instance.train = Mock(return_value=Mock()) + mock_trainer.return_value = mock_trainer_instance + mock_super.return_value = None + + step = _RepackModelStep( + name="repack-step", + sagemaker_session=mock_session, + role="arn:aws:iam::123456789012:role/SageMakerRole", + model_data="s3://bucket/model.tar.gz", + entry_point=temp_entry_point, + source_dir=source_dir_uri, + ) + + mock_save_model.assert_called_once() + upload_target = mock_save_model.call_args[0][0] + assert upload_target != source_dir_uri + assert upload_target == ( + "s3://caller-bucket/caller/prefix/sourcedir" + STAGED_SOURCE_DIR_SUFFIX + ) + assert step._staged_source_dir == upload_target + # the caller's URI is still what the repack script is told about + assert step._source_dir == source_dir_uri + def test_properties_returns_parent_properties(self, mock_session, temp_entry_point): """Test properties returns parent class properties.""" with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 10a29fa5f8..9aef5b6c70 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -2093,6 +2093,11 @@ def _build_for_passthrough(self) -> Model: model_artifact_uri = self.model_path elif isinstance(self.s3_model_data_url, str) and self.s3_model_data_url.startswith("s3://"): model_artifact_uri = self.s3_model_data_url + elif is_pipeline_variable(self.s3_model_data_url): + # The artifact does not exist yet: it is produced by an upstream + # pipeline step. ``_upload_code(repack=True)`` defers the repack to a + # ``_RepackModelStep`` created by ``ModelStep``. + model_artifact_uri = self.s3_model_data_url has_source_code = bool( getattr(self, "entry_point", None) and getattr(self, "source_dir", None) diff --git a/sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py b/sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py index 177698b8dd..7879cd732c 100644 --- a/sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py +++ b/sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py @@ -31,6 +31,8 @@ from sagemaker.serve.mode.function_pointers import Mode from sagemaker.core.training.configs import SourceCode from sagemaker.core.helper.session_helper import Session, get_execution_role +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.core.workflow.pipeline_context import PipelineSession MODEL_NAME_PREFIX = "mb-passthrough-repack" @@ -139,3 +141,68 @@ def test_build_repacks_source_code_into_artifact(): s3_client.delete_object(Bucket=bucket, Key=k) except Exception: # noqa pass + + +def test_build_preserves_source_code_with_pipeline_variable_model_uri(): + """build() under a PipelineSession with a pipeline-variable model artifact must + keep source_code and defer the repack to pipeline execution time. + + Regression for issue #6450. The artifact is produced by an upstream step, so + the URI is a PipelineVariable rather than a concrete ``s3://`` string; the + passthrough path used to drop source_code for that case. No deploy, no + pipeline execution, and no eager repack upload - runs in seconds. + + Note: generating the ``_RepackModelStep`` inside ``ModelStep`` from this + recorded request is covered separately by PR #6338. + """ + session = Session() + region = session.boto_region_name + role = get_execution_role(sagemaker_session=session) + pipeline_session = PipelineSession(boto_session=boto3.Session(region_name=region)) + + unique_id = uuid.uuid4().hex[:8] + + from sagemaker.core import image_uris + + image_uri = image_uris.retrieve( + framework="sklearn", + region=region, + version="1.2-1", + instance_type="ml.m5.large", + image_scope="inference", + ) + + src_dir = _make_source_code_dir() + model_data = ParameterString( + name="ModelDataUrl", default_value=f"s3://{session.default_bucket()}/model.tar.gz" + ) + + model_builder = ModelBuilder( + image_uri=image_uri, + source_code=SourceCode(source_dir=src_dir, entry_script="inference.py"), + s3_model_data_url=model_data, + role_arn=role, + instance_type="ml.m5.large", + sagemaker_session=pipeline_session, + ) + model_builder.model_path = f"/tmp/sagemaker/model-builder/{unique_id}" + os.makedirs(model_builder.model_path, exist_ok=True) + model_builder.dependencies = [] + + step_args = model_builder.build( + model_name=f"{MODEL_NAME_PREFIX}-pv-{unique_id}", mode=Mode.SAGEMAKER_ENDPOINT + ) + + # source_code must survive the passthrough path + assert model_builder.entry_point == "inference.py" + assert model_builder.source_dir == src_dir + + # the repack is deferred to pipeline execution, so no new artifact is uploaded now + assert model_builder.repacked_model_data is model_data + assert step_args.need_runtime_repack == {id(model_builder)} + assert step_args.runtime_repack_output_prefix is not None + + # script-mode env vars must be wired up on the recorded container definition + env = step_args.create_model_request["PrimaryContainer"]["Environment"] + assert env.get("SAGEMAKER_PROGRAM") == "inference.py" + assert env.get("SAGEMAKER_SUBMIT_DIRECTORY") == "/opt/ml/model/code" diff --git a/sagemaker-serve/tests/unit/test_model_builder_methods.py b/sagemaker-serve/tests/unit/test_model_builder_methods.py index 0704d22328..0bcd63ca11 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_methods.py +++ b/sagemaker-serve/tests/unit/test_model_builder_methods.py @@ -23,6 +23,7 @@ from sagemaker.core.deserializers import JSONDeserializer, TorchTensorDeserializer from sagemaker.serve.mode.function_pointers import Mode from sagemaker.core.training.configs import SourceCode +from sagemaker.core.workflow.parameters import ParameterString class TestModelBuilderSimpleMethods: @@ -535,6 +536,70 @@ def test_repack_branch_makes_is_repack_true(self, mock_create): assert builder.is_repack() is True + @patch.object(ModelBuilder, "_create_model") + def test_repacks_and_preserves_script_mode_vars_with_pipeline_variable(self, mock_create): + """source_code + PipelineVariable artifact -> repack path, script-mode vars kept. + + Regression for issue #6450: the artifact is produced by an upstream + pipeline step, so s3_model_data_url is a PipelineVariable rather than a + concrete ``s3://`` string. The passthrough path used to drop source_code + entirely for that case. + """ + mock_create.return_value = Mock() + + model_data = ParameterString(name="ModelDataUrl", default_value="s3://bucket/model.tar.gz") + builder = self._make_builder(s3_model_data_url=model_data) + + builder._build_for_passthrough() + + # script-mode vars must NOT be cleared (that would skip repack) + assert builder.entry_point == "inference.py" + assert builder.source_dir == "./code" + assert builder.dependencies is not None + assert builder.s3_upload_path is None + # the pipeline variable is left intact, not stringified or replaced + assert builder.s3_model_data_url is model_data + mock_create.assert_called_once() + + @patch.object(ModelBuilder, "_create_model") + def test_pipeline_variable_repack_branch_makes_is_repack_true(self, mock_create): + """A PipelineVariable artifact must reach _upload_code(repack=True) too. + + Regression for issue #6450. + """ + mock_create.return_value = Mock() + + builder = self._make_builder( + s3_model_data_url=ParameterString( + name="ModelDataUrl", default_value="s3://bucket/model.tar.gz" + ) + ) + + builder._build_for_passthrough() + + assert builder.is_repack() is True + + @patch.object(ModelBuilder, "_create_model") + def test_pipeline_variable_without_source_code_still_clears_vars(self, mock_create): + """No source_code + PipelineVariable artifact -> pure image passthrough (unchanged).""" + mock_create.return_value = Mock() + + builder = ModelBuilder( + image_uri="123456789.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + mode=Mode.SAGEMAKER_ENDPOINT, + role_arn="arn:aws:iam::123456789012:role/TestRole", + s3_model_data_url=ParameterString( + name="ModelDataUrl", default_value="s3://bucket/model.tar.gz" + ), + sagemaker_session=self._make_mock_session(), + ) + + builder._build_for_passthrough() + + assert builder.entry_point is None + assert builder.source_dir is None + mock_create.assert_called_once() + @patch.object(ModelBuilder, "_create_model") def test_pure_image_passthrough_without_source_code_still_clears_vars(self, mock_create): """No source_code -> pure image passthrough clears script-mode vars (unchanged)."""