Enable private docker registry support for ModelBuilder (#4399)

* Initial commit for image_config support in ModelBuilder

* Add vpc_config to model builder

* fix black format

* fix flake8

* try to fix ut

* try to fix ut

* fix constants import in UT

* fix ut

* fix test_model_builder ut

* fix test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mode

* Add documentation

* fix black-format, flake8 and unnecessary usecase with jumpstart models

---------

Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: jiapinw

src/sagemaker/serve/builder/djl_builder.py

@@ -83,6 +83,8 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self.image_config = None
+        self.vpc_config = None
         self._original_deploy = None
         self.secret_key = None
         self.engine = None
@@ -138,6 +140,8 @@ def _create_djl_model(self) -> Type[Model]:
             "source_dir": code_dir,
             "env": self.env_vars,
             "hf_hub_token": self.env_vars.get("HUGGING_FACE_HUB_TOKEN"),
+            "image_config": self.image_config,
+            "vpc_config": self.vpc_config,
         }
 
         if self.engine == _DjlEngine.DEEPSPEED:

src/sagemaker/serve/builder/model_builder.py

@@ -54,6 +54,7 @@
 from sagemaker.serve.validations.check_image_and_hardware_type import (
     validate_image_uri_and_hardware,
 )
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.huggingface.llm_utils import get_huggingface_model_metadata
 
 logger = logging.getLogger(__name__)
@@ -81,7 +82,6 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
 
             * ``Mode.SAGEMAKER_ENDPOINT``: Launch on a SageMaker endpoint
             * ``Mode.LOCAL_CONTAINER``: Launch locally with a container
