snowflake-ml-python 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

snowflake/ml/_internal/env_utils.py

@@ -310,7 +310,8 @@ def validate_requirements_in_snowflake_conda_channel(
                 FROM information_schema.packages
                 WHERE ({pkg_names_str})
                 AND language = 'python'
-                AND runtime_version = '{parsed_python_version.major}.{parsed_python_version.minor}';
+                AND (runtime_version = '{parsed_python_version.major}.{parsed_python_version.minor}'
+                    OR runtime_version is null);
                 """
             )
         else:

snowflake/ml/_internal/file_utils.py

@@ -59,10 +59,12 @@ def zip_file_or_directory_to_stream(
     Raises:
         FileNotFoundError: Raised when the given path does not exist.
         ValueError: Raised when the leading path is not a actual leading path of path
+        ValueError: Raised when the arcname cannot be encoded using ASCII.
 
     Yields:
         A bytes IO stream containing the zip file.
     """
+    # TODO(SNOW-862576): Should remove check on ASCII encoding after SNOW-862576 fixed.
     if not os.path.exists(path):
         raise FileNotFoundError(f"{path} is not found")
     if leading_path and not path.startswith(leading_path):
@@ -76,23 +78,35 @@ def zip_file_or_directory_to_stream(
             if os.path.realpath(path) != os.path.realpath(start_path):
                 cur_path = os.path.dirname(path)
                 while os.path.realpath(cur_path) != os.path.realpath(start_path):
-                    zf.writestr(f"{os.path.relpath(cur_path, start_path)}/", "")
+                    arcname = os.path.relpath(cur_path, start_path)
+                    if not _able_ascii_encode(arcname):
+                        raise ValueError(f"File name {arcname} cannot be encoded using ASCII. Please rename.")
+                    zf.write(cur_path, arcname)
                     cur_path = os.path.dirname(cur_path)
 
             if os.path.isdir(path):
-                for dirname, _, files in os.walk(path):
+                for dirpath, _, files in os.walk(path):
                     # ignore __pycache__
-                    if ignore_generated_py_file and "__pycache__" in dirname:
+                    if ignore_generated_py_file and "__pycache__" in dirpath:
                         continue
-                    zf.write(dirname, os.path.relpath(dirname, start_path))
+                    arcname = os.path.relpath(dirpath, start_path)
+                    if not _able_ascii_encode(arcname):
+                        raise ValueError(f"File name {arcname} cannot be encoded using ASCII. Please rename.")
+                    zf.write(dirpath, arcname)
                     for file in files:
                         # ignore generated python files
                         if ignore_generated_py_file and file.endswith(GENERATED_PY_FILE_EXT):
                             continue
-                        filename = os.path.join(dirname, file)
-                        zf.write(filename, os.path.relpath(filename, start_path))
+                        file_path = os.path.join(dirpath, file)
+                        arcname = os.path.relpath(file_path, start_path)
+                        if not _able_ascii_encode(arcname):
+                            raise ValueError(f"File name {arcname} cannot be encoded using ASCII. Please rename.")
+                        zf.write(file_path, arcname)
             else:
-                zf.write(path, os.path.relpath(path, start_path))
+                arcname = os.path.relpath(path, start_path)
+                if not _able_ascii_encode(arcname):
+                    raise ValueError(f"File name {arcname} cannot be encoded using ASCII. Please rename.")
+                zf.write(path, arcname)
 
         yield input_stream
 
@@ -145,3 +159,11 @@ def get_all_modules(dirname: str, prefix: str = "") -> List[pkgutil.ModuleInfo]:
     for dirname in subdirs:
         modules.extend(get_all_modules(dirname, prefix=f"{prefix}.{dirname}" if prefix else dirname))
     return modules
+
+
+def _able_ascii_encode(s: str) -> bool:
+    try:
+        s.encode("ascii", errors="strict")
+        return True
+    except UnicodeEncodeError:
+        return False

snowflake/ml/_internal/telemetry.py

@@ -6,7 +6,6 @@ import enum
 import functools
 import inspect
 import operator
-import threading
 import types
 from typing import (
     Any,
@@ -29,7 +28,6 @@ from snowflake.ml._internal import env
 from snowflake.snowpark import dataframe, exceptions, session
 from snowflake.snowpark._internal import utils
 
-_rlock = threading.RLock()
 _log_counter = 0
 _FLUSH_SIZE = 10
 
@@ -308,12 +306,11 @@ def send_api_usage_telemetry(
                 return res
             finally:
                 telemetry.send_function_usage_telemetry(**telemetry_args)
-                with _rlock:
-                    global _log_counter
-                    _log_counter += 1
-                    if _log_counter >= _FLUSH_SIZE or "error" in telemetry_args:
-                        telemetry.send_batch()
-                        _log_counter = 0
+                global _log_counter
+                _log_counter += 1
+                if _log_counter >= _FLUSH_SIZE or "error" in telemetry_args:
+                    telemetry.send_batch()
+                    _log_counter = 0
 
         return cast(Callable[_Args, _ReturnValue], wrap)
 

snowflake/ml/_internal/utils/uri.py

@@ -1,4 +1,4 @@
-import os.path
+import posixpath
 from typing import Optional
 from urllib.parse import ParseResult, urlparse, urlunparse
 
@@ -35,7 +35,12 @@ def get_snowflake_stage_path_from_uri(uri: str) -> Optional[str]:
     if not is_snowflake_stage_uri(uri):
         return None
     uri_components = urlparse(uri)
-    return os.path.join(uri_components.netloc.strip("/"), uri_components.path.strip("/")).rstrip("/")
+    # posixpath.join will drop other components if any of arguments is absolute path.
+    # The path we get is actually absolute (starting with '/'), however, since we concat them to stage location,
+    # it should not.
+    return posixpath.normpath(
+        posixpath.join(posixpath.normpath(uri_components.netloc), posixpath.normpath(uri_components.path.lstrip("/")))
+    )
 
 
 def get_uri_scheme(uri: str) -> str:

snowflake/ml/model/_deploy_client/image_builds/base_image_builder.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+
+
+class ImageBuilder(ABC):
+    """
+    Abstract class encapsulating image building and upload to model registry.
+    """
+
+    @abstractmethod
+    def build_and_upload_image(self) -> str:
+        """Builds and uploads an image to the model registry.
+
+        Returns: Full image path.
+        """
+        pass

snowflake/ml/model/_deploy_client/image_builds/client_image_builder.py

@@ -0,0 +1,259 @@
+import base64
+import json
+import logging
+import os
+import posixpath
+import subprocess
+import tempfile
+import zipfile
+from enum import Enum
+from typing import List
+
+import yaml
+
+from snowflake import snowpark
+from snowflake.ml._internal.utils import query_result_checker
+from snowflake.ml.model._deploy_client.image_builds import (
+    base_image_builder,
+    docker_context,
+)
+from snowflake.ml.model._deploy_client.utils import constants
+
+
+class Platform(Enum):
+    LINUX_AMD64 = "linux/amd64"
+
+
+class ClientImageBuilder(base_image_builder.ImageBuilder):
+    """
+    Client-side image building and upload to model registry.
+
+    Usage requirements:
+    Requires prior installation and running of Docker with BuildKit. See installation instructions in
+        https://docs.docker.com/engine/install/
+
+
+    """
+
+    def __init__(
+        self, *, id: str, image_repo: str, model_zip_stage_path: str, session: snowpark.Session, use_gpu: bool = False
+    ) -> None:
+        """Initialization
+
+        Args:
+            id: A hexadecimal string used for naming the image tag.
+            image_repo: Path to image repository.
+            model_zip_stage_path: Path to model zip file in stage.
+            use_gpu: Boolean flag for generating the CPU or GPU base image.
+            session: Snowpark session
+        """
+        self.image_tag = "/".join([image_repo.rstrip("/"), id]) + ":latest"
+        self.image_repo = image_repo
+        self.model_zip_stage_path = model_zip_stage_path
+        self.use_gpu = use_gpu
+        self.session = session
+
+    def build_and_upload_image(self) -> str:
+        """
+        Builds and uploads an image to the model registry.
+        """
+
+        def _setup_docker_config(docker_config_dir: str) -> None:
+            """Set up a temporary docker config, which is used for running all docker commands.
+
+            Args:
+                docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
+            """
+            ctx = self.session._conn._conn
+            assert ctx._rest, "SnowflakeRestful is not set in session"
+            token_data = ctx._rest._token_request("ISSUE")
+            snowpark_session_token = token_data["data"]["sessionToken"]
+            token_obj = {"token": snowpark_session_token}
+            credentials = f"0sessiontoken:{json.dumps(token_obj)}"
+            encoded_credentials = base64.b64encode(credentials.encode("utf-8")).decode("utf-8")
+            content = {"auths": {self.image_tag: {"auth": encoded_credentials}}}
+            config_path = os.path.join(docker_config_dir, "config.json")
+            with open(config_path, "w", encoding="utf-8") as file:
+                json.dump(content, file)
+
+        self.validate_docker_client_env()
+
+        query_result = (
+            query_result_checker.SqlResultValidator(
+                self.session,
+                query="SHOW PARAMETERS LIKE 'PYTHON_CONNECTOR_QUERY_RESULT_FORMAT' IN SESSION",
+            )
+            .has_dimensions(expected_rows=1)
+            .validate()
+        )
+        prev_format = query_result[0].value
+
+        with tempfile.TemporaryDirectory() as config_dir:
+            try:
+                # Workaround for SNOW-841699: Fail to authenticate to image registry with session token generated from
+                # Snowpark. Need to temporarily set the json query format in order to process GS token response.
+                self.session.sql("ALTER SESSION SET PYTHON_CONNECTOR_QUERY_RESULT_FORMAT = 'json'").collect()
+                _setup_docker_config(config_dir)
+                self._build(config_dir)
+                self._upload(config_dir)
+            finally:
+                self.session.sql(f"ALTER SESSION SET PYTHON_CONNECTOR_QUERY_RESULT_FORMAT = '{prev_format}'").collect()
+                commands = ["docker", "--config", config_dir, "rmi", self.image_tag]
+                logging.info(f"Removing local image: {self.image_tag}")
+                self._run_docker_commands(commands)
+        return self.image_tag
+
+    def validate_docker_client_env(self) -> None:
+        """Ensure docker client is running and BuildKit is enabled. Note that Buildx always uses BuildKit.
+        - Ensure docker daemon is running through the "docker info" command on shell. When docker daemon is running,
+        return code will be 0, else return code will be 1.
+        - Ensure BuildKit is enabled by checking "docker buildx version".
+
+        Raises:
+            ConnectionError: Occurs when Docker is not installed or is not running.
+
+        """
+        info_command = "docker info"
+        buildx_command = "docker buildx version"
+
+        try:
+            subprocess.check_call(info_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True)
+        except subprocess.CalledProcessError:
+            raise ConnectionError("Failed to initialize Docker client. Please ensure Docker is installed and running.")
+
+        try:
+            subprocess.check_call(buildx_command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True)
+        except subprocess.CalledProcessError:
+            raise ConnectionError(
+                "Please ensured Docker is installed with BuildKit by following "
+                "https://docs.docker.com/build/buildkit/#getting-started"
+            )
+
+    def _extract_model_zip(self, context_dir: str) -> str:
+        """Extract a zip file into the specified directory.
+
+        Args:
+            context_dir: Directory to extract the zip to.
+
+        Returns:
+            The extracted model directory.
+        """
+
+        local_model_zip_path = os.path.join(context_dir, posixpath.basename(self.model_zip_stage_path))
+        if zipfile.is_zipfile(local_model_zip_path):
+            extracted_model_dir = os.path.join(context_dir, constants.MODEL_DIR)
+            with zipfile.ZipFile(local_model_zip_path, "r") as model_zip:
+                if len(model_zip.namelist()) > 1:
+                    model_zip.extractall(extracted_model_dir)
+            conda_path = os.path.join(extracted_model_dir, "env", "conda.yaml")
+
+            def remove_snowml_from_conda() -> None:
+                with open(conda_path, encoding="utf-8") as file:
+                    conda_yaml = yaml.safe_load(file)
+
+                dependencies = conda_yaml["dependencies"]
+                dependencies = [dep for dep in dependencies if not dep.startswith("snowflake-ml-python")]
+
+                conda_yaml["dependencies"] = dependencies
+
+                with open(conda_path, "w", encoding="utf-8") as file:
+                    yaml.dump(conda_yaml, file)
+
+            # TODO(shchen): Remove once SNOW-840411 is landed.
+            remove_snowml_from_conda()
+        return extracted_model_dir
+
+    def _build(self, docker_config_dir: str) -> None:
+        """Constructs the Docker context directory and then builds a Docker image based on that context.
+
+        Args:
+            docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
+        """
+
+        with tempfile.TemporaryDirectory() as context_dir:
+            # Download the model zip file that is already uploaded to stage during model registry log_model step.
+            # This is needed in order to obtain the conda and requirement file inside the model zip.
+            self.session.file.get(self.model_zip_stage_path, context_dir)
+
+            extracted_model_dir = self._extract_model_zip(context_dir)
+
+            dc = docker_context.DockerContext(
+                context_dir=context_dir, model_dir=extracted_model_dir, use_gpu=self.use_gpu
+            )
+            dc.build()
+            self._build_image_from_context(context_dir=context_dir, docker_config_dir=docker_config_dir)
+
+    def _run_docker_commands(self, commands: List[str]) -> None:
+        """Run docker commands in a new child process.
+
+        Args:
+            commands: List of commands to run.
+
+        Raises:
+            RuntimeError: Occurs when docker commands failed to execute.
+        """
+        proc = subprocess.Popen(commands, cwd=os.getcwd(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        output_lines = []
+
+        if proc.stdout:
+            for line in iter(proc.stdout.readline, ""):
+                output_lines.append(line)
+                logging.info(line)
+
+        if proc.wait():
+            raise RuntimeError(f"Docker build failed: {''.join(output_lines)}")
+
+    def _build_image_from_context(
+        self, context_dir: str, docker_config_dir: str, *, platform: Platform = Platform.LINUX_AMD64
+    ) -> None:
+        """Builds a Docker image based on provided context.
+
+        Args:
+            context_dir: Path to context directory.
+            docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
+            platform: Target platform for the build output, in the format "os[/arch[/variant]]".
+        """
+
+        commands = [
+            "docker",
+            "--config",
+            docker_config_dir,
+            "buildx",
+            "build",
+            "--platform",
+            platform.value,
+            "--tag",
+            f"{self.image_tag}",
+            context_dir,
+        ]
+
+        self._run_docker_commands(commands)
+
+    def _upload(self, docker_config_dir: str) -> None:
+        """
+        Uploads image to the image registry. This process requires a "docker login" followed by a "docker push". Remove
+        local image at the end of the upload operation to save up local space. Image cache is kept for more performant
+        built experience at the cost of small storage footprint.
+
+        For image registry authentication, we will use a session token obtained from the Snowpark session object.
+        The token authentication mechanism is automatically used when the username is set to "0sessiontoken" according
+        to the registry implementation detailed in the following link:
+        https://github.com/snowflakedb/snowflake-image-registry/blob/277435c6fd79db2df9f863aa9d04dc875e034d85
+        /AuthAdapter/src/main/java/com/snowflake/registry/service/AuthHeader.java#L122
+
+        By default, Docker overwrites the local Docker config file "/.docker/config.json" whenever a docker login
+        occurs. However, to ensure better isolation between Snowflake-managed Docker credentials and the user's own
+        Docker credentials, we will not use the default Docker config. Instead, we will write the username and session
+        token to a temporary file and use "docker --config" so that it only applies to the specific Docker command being
+        executed, without affecting the user's local Docker setup. The credential file will be automatically removed
+        at the end of upload operation.
+
+        Args:
+            docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
+        """
+        commands = ["docker", "--config", docker_config_dir, "login", self.image_tag]
+        self._run_docker_commands(commands)
+
+        logging.info(f"Pushing image to image repo {self.image_tag}")
+        commands = ["docker", "--config", docker_config_dir, "push", self.image_tag]
+        self._run_docker_commands(commands)

snowflake/ml/model/_deploy_client/image_builds/docker_context.py

@@ -0,0 +1,89 @@
+import importlib
+import os
+import shutil
+import string
+from abc import ABC
+
+from snowflake.ml.model._deploy_client.utils import constants
+
+
+class DockerContext(ABC):
+    """
+    Constructs the Docker context directory required for image building.
+    """
+
+    def __init__(self, context_dir: str, model_dir: str, *, use_gpu: bool = False) -> None:
+        """Initialization
+
+        Args:
+            context_dir: Path to context directory.
+            model_dir: Path to local model directory.
+            use_gpu: Boolean flag for generating the CPU or GPU base image.
+        """
+        self.context_dir = context_dir
+        self.model_dir = model_dir
+        # TODO(shchen): SNOW-825995, Define dockerfile template used for model deployment. use_gpu will be used.
+        self.use_gpu = use_gpu
+
+    def build(self) -> None:
+        """
+        Generates and/or moves resources into the Docker context directory.Rename the random model directory name to
+        constant "model_dir" instead for better readability.
+        """
+        self._generate_inference_code()
+        self._copy_entrypoint_script_to_docker_context()
+        self._copy_snowml_source_code_to_docker_context()
+        self._generate_docker_file()
+
+    def _copy_snowml_source_code_to_docker_context(self) -> None:
+        """Copy the entire snowflake/ml source code to docker context. This will be particularly useful for CI tests
+        against latest changes.
+
+        Note that we exclude the experimental directory mainly for development scenario; as experimental directory won't
+        be included in the release.
+        """
+        snow_ml_source_dir = list(importlib.import_module("snowflake.ml").__path__)[0]
+        shutil.copytree(
+            snow_ml_source_dir,
+            os.path.join(self.context_dir, "snowflake", "ml"),
+            ignore=shutil.ignore_patterns("*.pyc", "experimental"),
+        )
+
+    def _copy_entrypoint_script_to_docker_context(self) -> None:
+        """Copy gunicorn_run.sh entrypoint to docker context directory."""
+        path = os.path.join(os.path.dirname(__file__), constants.ENTRYPOINT_SCRIPT)
+        assert os.path.exists(path), f"Run script file missing at path: {path}"
+        shutil.copy(path, os.path.join(self.context_dir, constants.ENTRYPOINT_SCRIPT))
+
+    def _generate_docker_file(self) -> None:
+        """
+        Generates dockerfile based on dockerfile template.
+        """
+        docker_file_path = os.path.join(self.context_dir, "Dockerfile")
+        docker_file_template = os.path.join(os.path.dirname(__file__), "templates/dockerfile_template")
+
+        with open(docker_file_path, "w", encoding="utf-8") as dockerfile, open(
+            docker_file_template, encoding="utf-8"
+        ) as template:
+            dockerfile_content = string.Template(template.read()).safe_substitute(
+                {
+                    # TODO(shchen): SNOW-835411, Support overwriting base image
+                    "base_image": "mambaorg/micromamba:focal-cuda-11.7.1"
+                    if self.use_gpu
+                    else "mambaorg/micromamba:1.4.3",
+                    "model_dir": constants.MODEL_DIR,
+                    "inference_server_dir": constants.INFERENCE_SERVER_DIR,
+                    "entrypoint_script": constants.ENTRYPOINT_SCRIPT,
+                }
+            )
+            dockerfile.write(dockerfile_content)
+
+    def _generate_inference_code(self) -> None:
+        """
+        Generates inference code based on the app template and creates a folder named 'server' to house the inference
+        server code.
+        """
+        inference_server_folder_path = os.path.join(os.path.dirname(__file__), constants.INFERENCE_SERVER_DIR)
+        destination_folder_path = os.path.join(self.context_dir, constants.INFERENCE_SERVER_DIR)
+        ignore_patterns = shutil.ignore_patterns("BUILD.bazel", "*test.py", "*.\\.*", "__pycache__")
+        shutil.copytree(inference_server_folder_path, destination_folder_path, ignore=ignore_patterns)

snowflake/ml/model/_deploy_client/image_builds/gunicorn_run.sh

@@ -0,0 +1,24 @@
+#!/bin/sh
+set -eu
+
+OS=$(uname)
+
+if [ "${OS}" = "Linux" ]; then
+    NUM_CORES=$(nproc)
+elif [ "${OS}" = "Darwin" ]; then
+    # macOS
+    NUM_CORES=$(sysctl -n hw.ncpu)
+elif [ "${OS}" = "Windows" ]; then
+    NUM_CORES=$(wmic cpu get NumberOfCores | grep -Eo '[0-9]+')
+else
+    echo "Unsupported operating system: ${OS}"
+    exit 1
+fi
+
+# Based on the Gunicorn documentation, set the number of workers to number_of_cores * 2 + 1. This assumption is
+# based on an ideal scenario where one core is handling two processes simultaneously, while one process is dedicated to
+# IO operations and the other process is performing compute tasks.
+NUM_WORKERS=$((NUM_CORES * 2 + 1))
+echo "Number of CPU cores: $NUM_CORES"
+echo "Setting number of workers to $NUM_WORKERS"
+exec /opt/conda/bin/gunicorn --preload -w "$NUM_WORKERS" -k uvicorn.workers.UvicornWorker -b 0.0.0.0:5000 inference_server.main:app

snowflake/ml/model/_deploy_client/image_builds/inference_server/main.py

@@ -0,0 +1,118 @@
+import logging
+import os
+import tempfile
+import zipfile
+
+import pandas as pd
+from starlette import applications, requests, responses, routing
+
+logger = logging.getLogger(__name__)
+loaded_model = None
+
+
+def _run_setup() -> None:
+    """Set up logging and load model into memory."""
+    # Align the application logger's handler with Gunicorn's to capture logs from all processes.
+    gunicorn_logger = logging.getLogger("gunicorn.error")
+    logger.handlers = gunicorn_logger.handlers
+    logger.setLevel(gunicorn_logger.level)
+
+    from snowflake.ml.model import _model as model_api
+
+    global loaded_model
+
+    MODEL_ZIP_STAGE_PATH = os.getenv("MODEL_ZIP_STAGE_PATH")
+    assert MODEL_ZIP_STAGE_PATH, "Missing environment variable MODEL_ZIP_STAGE_PATH"
+    root_path = os.path.abspath(os.sep)
+    model_zip_stage_path = os.path.join(root_path, MODEL_ZIP_STAGE_PATH)
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        if zipfile.is_zipfile(model_zip_stage_path):
+            extracted_dir = os.path.join(tmp_dir, "extracted_model_dir")
+            logger.info(f"Extracting model zip from {model_zip_stage_path} to {extracted_dir}")
+            with zipfile.ZipFile(model_zip_stage_path, "r") as model_zip:
+                if len(model_zip.namelist()) > 1:
+                    model_zip.extractall(extracted_dir)
+        else:
+            raise RuntimeError(f"No model zip found at stage path: {model_zip_stage_path}")
+        logger.info(f"Loading model from {extracted_dir} into memory")
+        loaded_model, _ = model_api._load_model_for_deploy(model_dir_path=extracted_dir)
+        logger.info("Successfully loaded model into memory")
+
+
+async def ready(request: requests.Request) -> responses.JSONResponse:
+    """Endpoint to check if the application is ready."""
+    return responses.JSONResponse({"status": "ready"})
+
+
+async def predict(request: requests.Request) -> responses.JSONResponse:
+    """Endpoint to make predictions based on input data.
+
+    Args:
+        request: The input data is expected to be in the following JSON format:
+            {
+                "data": [
+                    [0, 5.1, 3.5, 4.2, 1.3],
+                    [1, 4.7, 3.2, 4.1, 4.2]
+            }
+            Each row is represented as a list, where the first element denotes the index of the row.
+
+    Returns:
+        Two possible responses:
+        For success, return a JSON response {"data": [[0, 1], [1, 2]]}, where the first element of each resulting list
+            denotes the index of the row, and the rest of the elements represent the prediction results for that row.
+        For an error, return {"error": error_message, "status_code": http_response_status_code}.
+    """
+    try:
+        input = await request.json()
+        assert "data" in input, "missing data field in the request input"
+        # The expression x[1:] is used to exclude the index of the data row.
+        input_data = [x[1:] for x in input.get("data")]
+        x = pd.DataFrame(input_data)
+        assert len(input_data) != 0 and not all(not row for row in input_data), "empty data"
+    except Exception as e:
+        error_message = f"Input data malformed: {str(e)}"
+        return responses.JSONResponse({"error": error_message}, status_code=400)
+
+    assert loaded_model
+
+    try:
+        # TODO(shchen): SNOW-835369, Support target method in inference server (Multi-task model).
+        # Mypy ignore will be fixed along with the above ticket.
+        predictions = loaded_model.predict(x)  # type: ignore[attr-defined]
+        result = predictions.to_records(index=True).tolist()
+        response = {"data": result}
+        return responses.JSONResponse(response)
+    except Exception as e:
+        error_message = f"Prediction failed: {str(e)}"
+        return responses.JSONResponse({"error": error_message}, status_code=400)
+
+
+def _in_test_mode() -> bool:
+    """Check if the code is running in test mode.
+
+    Specifically, it checks for the presence of
+    - "PYTEST_CURRENT_TEST" environment variable, which is automatically set by Pytest when running tests, and
+    - "TEST_WORKSPACE" environment variable, which is set by Bazel test, and
+    - "TEST_SRCDIR" environment variable, which is set by the Absl test.
+
+    Returns:
+        True if in test mode; otherwise, returns False
+    """
+    is_running_under_py_test = "PYTEST_CURRENT_TEST" in os.environ
+    is_running_under_bazel_test = "TEST_WORKSPACE" in os.environ
+    is_running_under_absl_test = "TEST_SRCDIR" in os.environ
+    return is_running_under_py_test or is_running_under_bazel_test or is_running_under_absl_test
+
+
+def run_app() -> applications.Starlette:
+    if not _in_test_mode():
+        _run_setup()
+    routes = [
+        routing.Route("/health", endpoint=ready, methods=["GET"]),
+        routing.Route("/predict", endpoint=predict, methods=["POST"]),
+    ]
+    return applications.Starlette(routes=routes)
+
+
+app = run_app()