diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 58761105e1..11853c8a11 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -1879,6 +1879,42 @@ def _build_for_passthrough(self) -> Model: self.secret_key = "" + model_artifact_uri = None + if self.model_path and self.model_path.startswith("s3://"): + 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 + + has_source_code = bool( + getattr(self, "entry_point", None) and getattr(self, "source_dir", None) + ) + + # Repack source_code into the model artifact so build() produces a + # self-contained model.tar.gz (code under code/). + if has_source_code and model_artifact_uri: + if not ( + isinstance(self.s3_model_data_url, str) + and self.s3_model_data_url.startswith("s3://") + ): + self.s3_model_data_url = model_artifact_uri + self.s3_upload_path = None + + if self.mode in LOCAL_MODES: + self._prepare_for_mode() + + return self._create_model() + + if getattr(self, "entry_point", None) and model_artifact_uri: + # entry_point provided without a source_dir: repack cannot bundle the + # code, so it would be dropped. Warn instead of silently ignoring it. + logger.warning( + "source_code was provided without a source_dir; the inference code " + "will not be repacked into the model artifact. Provide " + "SourceCode(source_dir=...) to bundle custom inference code." + ) + if self.model_path and self.model_path.startswith("s3://"): self.s3_upload_path = self.model_path else: @@ -2282,6 +2318,8 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None: script_name=os.path.basename(self.entry_point), ) + repack_dependencies = self.script_dependencies or [] + logger.info( "Repacking model artifact (%s), script artifact " "(%s), and dependencies (%s) " @@ -2289,14 +2327,14 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None: "This may take some time depending on model size...", self.s3_model_data_url, self.source_dir, - self.dependencies, + repack_dependencies, repacked_model_data, ) repack_model( inference_script=self.entry_point, source_directory=self.source_dir, - dependencies=self.dependencies, + dependencies=repack_dependencies, model_uri=self.s3_model_data_url, repacked_model_uri=repacked_model_data, sagemaker_session=self.sagemaker_session, 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 new file mode 100644 index 0000000000..652997deb5 --- /dev/null +++ b/sagemaker-serve/tests/integ/test_passthrough_source_code_repack_integration.py @@ -0,0 +1,142 @@ +# 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. +"""Integration test for ModelBuilder passthrough source_code repack. + +Covers the fix where an image_uri build with a model artifact and custom +source_code repacks the code into the artifact (instead of silently dropping +it). This calls build() only (no deploy) so it runs in seconds. +""" +from __future__ import absolute_import + +import io +import os +import tarfile +import tempfile +import uuid + +import boto3 +import pytest + +from sagemaker.serve.model_builder import ModelBuilder +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 + +MODEL_NAME_PREFIX = "mb-passthrough-repack" + + +def _upload_raw_artifact(s3_client, bucket, prefix): + """Upload a minimal model.tar.gz containing only a model file (no code/).""" + tmp = tempfile.mkdtemp() + with open(os.path.join(tmp, "model.json"), "w") as f: + f.write('{"weights": [1, 1, 1, 1]}') + tar_path = os.path.join(tmp, "model.tar.gz") + with tarfile.open(tar_path, "w:gz") as t: + t.add(os.path.join(tmp, "model.json"), arcname="model.json") + key = f"{prefix}/model.tar.gz" + s3_client.upload_file(tar_path, bucket, key) + return f"s3://{bucket}/{key}", key + + +def _make_source_code_dir(): + d = tempfile.mkdtemp() + with open(os.path.join(d, "inference.py"), "w") as f: + f.write( + "def model_fn(model_dir):\n return None\n" + "def predict_fn(data, model):\n return data\n" + ) + with open(os.path.join(d, "requirements.txt"), "w") as f: + f.write("joblib\n") + return d + + +def _tar_members(s3_client, s3_uri): + _, _, rest = s3_uri.partition("s3://") + bucket, _, key = rest.partition("/") + body = s3_client.get_object(Bucket=bucket, Key=key)["Body"].read() + return tarfile.open(fileobj=io.BytesIO(body), mode="r:gz").getnames() + + +@pytest.mark.slow_test +def test_build_repacks_source_code_into_artifact(): + """build() with image_uri + model artifact + source_code repacks code/ into + the model.tar.gz. No deploy - runs in seconds.""" + session = Session() + region = session.boto_region_name + bucket = session.default_bucket() + role = get_execution_role(sagemaker_session=session) + s3_client = boto3.client("s3", region_name=region) + + unique_id = uuid.uuid4().hex[:8] + prefix = f"{MODEL_NAME_PREFIX}/{unique_id}" + s3_keys = [] + core_model = None + + try: + 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", + ) + + artifact_uri, key = _upload_raw_artifact(s3_client, bucket, prefix) + s3_keys.append(key) + src_dir = _make_source_code_dir() + + model_builder = ModelBuilder( + image_uri=image_uri, + source_code=SourceCode(source_dir=src_dir, entry_script="inference.py"), + s3_model_data_url=artifact_uri, + role_arn=role, + instance_type="ml.m5.large", + sagemaker_session=session, + ) + model_builder.model_path = f"/tmp/sagemaker/model-builder/{unique_id}" + os.makedirs(model_builder.model_path, exist_ok=True) + model_builder.dependencies = [] + + core_model = model_builder.build( + model_name=f"{MODEL_NAME_PREFIX}-{unique_id}", mode=Mode.SAGEMAKER_ENDPOINT + ) + + # A repack must have produced a new artifact (not the raw one) + repacked = model_builder.repacked_model_data + assert repacked is not None + assert repacked != artifact_uri + s3_keys.append(repacked.partition(f"s3://{bucket}/")[2]) + + # The repacked artifact must contain the inference code under code/ + members = _tar_members(s3_client, repacked) + assert any(m.endswith("code/inference.py") for m in members), members + assert any(m.endswith("model.json") for m in members), members + + # Script-mode env vars must be wired up on the model container + env = core_model.primary_container.environment or {} + assert env.get("SAGEMAKER_PROGRAM") == "inference.py" + assert env.get("SAGEMAKER_SUBMIT_DIRECTORY") == "/opt/ml/model/code" + + finally: + if core_model is not None: + try: + core_model.delete() + except Exception: # noqa + pass + for k in s3_keys: + try: + s3_client.delete_object(Bucket=bucket, Key=k) + except Exception: # noqa + pass diff --git a/sagemaker-serve/tests/unit/test_model_builder_build.py b/sagemaker-serve/tests/unit/test_model_builder_build.py index bfcb3bc19a..be6cdd06de 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_build.py +++ b/sagemaker-serve/tests/unit/test_model_builder_build.py @@ -386,10 +386,33 @@ def test_upload_code_local_mode(self, mock_determine): self.assertIsNone(builder.uploaded_code) - @unittest.skip("Complex mocking required for repack_model with file system operations") - def test_upload_code_with_repack(self): - """Test _upload_code with repacking.""" - pass + @patch("sagemaker.serve.model_builder.repack_model") + @patch("sagemaker.core.workflow.is_pipeline_variable", return_value=False) + @patch("sagemaker.core.s3.determine_bucket_and_prefix") + def test_upload_code_repack_uses_script_dependencies( + self, mock_determine, mock_is_pipeline, mock_repack + ): + """_upload_code(repack=True) must pass script_dependencies (a list) to + repack_model, not the deprecated self.dependencies auto-detect dict.""" + mock_determine.return_value = ("test-bucket", "test-prefix") + + builder = ModelBuilder( + model=Mock(), + role_arn="arn:aws:iam::123456789012:role/TestRole", + sagemaker_session=self.mock_session, + ) + builder.bucket = None + builder.entry_point = "inference.py" + builder.source_dir = "/path/to/code" + builder.script_dependencies = ["/path/to/requirements.txt"] + builder.dependencies = {"auto": True} # deprecated dict, must NOT be used + builder.s3_model_data_url = "s3://bucket/model.tar.gz" + builder.model_kms_key = None + + builder._upload_code("test-prefix", repack=True) + + mock_repack.assert_called_once() + assert mock_repack.call_args.kwargs["dependencies"] == ["/path/to/requirements.txt"] @unittest.skip("Complex file system mocking - os.stat requires real file paths") @patch('sagemaker.core.s3.determine_bucket_and_prefix') diff --git a/sagemaker-serve/tests/unit/test_model_builder_methods.py b/sagemaker-serve/tests/unit/test_model_builder_methods.py index 400eb70047..a58c862502 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.constants import Framework from sagemaker.serve.mode.function_pointers import Mode +from sagemaker.core.training.configs import SourceCode class TestModelBuilderSimpleMethods: @@ -455,6 +456,119 @@ def test_build_for_passthrough_sets_s3_upload_path_none_when_no_model_path( assert builder.s3_upload_path is None +class TestBuildForPassthroughSourceCodeRepack: + """Tests for _build_for_passthrough() repacking source_code into the model artifact. + + When source_code is supplied together with a model artifact, build() must repack + the code into the artifact instead of dropping it. + """ + + def _make_mock_session(self): + mock_session = Mock() + mock_session.boto_region_name = "us-west-2" + mock_session.config = {} + mock_session.boto_session = Mock() + mock_session.boto_session.region_name = "us-west-2" + return mock_session + + def _make_builder(self, **kwargs): + return 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", + source_code=SourceCode(source_dir="./code", entry_script="inference.py"), + sagemaker_session=self._make_mock_session(), + **kwargs, + ) + + @patch.object(ModelBuilder, "_create_model") + def test_repacks_and_preserves_script_mode_vars_with_s3_model_data_url(self, mock_create): + """source_code + s3_model_data_url artifact -> repack path, script-mode vars kept.""" + mock_create.return_value = Mock() + + builder = self._make_builder(s3_model_data_url="s3://bucket/model.tar.gz") + + 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.s3_upload_path is None + mock_create.assert_called_once() + + @patch.object(ModelBuilder, "_create_model") + def test_bridges_model_path_to_s3_model_data_url(self, mock_create): + """When only model_path (S3) is given, it is bridged to s3_model_data_url for repack.""" + mock_create.return_value = Mock() + + builder = self._make_builder() + builder.model_path = "s3://bucket/model.tar.gz" + + builder._build_for_passthrough() + + assert builder.s3_model_data_url == "s3://bucket/model.tar.gz" + assert builder.entry_point == "inference.py" + assert builder.source_dir == "./code" + mock_create.assert_called_once() + + @patch.object(ModelBuilder, "_create_model") + def test_repack_branch_makes_is_repack_true(self, mock_create): + """After the repack branch runs, is_repack() must return True so + _prepare_container_def_base triggers _upload_code(repack=True).""" + mock_create.return_value = Mock() + + builder = self._make_builder(s3_model_data_url="s3://bucket/model.tar.gz") + + builder._build_for_passthrough() + + assert builder.is_repack() is True + + @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).""" + 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="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_entry_point_without_source_dir_warns_and_does_not_repack(self, mock_create): + """entry_script without source_dir cannot be repacked: warn, don't repack.""" + 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", + source_code=SourceCode(entry_script="inference.py"), + s3_model_data_url="s3://bucket/model.tar.gz", + sagemaker_session=self._make_mock_session(), + ) + # sanity: no source_dir -> repack branch is not eligible + assert builder.source_dir is None + assert builder.entry_point == "inference.py" + + with patch("sagemaker.serve.model_builder.logger") as mock_logger: + builder._build_for_passthrough() + mock_logger.warning.assert_called_once() + + # fall-through nulls script-mode vars (no repack occurred) + assert builder.entry_point is None + assert builder.source_dir is None + mock_create.assert_called_once() + + class TestGetDockerClient: """Tests for _get_docker_client() Studio Docker client initialization.