Introduce changes for InProcess Mode

Author: Bryannah Hernandez

src/sagemaker/serve/mode/in_process_mode.py

@@ -0,0 +1,164 @@
+"""Module that defines the LocalContainerMode class"""
+
+from __future__ import absolute_import
+from pathlib import Path
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, Type
+import base64
+import time
+import subprocess
+import docker
+
+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.logging_agent import pull_logs
+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_INTERVAL_SEC = 5
+
+_PING_HEALTH_CHECK_FAIL_MSG = (
+    "Container did not pass the ping health check. "
+    + "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 local 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.ecr = session.boto_session.client("ecr")
+        self.model_server = model_server
+        self.client = None
+        self.container = None
+        self.secret_key = None
+        self._ping_container = None
+        self._invoke_serving = None
+
+    def load(self, model_path: str = None):
+        """Placeholder docstring"""
+        path = Path(model_path if model_path else self.model_path)
+        if not path.exists():
+            raise Exception("model_path does not exist")
+        if not path.is_dir():
+            raise Exception("model_path is not a valid directory")
+
+        return self.inference_spec.load(str(path))
+
+    def prepare(self):
+        """Placeholder docstring"""
+
+    def create_server(
+        self,
+        image: str,
+        container_timeout_seconds: int,
+        secret_key: str,
+        predictor: PredictorBase,
+        env_vars: Dict[str, str] = None,
+        model_path: str = None,
+    ):
+        """Placeholder docstring"""
+
+        self._pull_image(image=image)
+
+        self.destroy_server()
+
+        logger.info("Waiting for model server %s to start up...", self.model_server)
+
+        if self.model_server == ModelServer.MMS:
+            self._start_serving(
+                client=self.client,
+                image=image,
+                model_path=model_path if model_path else self.model_path,
+                secret_key=secret_key,
+                env_vars=env_vars if env_vars else self.env_vars,
+            )
+            self._ping_container = self._multi_model_server_deep_ping
+
+        # allow some time for container to be ready
+        time.sleep(10)
+
+        log_generator = self.container.logs(follow=True, stream=True)
+        time_limit = datetime.now() + timedelta(seconds=container_timeout_seconds)
+        healthy = False
+        while True:
+            now = datetime.now()
+            final_pull = now > time_limit
+            pull_logs(
+                (x.decode("UTF-8").rstrip() for x in log_generator),
+                log_generator.close,
+                datetime.now() + timedelta(seconds=_PING_HEALTH_CHECK_INTERVAL_SEC),
+                now > time_limit,
+            )
+
+            if final_pull:
+                break
+
+            # allow some time for container to be ready
+            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)
+
+    def destroy_server(self):
+        """Placeholder docstring"""
+        if self.container:
+            try:
+                logger.debug("Stopping currently running container...")
+                self.container.kill()
+            except docker.errors.APIError as exc:
+                if exc.response.status_code < 400 or exc.response.status_code > 499:
+                    raise Exception("Error encountered when cleaning up local container") from exc
+            self.container = None
+
+    def _pull_image(self, image: str):
+        """Placeholder docstring"""
+        try:
+            encoded_token = (
+                self.ecr.get_authorization_token()
+                .get("authorizationData")[0]
+                .get("authorizationToken")
+            )
+            decoded_token = base64.b64decode(encoded_token).decode("utf-8")
+            username, password = decoded_token.split(":")
+            ecr_uri = image.split("/")[0]
+            login_command = ["docker", "login", "-u", username, "-p", password, ecr_uri]
+            subprocess.run(login_command, check=True, capture_output=True)
+        except subprocess.CalledProcessError as e:
+            logger.warning("Unable to login to ecr: %s", e)
+
+        self.client = docker.from_env()
+        try:
+            logger.info("Pulling image %s from repository...", image)
+            self.client.images.pull(image)
+        except docker.errors.NotFound as e:
+            raise ValueError("Could not find remote image to pull") from e

src/sagemaker/serve/model_server/multi_model_server/server.py

@@ -19,6 +19,55 @@
 logger = logging.getLogger(__name__)
 
 
+class InProcessMultiModelServer:
+    """In Process Mode Multi Model server instance"""
+
+    def _start_serving(
+        self,
+        client: object,
+        image: str,
+        model_path: str,
+        secret_key: str,
+        env_vars: dict,
+    ):
+        """Placeholder docstring"""
+        self.container = client.containers.run(
+            image,
+            "serve",
+            network_mode="host",
+            detach=True,
+            auto_remove=True,
+            volumes={
+                Path(model_path).joinpath("code"): {
+                    "bind": MODE_DIR_BINDING,
+                    "mode": "rw",
+                },
+            },
+            environment={
+                "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                "SAGEMAKER_PROGRAM": "inference.py",
+                "SAGEMAKER_SERVE_SECRET_KEY": secret_key,
+                "LOCAL_PYTHON": platform.python_version(),
+                **env_vars,
+            },
+        )
+
+    def _invoke_multi_model_server_serving(self, request: object, content_type: str, accept: str):
+        """Placeholder docstring"""
+        try:
+            response = requests.post(
+                f"http://{get_docker_host()}:8080/invocations",
+                data=request,
+                headers={"Content-Type": content_type, "Accept": accept},
+                timeout=600,
+            )
+            response.raise_for_status()
+            return response.content
+        except Exception as e:
+            raise Exception("Unable to send request to the local container server") from e
+
+        return (True, response)
+
 class LocalMultiModelServer:
     """Local Multi Model server instance"""