Merge branch 'master' into processing-job-codeartifact-support

Author: akuma12

requirements/extras/test_requirements.txt

@@ -23,12 +23,7 @@ pyvis==0.2.1
 pandas>=1.3.5,<1.5
 scikit-learn==1.3.0
 cloudpickle==2.2.1
-scipy==1.10.1
-urllib3>=1.26.8,<3.0.0
-docker>=5.0.2,<7.0.0
 PyYAML==6.0
-pyspark==3.3.1
-sagemaker-feature-store-pyspark-3.3
 # TODO find workaround
 xgboost>=1.6.2,<=1.7.6
 pillow>=10.0.1,<=11
@@ -39,4 +34,3 @@ tritonclient[http]<2.37.0
 onnx==1.14.1
 # tf2onnx==1.15.1
 nbformat>=5.9,<6
-accelerate>=0.24.1,<=0.27.0

setup.py

@@ -85,6 +85,8 @@ def read_requirements(filename):
 extras["all"] = [item for group in extras.values() for item in group]
 # Tests specific dependencies (do not need to be included in 'all')
 test_dependencies = read_requirements("requirements/extras/test_requirements.txt")
+# test dependencies are a superset of testing and extra dependencies
+test_dependencies.extend(extras["all"])
 # remove torch and torchvision if python version is not 3.10
 if sys.version_info.minor != 10:
     test_dependencies = [

src/sagemaker/clarify.py

@@ -20,7 +20,6 @@
 
 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
@@ -43,7 +42,11 @@
 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.utils.hardware_detector import (
+    _get_gpu_info,
+    _get_gpu_info_fallback,
+    _total_inference_model_size_mib,
+)
 from sagemaker.serve.detector.image_detector import (
     auto_detect_container,
     _detect_framework_and_version,
@@ -70,11 +73,8 @@
     ModelServer.DJL_SERVING,
 }
 
-MIB_CONVERSION_FACTOR = 0.00000095367431640625
-MEMORY_BUFFER_MULTIPLIER = 1.2  # 20% buffer
-
 
-# pylint: disable=attribute-defined-outside-init
+# pylint: disable=attribute-defined-outside-init, disable=E1101
 @dataclass
 class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
     """Class that builds a deployable model.
@@ -719,39 +719,22 @@ def _schema_builder_init(self, model_task: str):
         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 = self.env_vars.get("dtypes", "float32")
-        parser = estimate_command_parser()
-        args = parser.parse_args([self.model, "--dtypes", dtypes])
-
-        output = gather_data(
-            args
-        )  # "dtype", "Largest Layer", "Total Size Bytes", "Training using Adam"
-
-        if output is None:
-            raise ValueError(f"Could not get Model size for {self.model}")
-
-        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:
+            if (
+                _total_inference_model_size_mib(self.model, self.env_vars.get("dtypes", "float32"))
+                <= 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(),
+                    _total_inference_model_size_mib(
+                        self.model, self.env_vars.get("dtypes", "float32")
+                    ),
                     single_gpu_size_mib,
                 )
                 return True

src/sagemaker/serve/builder/model_builder.py

@@ -208,12 +208,18 @@ def _get_inverse(self, obj):
 
     def __repr__(self):
         """Placeholder docstring"""
+        if hasattr(self, "input_serializer") and hasattr(self, "output_serializer"):
+            return (
+                f"SchemaBuilder(\n"
+                f"input_serializer={self.input_serializer}\n"
+                f"output_serializer={self.output_serializer}\n"
+                f"input_deserializer={self.input_deserializer._deserializer}\n"
+                f"output_deserializer={self.output_deserializer._deserializer})"
+            )
         return (
             f"SchemaBuilder(\n"
-            f"input_serializer={self.input_serializer}\n"
-            f"output_serializer={self.output_serializer}\n"
-            f"input_deserializer={self.input_deserializer._deserializer}\n"
-            f"output_deserializer={self.output_deserializer._deserializer})"
+            f"custom_input_translator={self.custom_input_translator}\n"
+            f"custom_output_translator={self.custom_output_translator}\n"
         )
 
     def generate_marshalling_map(self) -> dict:

src/sagemaker/serve/builder/schema_builder.py

@@ -18,12 +18,18 @@
 
 from botocore.exceptions import ClientError
 
+from accelerate.commands.estimate import estimate_command_parser, gather_data
 from sagemaker import Session
+from sagemaker.model import Model
 from sagemaker import instance_types_gpu_info
 
 logger = logging.getLogger(__name__)
 
 
+MIB_CONVERSION_FACTOR = 0.00000095367431640625
+MEMORY_BUFFER_MULTIPLIER = 1.2  # 20% buffer
+
+
 def _get_gpu_info(instance_type: str, session: Session) -> Tuple[int, int]:
     """Get GPU info for the provided instance
 
@@ -108,3 +114,24 @@ def _format_instance_type(instance_type: str) -> str:
 
     ec2_instance = ".".join(split_instance)
     return ec2_instance
+
+
+def _total_inference_model_size_mib(model: Model, dtype: str) -> int:
+    """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.
+    """
+    args = estimate_command_parser().parse_args([model, "--dtypes", dtype])
+
+    output = gather_data(
+        args
+    )  # "dtype", "Largest Layer", "Total Size Bytes", "Training using Adam"
+
+    if output is None:
+        raise ValueError(f"Could not get Model size for {model}")
+
+    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

src/sagemaker/serve/utils/hardware_detector.py

@@ -189,7 +189,7 @@ def __init__(
         sagemaker_runtime_client=None,
         sagemaker_featurestore_runtime_client=None,
         default_bucket=None,
-        settings=SessionSettings(),
+        settings=None,
         sagemaker_metrics_client=None,
         sagemaker_config: dict = None,
         default_bucket_prefix: str = None,
@@ -260,7 +260,7 @@ def __init__(
         self.resource_group_tagging_client = None
         self._config = None
         self.lambda_client = None
-        self.settings = settings
+        self.settings = settings if settings else SessionSettings()
 
         self._initialize(
             boto_session=boto_session,

src/sagemaker/session.py

@@ -181,7 +181,6 @@ def model_builder(request):
 #                 ), f"{caught_ex} was thrown when running pytorch squeezenet local container test"
 
 
-@pytest.mark.skip(reason="Failing test. Fix is pending.")
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,  # or NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
     reason="The goal of these test are to test the serving components of our feature",
@@ -222,8 +221,10 @@ def test_happy_pytorch_sagemaker_endpoint(
             )
             if caught_ex:
                 logger.exception(caught_ex)
+                ignore_if_worker_dies = "Worker died." in str(caught_ex)
+                # https://github.com/pytorch/serve/issues/3032
                 assert (
-                    False
+                    ignore_if_worker_dies
                 ), f"{caught_ex} was thrown when running pytorch squeezenet sagemaker endpoint test"
 
 

tests/integ/sagemaker/serve/test_serve_pt_happy.py

@@ -0,0 +1,65 @@
+# 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 os
+
+import pytest
+from botocore.config import Config
+
+from tests.integ import DATA_DIR
+from sagemaker import Session, get_execution_role
+
+CUSTOM_S3_OBJECT_KEY_PREFIX = "session-default-prefix"
+
+
+# Create a sagemaker_session in workflow scope to prevent race condition
+# with other tests. Some other tests may change the session `settings`.
+@pytest.fixture(scope="module")
+def sagemaker_session_for_pipeline(
+    sagemaker_client_config,
+    boto_session,
+):
+    sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=10)))
+    sagemaker_client = (
+        boto_session.client("sagemaker", **sagemaker_client_config)
+        if sagemaker_client_config
+        else None
+    )
+
+    return Session(
+        boto_session=boto_session,
+        sagemaker_client=sagemaker_client,
+        sagemaker_config={},
+        default_bucket_prefix=CUSTOM_S3_OBJECT_KEY_PREFIX,
+    )
+
+
+@pytest.fixture(scope="module")
+def smclient(sagemaker_session):
+    return sagemaker_session.boto_session.client("sagemaker")
+
+
+@pytest.fixture(scope="module")
+def role(sagemaker_session_for_pipeline):
+    return get_execution_role(sagemaker_session_for_pipeline)
+
+
+@pytest.fixture(scope="module")
+def region_name(sagemaker_session_for_pipeline):
+    return sagemaker_session_for_pipeline.boto_session.region_name
+
+
+@pytest.fixture(scope="module")
+def script_dir():
+    return os.path.join(DATA_DIR, "sklearn_processing")

tests/integ/sagemaker/workflow/conftest.py

@@ -19,7 +19,6 @@
 
 from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.processing import ProcessingInput
-from sagemaker.session import get_execution_role
 from sagemaker.sklearn.processing import SKLearnProcessor
 from sagemaker.dataset_definition.inputs import DatasetDefinition, AthenaDatasetDefinition
 from sagemaker.workflow.execution_variables import ExecutionVariables
@@ -33,33 +32,13 @@
 from tests.integ import DATA_DIR
 
 
-@pytest.fixture(scope="module")
-def region_name(sagemaker_session):
-    return sagemaker_session.boto_session.region_name
-
-
-@pytest.fixture(scope="module")
-def role(sagemaker_session):
-    return get_execution_role(sagemaker_session)
-
-
-@pytest.fixture(scope="module")
-def script_dir():
-    return os.path.join(DATA_DIR, "sklearn_processing")
-
-
 @pytest.fixture
 def pipeline_name():
     return f"my-pipeline-{int(time.time() * 10**7)}"
 
 
 @pytest.fixture
-def smclient(sagemaker_session):
-    return sagemaker_session.boto_session.client("sagemaker")
-
-
-@pytest.fixture
-def athena_dataset_definition(sagemaker_session):
+def athena_dataset_definition(sagemaker_session_for_pipeline):
     return DatasetDefinition(
         local_path="/opt/ml/processing/input/add",
         data_distribution_type="FullyReplicated",
@@ -69,15 +48,15 @@ def athena_dataset_definition(sagemaker_session):
             database="default",
             work_group="workgroup",
             query_string='SELECT * FROM "default"."s3_test_table_$STAGE_$REGIONUNDERSCORED";',
-            output_s3_uri=f"s3://{sagemaker_session.default_bucket()}/add",
+            output_s3_uri=f"s3://{sagemaker_session_for_pipeline.default_bucket()}/add",
             output_format="JSON",
             output_compression="GZIP",
         ),
     )
 
 
 def test_pipeline_execution_with_default_experiment_config(
-    sagemaker_session,
+    sagemaker_session_for_pipeline,
     smclient,
     role,
     sklearn_latest_version,
@@ -99,7 +78,7 @@ def test_pipeline_execution_with_default_experiment_config(
         instance_type=cpu_instance_type,
         instance_count=instance_count,
         command=["python3"],
-        sagemaker_session=sagemaker_session,
+        sagemaker_session=sagemaker_session_for_pipeline,
         base_job_name="test-sklearn",
     )
 
@@ -113,7 +92,7 @@ def test_pipeline_execution_with_default_experiment_config(
         name=pipeline_name,
         parameters=[instance_count],
         steps=[step_sklearn],
-        sagemaker_session=sagemaker_session,
+        sagemaker_session=sagemaker_session_for_pipeline,
     )
 
     try:
@@ -142,7 +121,7 @@ def test_pipeline_execution_with_default_experiment_config(
 
 
 def test_pipeline_execution_with_custom_experiment_config(
-    sagemaker_session,
+    sagemaker_session_for_pipeline,
     smclient,
     role,
     sklearn_latest_version,
@@ -164,7 +143,7 @@ def test_pipeline_execution_with_custom_experiment_config(
         instance_type=cpu_instance_type,
         instance_count=instance_count,
         command=["python3"],
-        sagemaker_session=sagemaker_session,
+        sagemaker_session=sagemaker_session_for_pipeline,
         base_job_name="test-sklearn",
     )
 
@@ -185,7 +164,7 @@ def test_pipeline_execution_with_custom_experiment_config(
             trial_name=Join(on="-", values=["my-trial", ExecutionVariables.PIPELINE_EXECUTION_ID]),
         ),
         steps=[step_sklearn],
-        sagemaker_session=sagemaker_session,
+        sagemaker_session=sagemaker_session_for_pipeline,
     )
 
     try: