change: Enhance model builder selection logic to include model size

Author: samruds

requirements/extras/huggingface_requirements.txt

@@ -0,0 +1,2 @@
+accelerate>=0.24.1,<=0.27.0
+huggingface_hub>=0.14.0,<=0.18.0

requirements/extras/test_requirements.txt

@@ -39,3 +39,5 @@ tritonclient[http]<2.37.0
 onnx==1.14.1
 # tf2onnx==1.15.1
 nbformat>=5.9,<6
+accelerate>=0.24.1,<=0.27.0
+huggingface_hub>=0.14.0,<=0.18.0

setup.py

@@ -79,6 +79,7 @@ def read_requirements(filename):
     "feature-processor": read_requirements(
         "requirements/extras/feature-processor_requirements.txt"
     ),
+    "huggingface": read_requirements("requirements/extras/huggingface_requirements.txt"),
 }
 # Meta dependency groups
 extras["all"] = [item for group in extras.values() for item in group]

src/sagemaker/serve/builder/model_builder.py

@@ -20,9 +20,11 @@
 
 from pathlib import Path
 
+from accelerate.commands.estimate import estimate_command_parser, gather_data
 from sagemaker import Session
 from sagemaker.model import Model
 from sagemaker.base_predictor import PredictorBase
+from sagemaker.djl_inference import defaults
 from sagemaker.serializers import NumpySerializer, TorchTensorSerializer
 from sagemaker.deserializers import JSONDeserializer, TorchTensorDeserializer
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
@@ -41,6 +43,7 @@
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
+from sagemaker.serve.utils.hardware_detector import _get_gpu_info, _get_gpu_info_fallback
 from sagemaker.serve.detector.image_detector import (
     auto_detect_container,
     _detect_framework_and_version,
@@ -67,6 +70,9 @@
     ModelServer.DJL_SERVING,
 }
 
+MIB_CONVERSION_FACTOR = 0.00000095367431640625
+MEMORY_BUFFER_MULTIPLIER = 1.2  # 20% buffer
+
 
 # pylint: disable=attribute-defined-outside-init
 @dataclass
@@ -569,7 +575,7 @@ def wrapper(*args, **kwargs):
     # It supports two modes of deployment
     # 1/ SageMaker Endpoint
     # 2/ Local launch with container