-
         shared_libs (List[str]): Any shared libraries you want to bring into
             the model packaging.
         dependencies (Optional[Dict[str, Any]): The dependencies of the model
@@ -122,6 +122,15 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
             ``invoke`` and ``load`` functions.
         image_uri (Optional[str]): The container image uri (which is derived from a
             SageMaker-based container).
+        image_config (dict[str, str] or dict[str, PipelineVariable]): Specifies
+            whether the image of model container is pulled from ECR, or private
+            registry in your VPC. By default it is set to pull model container
+            image from ECR. (default: None).
+        vpc_config ( Optional[Dict[str, List[Union[str, PipelineVariable]]]]):
+                The VpcConfig set on the model (default: None)
+                * 'Subnets' (List[Union[str, PipelineVariable]]): List of subnet ids.
+                * 'SecurityGroupIds' (List[Union[str, PipelineVariable]]]): List of security group
+                ids.
         model_server (Optional[ModelServer]): The model server to which to deploy.
             You need to provide this argument when you specify an ``image_uri``
             in order for model builder to build the artifacts correctly (according
@@ -204,6 +213,23 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
     image_uri: Optional[str] = field(
         default=None, metadata={"help": "Define the container image uri"}
     )
+    image_config: Optional[Dict[str, Union[str, PipelineVariable]]] = field(
+        default=None,
+        metadata={
+            "help": "Specifies whether the image of model container is pulled from ECR,"
+            " or private registry in your VPC. By default it is set to pull model "
+            "container image from ECR. (default: None)."
+        },
+    )
+    vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = field(
+        default=None,
+        metadata={
+            "help": "The VpcConfig set on the model (default: None)."
+            "* 'Subnets' (List[Union[str, PipelineVariable]]): List of subnet ids."
+            "* ''SecurityGroupIds'' (List[Union[str, PipelineVariable]]): List of"
+            " security group ids."
+        },
+    )
     model_server: Optional[ModelServer] = field(
         default=None, metadata={"help": "Define the model server to deploy to."}
     )
@@ -386,6 +412,8 @@ def _create_model(self):
         # TODO: we should create model as per the framework
         self.pysdk_model = Model(
             image_uri=self.image_uri,
+            image_config=self.image_config,
+            vpc_config=self.vpc_config,
             model_data=self.s3_upload_path,
             role=self.serve_settings.role_arn,
             env=self.env_vars,
@@ -543,15 +571,16 @@ def build(
         self,
         mode: Type[Mode] = None,
         role_arn: str = None,
-        sagemaker_session: str = None,
+        sagemaker_session: Optional[Session] = None,
     ) -> Type[Model]:
         """Create a deployable ``Model`` instance with ``ModelBuilder``.
 
         Args:
             mode (Type[Mode], optional): The mode. Defaults to ``None``.
             role_arn (str, optional): The IAM role arn. Defaults to ``None``.
-            sagemaker_session (str, optional): The SageMaker session to use
-              for the execution. Defaults to ``None``.
+            sagemaker_session (Optional[Session]): Session object which manages interactions
+                with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+                function creates one using the default AWS configuration chain.
 
         Returns:
             Type[Model]: A deployable ``Model`` object.
@@ -562,10 +591,7 @@ def build(
             self.mode = mode
         if role_arn:
             self.role_arn = role_arn
-        if sagemaker_session:
-            self.sagemaker_session = sagemaker_session
-        elif not self.sagemaker_session:
-            self.sagemaker_session = Session()
+        self.sagemaker_session = sagemaker_session or Session()
 
         self.sagemaker_session.settings._local_download_dir = self.model_path
 
@@ -607,7 +633,7 @@ def save(
         self,
         save_path: Optional[str] = None,
         s3_path: Optional[str] = None,
-        sagemaker_session: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
         role_arn: Optional[str] = None,
     ) -> Type[Model]:
         """WARNING: This function is expremental and not intended for production use.
@@ -618,7 +644,7 @@ def save(
             save_path (Optional[str]): The path where you want to save resources.
             s3_path (Optional[str]): The path where you want to upload resources.
         """
-        self.sagemaker_session = sagemaker_session if sagemaker_session else Session()
+        self.sagemaker_session = sagemaker_session or Session()
 
         if role_arn:
             self.role_arn = role_arn

src/sagemaker/serve/builder/tgi_builder.py

@@ -76,6 +76,8 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self.image_config = None
+        self.vpc_config = None
         self._original_deploy = None
         self.hf_model_config = None
         self._default_tensor_parallel_degree = None
@@ -134,6 +136,8 @@ def _create_tgi_model(self) -> Type[Model]:
 
         pysdk_model = HuggingFaceModel(
             image_uri=self.image_uri,
+            image_config=self.image_config,
+            vpc_config=self.vpc_config,
             env=self.env_vars,
             role=self.role_arn,
             sagemaker_session=self.sagemaker_session,

src/sagemaker/serve/model_server/triton/triton_builder.py

@@ -413,6 +413,8 @@ def _auto_detect_image_for_triton(self):
     def _create_triton_model(self) -> Type[Model]:
         self.pysdk_model = Model(
             image_uri=self.image_uri,
+            image_config=self.image_config,
+            vpc_config=self.vpc_config,
             model_data=self.s3_upload_path,
             role=self.serve_settings.role_arn,
             env=self.env_vars,

tests/unit/sagemaker/serve/builder/test_djl_builder.py

@@ -30,6 +30,7 @@
     LocalModelInvocationException,
 )
 from sagemaker.serve.utils.predictors import DjlLocalModePredictor
+from tests.unit.sagemaker.serve.constants import MOCK_IMAGE_CONFIG, MOCK_VPC_CONFIG
 
 mock_model_id = "TheBloke/Llama-2-7b-chat-fp16"
 mock_t5_model_id = "google/flan-t5-xxl"
@@ -113,6 +114,8 @@ def test_build_deploy_for_djl_local_container(
             schema_builder=mock_schema_builder,
             mode=Mode.LOCAL_CONTAINER,
             model_server=ModelServer.DJL_SERVING,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
         )
 
         builder._prepare_for_mode = MagicMock()
@@ -132,6 +135,8 @@ def test_build_deploy_for_djl_local_container(
         assert builder._default_max_new_tokens == 256
         assert builder.schema_builder.sample_input["parameters"]["max_new_tokens"] == 256
         assert builder.nb_instance_type == "ml.g5.24xlarge"
+        assert model.image_config == MOCK_IMAGE_CONFIG
+        assert model.vpc_config == MOCK_VPC_CONFIG
         assert "deepspeed" in builder.image_uri
 
         builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
@@ -176,6 +181,8 @@ def test_build_for_djl_local_container_faster_transformer(
             schema_builder=mock_schema_builder,
             mode=Mode.LOCAL_CONTAINER,
             model_server=ModelServer.DJL_SERVING,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
         )
         model = builder.build()
         builder.serve_settings.telemetry_opt_out = True
@@ -185,6 +192,8 @@ def test_build_for_djl_local_container_faster_transformer(
             model.generate_serving_properties()
             == mock_expected_fastertransformer_serving_properties
         )
+        assert model.image_config == MOCK_IMAGE_CONFIG
+        assert model.vpc_config == MOCK_VPC_CONFIG
         assert "fastertransformer" in builder.image_uri
 
     @patch(
@@ -212,11 +221,15 @@ def test_build_for_djl_local_container_deepspeed(
             schema_builder=mock_schema_builder,
             mode=Mode.LOCAL_CONTAINER,
             model_server=ModelServer.DJL_SERVING,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
         )
         model = builder.build()
         builder.serve_settings.telemetry_opt_out = True
 
         assert isinstance(model, DeepSpeedModel)
+        assert model.image_config == MOCK_IMAGE_CONFIG
+        assert model.vpc_config == MOCK_VPC_CONFIG
         assert model.generate_serving_properties() == mock_expected_deepspeed_serving_properties
         assert "deepspeed" in builder.image_uri
 

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -19,6 +19,7 @@
 from sagemaker.serve.builder.model_builder import ModelBuilder
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.types import ModelServer
+from tests.unit.sagemaker.serve.constants import MOCK_IMAGE_CONFIG, MOCK_VPC_CONFIG
 
 schema_builder = MagicMock()
 mock_inference_spec = Mock()
@@ -187,8 +188,10 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
+            and image_config == MOCK_IMAGE_CONFIG
+            and vpc_config == MOCK_VPC_CONFIG
             and model_data == model_data
             and role == mock_role_arn
             and env == ENV_VARS
@@ -205,6 +208,8 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
             model=mock_fw_model,
             model_server=ModelServer.TORCHSERVE,
             image_uri=mock_image_uri,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
         )
         build_result = builder.build(sagemaker_session=mock_session)
 
@@ -286,7 +291,7 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_1p_dlc_image_uri
             and model_data == model_data
             and role == mock_role_arn
@@ -391,7 +396,7 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data == model_data
             and role == mock_role_arn
@@ -487,7 +492,7 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data == model_data
             and role == mock_role_arn
@@ -591,7 +596,7 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data == model_data
             and role == mock_role_arn
@@ -692,7 +697,7 @@ def test_build_happy_path_with_local_container_mode(
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data is None
             and role == mock_role_arn
@@ -809,7 +814,7 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data is None
             and role == mock_role_arn
@@ -951,7 +956,7 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co
 
         mock_model_obj = Mock()
         mock_sdk_model.side_effect = (
-            lambda image_uri, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj
+            lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: mock_model_obj  # noqa E501
             if image_uri == mock_image_uri
             and model_data == model_data
             and role == mock_role_arn

tests/unit/sagemaker/serve/constants.py

@@ -11,3 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+
+
+MOCK_IMAGE_CONFIG = {"RepositoryAccessMode": "Vpc"}
+MOCK_VPC_CONFIG = {"Subnets": ["subnet-1234"], "SecurityGroupIds": ["sg123"]}

tests/unit/sagemaker/serve/model_server/triton/test_triton_builder.py

@@ -20,6 +20,8 @@
 from sagemaker.serve.mode.function_pointers import Mode
 import torch
 
+from tests.unit.sagemaker.serve.constants import MOCK_IMAGE_CONFIG, MOCK_VPC_CONFIG
+
 TRITON_IMAGE = "301217895009.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tritonserver:23.02-py3"
 MODEL_PATH = "/path/to/working/dir"
 S3_UPLOAD_PATH = "s3://path/to/bucket"
@@ -48,6 +50,8 @@ class TritonBuilderTests(TestCase):
     def prepare_triton_builder_for_model(self, triton_builder: Triton) -> Triton:
         triton_builder.model = MOCK_PT_MODEL
         triton_builder.image_uri = TRITON_IMAGE
+        triton_builder.image_config = MOCK_IMAGE_CONFIG
+        triton_builder.vpc_config = MOCK_VPC_CONFIG
         triton_builder.mode = Mode.LOCAL_CONTAINER
         triton_builder.schema_builder = pt_schema_builder
         triton_builder.model_path = MODEL_PATH
@@ -90,6 +94,8 @@ def test_build_for_triton_pt(self, mock_detect_fw, mock_get_gpus, mock_path, moc
 
         mock_model.assert_called_with(
             image_uri=TRITON_IMAGE,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
             model_data=S3_UPLOAD_PATH,
             role=ROLE_ARN,
             env=ENV_VAR,
@@ -122,6 +128,8 @@ def test_build_for_triton_tf(self, mock_detect_fw, mock_get_gpus, mock_path, moc
 
         mock_model.assert_called_with(
             image_uri=TRITON_IMAGE,
+            image_config=MOCK_IMAGE_CONFIG,
+            vpc_config=MOCK_VPC_CONFIG,
             model_data=S3_UPLOAD_PATH,
             role=ROLE_ARN,
             env=ENV_VAR,