feat: Pulling in dependencies (in_process mode) using conda environment

src/sagemaker/serve/builder/model_builder.py

@@ -36,6 +36,7 @@
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.mode.sagemaker_endpoint_mode import SageMakerEndpointMode
 from sagemaker.serve.mode.local_container_mode import LocalContainerMode
+from sagemaker.serve.mode.in_process_mode import InProcessMode
 from sagemaker.serve.detector.pickler import save_pkl, save_xgboost
 from sagemaker.serve.builder.serve_settings import _ServeSettings
 from sagemaker.serve.builder.djl_builder import DJL
@@ -410,7 +411,7 @@ def _prepare_for_mode(
             )
             self.env_vars.update(env_vars_sagemaker)
             return self.s3_upload_path, env_vars_sagemaker
-        if self.mode == Mode.LOCAL_CONTAINER:
+        elif self.mode == Mode.LOCAL_CONTAINER:
             # init the LocalContainerMode object
             self.modes[str(Mode.LOCAL_CONTAINER)] = LocalContainerMode(
                 inference_spec=self.inference_spec,
@@ -422,9 +423,22 @@ def _prepare_for_mode(
             )
             self.modes[str(Mode.LOCAL_CONTAINER)].prepare()
             return None
+        elif self.mode == Mode.IN_PROCESS:
+            # init the InProcessMode object
+            self.modes[str(Mode.IN_PROCESS)] = InProcessMode(
+                inference_spec=self.inference_spec,
+                schema_builder=self.schema_builder,
+                session=self.sagemaker_session,
+                model_path=self.model_path,
+                env_vars=self.env_vars,
+                model_server=self.model_server,
+            )
+            self.modes[str(Mode.IN_PROCESS)].prepare()
+            return None
 
         raise ValueError(
-            "Please specify mode in: %s, %s" % (Mode.LOCAL_CONTAINER, Mode.SAGEMAKER_ENDPOINT)
+            "Please specify mode in: %s, %s, %s"
+            % (Mode.LOCAL_CONTAINER, Mode.SAGEMAKER_ENDPOINT, Mode.IN_PROCESS)
         )
 
     def _get_client_translators(self):
@@ -603,10 +617,12 @@ def _overwrite_mode_in_deploy(self, overwrite_mode: str):
             s3_upload_path, env_vars_sagemaker = self._prepare_for_mode()
             self.pysdk_model.model_data = s3_upload_path
             self.pysdk_model.env.update(env_vars_sagemaker)
-
         elif overwrite_mode == Mode.LOCAL_CONTAINER:
             self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
             self._prepare_for_mode()
+        elif overwrite_mode == Mode.IN_PROCESS:
+            self.mode = self.pysdk_model.mode = Mode.IN_PROCESS
+            self._prepare_for_mode()
         else:
             raise ValueError("Mode %s is not supported!" % overwrite_mode)
 
@@ -796,9 +812,10 @@ def _initialize_for_mlflow(self, artifact_path: str) -> None:
         self.dependencies.update({"requirements": mlflow_model_dependency_path})
 
     # Model Builder is a class to build the model for deployment.
-    # It supports two modes of deployment
+    # It supports two* modes of deployment
     # 1/ SageMaker Endpoint
     # 2/ Local launch with container
+    # 3/ In process mode with Transformers server in beta release
     def build(  # pylint: disable=R0911
         self,
         mode: Type[Mode] = None,
@@ -896,8 +913,10 @@ def build(  # pylint: disable=R0911
 
     def _build_validations(self):
         """Validations needed for model server overrides, or auto-detection or fallback"""
-        if self.mode == Mode.IN_PROCESS:
-            raise ValueError("IN_PROCESS mode is not supported yet!")
+        if self.mode == Mode.IN_PROCESS and self.model_server is not ModelServer.MMS:
+            raise ValueError(
+                "IN_PROCESS mode is only supported for MMS/Transformers server in beta release."
+            )
 
         if self.inference_spec and self.model:
             raise ValueError("Can only set one of the following: model, inference_spec.")

src/sagemaker/serve/builder/requirements_manager.py

@@ -0,0 +1,98 @@
+# 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.
+"""Requirements Manager class to pull in client dependencies from a .txt or .yml file"""
+from __future__ import absolute_import
+import logging
+import os
+import subprocess
+
+logger = logging.getLogger(__name__)
+
+
+class RequirementsManager:
+    """Transformers build logic with ModelBuilder()"""
+
+    def detect_file_exists(self, dependencies: str = None) -> str:
+        """Creates snapshot of the user's environment
+
+        If a req.txt or conda.yml file is provided, it verifies their existence and
+        returns the local file path
+
+        Args:
+            dependencies (str): Local path where dependencies file exists.
+
+        Returns:
+            file path of the existing or generated dependencies file
+        """
+        dependencies = self._capture_from_local_runtime()
+
+        # Dependencies specified as either req.txt or conda_env.yml
+        if dependencies.endswith(".txt"):
+            self._install_requirements_txt()
+        elif dependencies.endswith(".yml"):
+            self._update_conda_env_in_path()
+        else:
+            raise ValueError(f'Invalid dependencies provided: "{dependencies}"')
+
+    def _install_requirements_txt(self):
+        """Install requirements.txt file using pip"""
+        logger.info("Running command to pip install")
+        subprocess.run("pip install -r requirements.txt", shell=True, check=True)
+        logger.info("Command ran successfully")
+
+    def _update_conda_env_in_path(self):
+        """Update conda env using conda yml file"""
+        logger.info("Updating conda env")
+        subprocess.run("conda env update -f conda_in_process.yml", shell=True, check=True)
+        logger.info("Conda env updated successfully")
+
+    def _get_active_conda_env_name(self) -> str:
+        """Returns the conda environment name from the set environment variable. None otherwise."""
+        return os.getenv("CONDA_DEFAULT_ENV")
+
+    def _get_active_conda_env_prefix(self) -> str:
+        """Returns the conda prefix from the set environment variable. None otherwise."""
+        return os.getenv("CONDA_PREFIX")
+
+    def _capture_from_local_runtime(self) -> str:
+        """Generates dependencies list from the user's local runtime.
+
+        Raises RuntimeEnvironmentError if not able to.
+
+        Currently supports: conda environments
+        """
+
+        # Try to capture dependencies from the conda environment, if any.
+        conda_env_name = self._get_active_conda_env_name()
+        logger.info("Found conda_env_name: '%s'", conda_env_name)
+        conda_env_prefix = None
+
+        if conda_env_name is None:
+            conda_env_prefix = self._get_active_conda_env_prefix()
+
+        if conda_env_name is None and conda_env_prefix is None:
+            raise ValueError("No conda environment seems to be active.")
+
+        if conda_env_name == "base":
+            logger.warning(
+                "We recommend using an environment other than base to "
+                "isolate your project dependencies from conda dependencies"
+            )
+
+        local_dependencies_path = os.path.join(os.getcwd(), "inf_env_snapshot.yml")
+
+        return local_dependencies_path
+
+
+if __name__ == "__main__":
+    RequirementsManager().detect_file_exists(dependencies="auto_capture")

src/sagemaker/serve/builder/transformers_builder.py

@@ -17,6 +17,7 @@
 from abc import ABC, abstractmethod
 from typing import Type
 from pathlib import Path
+import subprocess
 from packaging.version import Version
 
 from sagemaker.model import Model
@@ -35,15 +36,21 @@
 )
 from sagemaker.serve.detector.pickler import save_pkl
 from sagemaker.serve.utils.optimize_utils import _is_optimized
-from sagemaker.serve.utils.predictors import TransformersLocalModePredictor
+from sagemaker.serve.utils.predictors import (
+    TransformersLocalModePredictor,
+    TransformersInProcessModePredictor,
+)
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.huggingface.llm_utils import get_huggingface_model_metadata
+from sagemaker.serve.builder.requirements_manager import RequirementsManager
+
 
 logger = logging.getLogger(__name__)
 DEFAULT_TIMEOUT = 1800
+LOCAL_MODES = [Mode.LOCAL_CONTAINER, Mode.IN_PROCESS]
 
 
 """Retrieves images for different libraries - Pytorch, TensorFlow from HuggingFace hub
@@ -227,6 +234,22 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
             )
             return predictor
 
+        if self.mode == Mode.IN_PROCESS:
+            timeout = kwargs.get("model_data_download_timeout")
+
+            predictor = TransformersInProcessModePredictor(
+                self.modes[str(Mode.IN_PROCESS)], serializer, deserializer
+            )
+
+            self.modes[str(Mode.IN_PROCESS)].create_server(
+                self.image_uri,
+                timeout if timeout else DEFAULT_TIMEOUT,
+                None,
+                predictor,
+                self.pysdk_model.env,
+            )
+            return predictor
+
         if "mode" in kwargs:
             del kwargs["mode"]
         if "role" in kwargs:
@@ -274,7 +297,7 @@ def _build_transformers_env(self):
 
         self.pysdk_model = self._create_transformers_model()
 
-        if self.mode == Mode.LOCAL_CONTAINER:
+        if self.mode in LOCAL_MODES:
             self._prepare_for_mode()
 
         return self.pysdk_model
@@ -358,6 +381,9 @@ def _build_for_transformers(self):
             save_pkl(code_path, (self.inference_spec, self.schema_builder))
             logger.info("PKL file saved to file: %s", code_path)
 
+            if self.mode == Mode.IN_PROCESS:
+                self._create_conda_env()
+
             self._auto_detect_container()
 
             self.secret_key = prepare_for_mms(
@@ -376,3 +402,11 @@ def _build_for_transformers(self):
         if self.sagemaker_session:
             self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model
+
+    def _create_conda_env(self):
+        """Creating conda environment by running commands"""
+
+        try:
+            RequirementsManager().detect_file_exists(self)
+        except subprocess.CalledProcessError:
+            print("Failed to create and activate conda environment.")

src/sagemaker/serve/mode/in_process_mode.py

@@ -0,0 +1,89 @@
+"""Module that defines the InProcessMode class"""
+
+from __future__ import absolute_import
+from pathlib import Path
+import logging
+from typing import Dict, Type
+import time
+from datetime import datetime, timedelta
+
+from sagemaker.base_predictor import PredictorBase
+from sagemaker.serve.spec.inference_spec import InferenceSpec
+from sagemaker.serve.builder.schema_builder import SchemaBuilder
+from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.utils.exceptions import LocalDeepPingException
+from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
+from sagemaker.session import Session
+
+logger = logging.getLogger(__name__)
+
+_PING_HEALTH_CHECK_FAIL_MSG = (
+    "Ping health check did not pass. "
+    + "Please increase container_timeout_seconds or review your inference code."
+)
+
+
+class InProcessMode(
+    InProcessMultiModelServer,
+):
+    """A class that holds methods to deploy model to a container in process environment"""
+
+    def __init__(
+        self,
+        model_server: ModelServer,
+        inference_spec: Type[InferenceSpec],
+        schema_builder: Type[SchemaBuilder],
+        session: Session,
+        model_path: str = None,
+        env_vars: Dict = None,
+    ):
+        # pylint: disable=bad-super-call
+        super().__init__()
+
+        self.inference_spec = inference_spec
+        self.model_path = model_path
+        self.env_vars = env_vars
+        self.session = session
+        self.schema_builder = schema_builder
+        self.model_server = model_server
+        self._ping_container = None
+
+    def load(self, model_path: str = None):
+        """Loads model path, checks that path exists"""
+        path = Path(model_path if model_path else self.model_path)
+        if not path.exists():
+            raise ValueError("model_path does not exist")
+        if not path.is_dir():
+            raise ValueError("model_path is not a valid directory")
+
+        return self.inference_spec.load(str(path))
+
+    def prepare(self):
+        """Prepares the server"""
+
+    def create_server(
+        self,
+        predictor: PredictorBase,
+    ):
+        """Creating the server and checking ping health."""
+        logger.info("Waiting for model server %s to start up...", self.model_server)
+
+        if self.model_server == ModelServer.MMS:
+            self._ping_container = self._multi_model_server_deep_ping
+
+        time_limit = datetime.now() + timedelta(seconds=5)
+        while self._ping_container is not None:
+            final_pull = datetime.now() > time_limit
+
+            if final_pull:
+                break
+
+            time.sleep(10)
+
+            healthy, response = self._ping_container(predictor)
+            if healthy:
+                logger.debug("Ping health check has passed. Returned %s", str(response))
+                break
+
+        if not healthy:
+            raise LocalDeepPingException(_PING_HEALTH_CHECK_FAIL_MSG)