-    def build(
+    def build(  # pylint: disable=R0911
         self,
         mode: Type[Mode] = None,
         role_arn: str = None,
@@ -607,7 +613,11 @@ def build(
 
         self.serve_settings = self._get_serve_setting()
 
+<<<<<<< Updated upstream
         self._is_custom_image_uri = self.image_uri is not None
+=======
+        self._is_custom_image_uri = self.image_uri is None
+>>>>>>> Stashed changes
 
         if isinstance(self.model, str):
             if self._is_jumpstart_model_id():
@@ -618,13 +628,25 @@ def build(
                 hf_model_md = get_huggingface_model_metadata(
                     self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
                 )
+<<<<<<< Updated upstream
 
                 model_task = hf_model_md.get("pipeline_tag")
                 if self.schema_builder is None and model_task:
                     self._schema_builder_init(model_task)
 
                 if model_task == "text-generation":  # pylint: disable=R1705
                     return self._build_for_tgi()
+=======
+                if hf_model_md.get("pipeline_tag") == "text-generation":  # pylint: disable=R1705
+                    return self._build_for_tgi()
+                elif self._can_fit_on_single_gpu():
+                    return self._build_for_transformers()
+                elif (
+                    self.model in defaults.DEEPSPEED_RECOMMENDED_ARCHITECTURES
+                    or self.model in defaults.FASTER_TRANSFORMER_RECOMMENDED_ARCHITECTURES
+                ):
+                    return self._build_for_djl()
+>>>>>>> Stashed changes
                 else:
                     return self._build_for_transformers()
 
@@ -682,6 +704,7 @@ def validate(self, model_dir: str) -> Type[bool]:
 
         return get_metadata(model_dir)
 
+<<<<<<< Updated upstream
     def _schema_builder_init(self, model_task: str):
         """Initialize the schema builder
 
@@ -696,3 +719,69 @@ def _schema_builder_init(self, model_task: str):
             self.schema_builder = SchemaBuilder(sample_inputs, sample_outputs)
         except ValueError:
             raise TaskNotFoundException(f"Schema builder for {model_task} could not be found.")
+=======
+    def _total_inference_model_size_mib(self):
+        """Calculates the model size from HF accelerate
+
+        This function gets the model size from accelerate. It also adds a
+        padding and converts to size MiB. When performing inference, expect
+        to add up to an additional 20% to the given model size as found by EleutherAI.
+        """
+        dtypes = "float32"
+        try:
+            if self.env_vars.get("dtypes"):
+                dtypes = self.env_vars.get("dtypes")
+
+            parser = estimate_command_parser()
+            args = parser.parse_args([self.model, "--dtypes", dtypes])
+        except ValueError:
+            logging.error("Args specified incorrect for model %s", self.model)
+
+        output = gather_data(
+            args
+        )  # "dtype", "Largest Layer", "Total Size Bytes", "Training using Adam"
+
+        total_memory_size_mib = MEMORY_BUFFER_MULTIPLIER * output[0][2] * MIB_CONVERSION_FACTOR
+        logger.info("Total memory size MIB: %s", total_memory_size_mib)
+        return total_memory_size_mib
+
+    def _can_fit_on_single_gpu(self) -> Type[bool]:
+        """Check if model can fit on a single GPU
+
+        If the size of the model is <= single gpu memory size, returns True else False
+        """
+        try:
+            single_gpu_size_mib = self._try_fetch_gpu_info()
+            if self._total_inference_model_size_mib() <= single_gpu_size_mib:
+                logger.info(
+                    "Total inference model size MIB %s, single GPU size for instance MIB %s",
+                    self._total_inference_model_size_mib(),
+                    single_gpu_size_mib,
+                )
+            return True
+        except ValueError:
+            logger.info("Unable to determine single GPU size for instance %s", self.instance_type)
+            return False
+
+    def _try_fetch_gpu_info(self):
+        """Get GPU info
+
+        This function gets the GPU info or fallback to set the size of a single GPU
+        """
+        try:
+            gpu_info = _get_gpu_info(self.instance_type, self.sagemaker_session)
+            logger.info("GPU info %s for instance %s", gpu_info, self.instance_type)
+            return gpu_info[1] / gpu_info[0]
+        except ValueError:
+            pass
+        try:
+            gpu_fallback = _get_gpu_info_fallback(
+                self.instance_type, self.sagemaker_session.boto_region_name
+            )
+            logger.info("GPU fallback picked up %s", gpu_fallback)
+            return gpu_fallback[1] / gpu_fallback[0]
+        except ValueError:
+            raise ValueError(
+                f"Unable to determine single GPU size for instance: [{self.instance_type}]"
+            )
+>>>>>>> Stashed changes

tests/integ/sagemaker/serve/test_model_builder_gpu.py

@@ -0,0 +1,174 @@
+# 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.
+from __future__ import absolute_import
+
+import pytest
+from sagemaker.serve import Mode
+from sagemaker.serve.builder.model_builder import ModelBuilder
+from sagemaker.serve.builder.schema_builder import SchemaBuilder
+from tests.integ.sagemaker.serve.constants import (
+    HF_DIR,
+    PYTHON_VERSION_IS_NOT_310,
+    SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
+)
+from tests.integ.timeout import timeout
+from tests.integ.utils import cleanup_model_resources
+import logging
+
+logger = logging.getLogger(__name__)
+
+model_id = "bert-base-uncased"
+
+sample_input = {"inputs": "Hello I'm a [MASK] model."}
+
+sample_output = [
+    {
+        "score": 0.10731109976768494,
+        "token": 4827,
+        "token_str": "fashion",
+        "sequence": "hello i'm a fashion model.",
+    },
+    {
+        "score": 0.08774465322494507,
+        "token": 2535,
+        "token_str": "role",
+        "sequence": "hello i'm a role model.",
+    },
+    {
+        "score": 0.05338414013385773,
+        "token": 2047,
+        "token_str": "new",
+        "sequence": "hello i'm a new model.",
+    },
+    {
+        "score": 0.04667224362492561,
+        "token": 3565,
+        "token_str": "super",
+        "sequence": "hello i'm a super model.",
+    },
+    {
+        "score": 0.027096163481473923,
+        "token": 2986,
+        "token_str": "fine",
+        "sequence": "hello i'm a fine model.",
+    },
+]
+
+
+@pytest.fixture
+def model_input():
+    return {"inputs": "The man worked as a [MASK]."}
+
+
+@pytest.fixture
+def model_builder_model_schema_builder():
+    return ModelBuilder(
+        model_path=HF_DIR, model=model_id, schema_builder=SchemaBuilder(sample_input, sample_output)
+    )
+
+
+@pytest.fixture
+def model_builder(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="Testing feature",
+)
+@pytest.mark.parametrize("model_builder", ["model_builder_model_schema_builder"], indirect=True)
+def test_non_text_generation_model_single_GPU(sagemaker_session, model_builder, model_input):
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
+    caught_ex = None
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            model = model_builder.build(role_arn=role_arn, sagemaker_session=sagemaker_session)
+            logger.info("Running in SAGEMAKER_ENDPOINT mode")
+            predictor = model.deploy(
+                mode=Mode.SAGEMAKER_ENDPOINT,
+                instance_type="ml.g4dn.xlarge",
+                initial_instance_count=1,
+            )
+            logger.info("Endpoint successfully deployed.")
+            prediction = predictor.predict(model_input)
+            assert prediction is not None
+
+            endpoint_name = predictor.endpoint_name
+            sagemaker_client = sagemaker_session.boto_session.client("sagemaker")
+            endpoint_config_name = sagemaker_client.describe_endpoint(EndpointName=endpoint_name)[
+                "EndpointConfigName"
+            ]
+            actual_instance_type = sagemaker_client.describe_endpoint_config(
+                EndpointConfigName=endpoint_config_name
+            )["ProductionVariants"][0]["InstanceType"]
+            assert "ml.g4dn.xlarge" == actual_instance_type
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                logger.exception(caught_ex)
+                assert (
+                    False
+                ), f"Exception {caught_ex} was thrown when running model builder single GPU test"
+
+
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="Testing feature",
+)
+@pytest.mark.parametrize("model_builder", ["model_builder_model_schema_builder"], indirect=True)
+def test_non_text_generation_model_multi_GPU(sagemaker_session, model_builder, model_input):
+    iam_client = sagemaker_session.boto_session.client("iam")
+    role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
+    caught_ex = None
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        try:
+            model = model_builder.build(role_arn=role_arn, sagemaker_session=sagemaker_session)
+            logger.info("Running in SAGEMAKER_ENDPOINT mode")
+            predictor = model.deploy(
+                mode=Mode.SAGEMAKER_ENDPOINT,
+                instance_type="ml.g4dn.12xlarge",
+                initial_instance_count=1,
+            )
+            logger.info("Endpoint successfully deployed.")
+            prediction = predictor.predict(model_input)
+            assert prediction is not None
+
+            endpoint_name = predictor.endpoint_name
+            sagemaker_client = sagemaker_session.boto_session.client("sagemaker")
+            endpoint_config_name = sagemaker_client.describe_endpoint(EndpointName=endpoint_name)[
+                "EndpointConfigName"
+            ]
+            actual_instance_type = sagemaker_client.describe_endpoint_config(
+                EndpointConfigName=endpoint_config_name
+            )["ProductionVariants"][0]["InstanceType"]
+            assert "ml.g4dn.12xlarge" == actual_instance_type
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                logger.exception(caught_ex)
+                assert (
+                    False
+                ), f"Exception {caught_ex} was thrown when running model builder multi GPU test"