snowflake-ml-python 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl

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

@@ -1,215 +0,0 @@
-import logging
-import os
-import posixpath
-from string import Template
-from typing import List
-
-import importlib_resources
-
-from snowflake import snowpark
-from snowflake.ml._internal import file_utils
-from snowflake.ml._internal.container_services.image_registry import (
-    registry_client as image_registry_client,
-)
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
-from snowflake.ml._internal.utils import identifier
-from snowflake.ml.model._deploy_client import image_builds
-from snowflake.ml.model._deploy_client.image_builds import base_image_builder
-from snowflake.ml.model._deploy_client.utils import constants, snowservice_client
-
-logger = logging.getLogger(__name__)
-
-
-class ServerImageBuilder(base_image_builder.ImageBuilder):
-    """
-    Server-side image building and upload to model registry.
-    """
-
-    def __init__(
-        self,
-        *,
-        context_dir: str,
-        full_image_name: str,
-        image_repo: str,
-        session: snowpark.Session,
-        artifact_stage_location: str,
-        compute_pool: str,
-        job_name: str,
-        external_access_integrations: List[str],
-    ) -> None:
-        """Initialization
-
-        Args:
-            context_dir: Local docker context dir.
-            full_image_name: Full image name consists of image name and image tag.
-            image_repo: Path to image repository.
-            session: Snowpark session
-            artifact_stage_location: Spec file and future deployment related artifacts will be stored under
-                {stage}/models/{model_id}
-            compute_pool: The compute pool used to run docker image build workload.
-            job_name: job_name to use.
-            external_access_integrations: EAIs for network connection.
-        """
-        self.context_dir = context_dir
-        self.image_repo = image_repo
-        self.full_image_name = full_image_name
-        self.session = session
-        self.artifact_stage_location = artifact_stage_location
-        self.compute_pool = compute_pool
-        self.external_access_integrations = external_access_integrations
-        self.job_name = job_name
-        self.client = snowservice_client.SnowServiceClient(session)
-
-        assert artifact_stage_location.startswith(
-            "@"
-        ), f"stage path should start with @, actual: {artifact_stage_location}"
-
-    def build_and_upload_image(self) -> None:
-        """
-        Builds and uploads an image to the model registry.
-        """
-        logger.info("Starting server-side image build")
-        self._build_image_in_remote_job()
-
-    def _build_image_in_remote_job(self) -> None:
-        context_tarball_stage_location = f"{self.artifact_stage_location}/{constants.CONTEXT}.tar.gz"
-        spec_stage_location = f"{self.artifact_stage_location}/{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}.yaml"
-        kaniko_shell_script_stage_location = f"{self.artifact_stage_location}/{constants.KANIKO_SHELL_SCRIPT_NAME}"
-
-        self._compress_and_upload_docker_context_tarball(context_tarball_stage_location=context_tarball_stage_location)
-
-        self._construct_and_upload_docker_entrypoint_script(
-            context_tarball_stage_location=context_tarball_stage_location
-        )
-
-        # This is more of a workaround to support non-spcs-registry images.
-        # TODO[shchen] remove such logic when first-party-image is supported on snowservice registry.
-        # The regular Kaniko image doesn't include a shell; only the debug image comes with a shell. We need a shell
-        # as we use an sh script to launch Kaniko
-        kaniko_image = "/".join([self.image_repo.rstrip("/"), constants.KANIKO_IMAGE])
-        registry_client = image_registry_client.ImageRegistryClient(self.session, kaniko_image)
-        if registry_client.image_exists(kaniko_image):
-            logger.debug(f"Kaniko image already existed at {kaniko_image}, skipping uploading")
-        else:
-            # Following Digest is corresponding to v1.16.0-debug tag. Note that we cannot copy from image that contains
-            # tag as the underlying image blob copying API supports digest only.
-            registry_client.copy_image(
-                source_image_with_digest="gcr.io/kaniko-project/executor@sha256:"
-                "b8c0977f88f24dbd7cbc2ffe5c5f824c410ccd0952a72cc066efc4b6dfbb52b6",
-                dest_image_with_tag=kaniko_image,
-            )
-        self._construct_and_upload_job_spec(
-            base_image=kaniko_image,
-            kaniko_shell_script_stage_location=kaniko_shell_script_stage_location,
-        )
-        self._launch_kaniko_job(spec_stage_location)
-
-    def _construct_and_upload_docker_entrypoint_script(self, context_tarball_stage_location: str) -> None:
-        """Construct a shell script that invokes logic to uncompress the docker context tarball, then invoke Kaniko
-        executor to build images and push to image registry; the script will also ensure the docker credential(used to
-        authenticate to image registry) stays up-to-date when session token refreshes.
-
-        Args:
-            context_tarball_stage_location: Path context directory stage location.
-        """
-        kaniko_shell_script_template = (
-            importlib_resources.files(image_builds)
-            .joinpath(f"templates/{constants.KANIKO_SHELL_SCRIPT_TEMPLATE}")
-            .read_text("utf-8")
-        )
-
-        kaniko_shell_file = os.path.join(self.context_dir, constants.KANIKO_SHELL_SCRIPT_NAME)
-
-        with file_utils.open_file(kaniko_shell_file, "w+") as script_file:
-            normed_artifact_stage_path = posixpath.normpath(identifier.remove_prefix(self.artifact_stage_location, "@"))
-            params = {
-                # Remove @ in the beginning, append "/" to denote root directory.
-                "tar_from": "/" + posixpath.normpath(identifier.remove_prefix(context_tarball_stage_location, "@")),
-                # Remove @ in the beginning, append "/" to denote root directory.
-                "tar_to": "/" + normed_artifact_stage_path,
-                "context_dir": f"dir:///{normed_artifact_stage_path}/{constants.CONTEXT}",
-                "image_repo": self.image_repo,
-                # All models will be sharing the same layer cache from the image_repo/cache directory.
-                "cache_repo": f"{self.image_repo.rstrip('/')}/cache",
-                "image_destination": self.full_image_name,
-            }
-            template = Template(kaniko_shell_script_template)
-            script = template.safe_substitute(params)
-            script_file.write(script)
-            logger.debug(f"script content: \n\n {script}")
-        self.session.file.put(
-            local_file_name=kaniko_shell_file,
-            stage_location=self.artifact_stage_location,
-            auto_compress=False,
-            overwrite=True,
-        )
-
-    def _compress_and_upload_docker_context_tarball(self, context_tarball_stage_location: str) -> None:
-        try:
-            with file_utils._create_tar_gz_stream(
-                source_dir=self.context_dir, arcname=constants.CONTEXT
-            ) as input_stream:
-                self.session.file.put_stream(
-                    input_stream=input_stream,
-                    stage_location=context_tarball_stage_location,
-                    auto_compress=False,
-                    overwrite=True,
-                )
-        except Exception as e:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
-                original_exception=RuntimeError(
-                    "Exception occurred when compressing docker context dir as tarball and upload to stage."
-                ),
-            ) from e
-
-    def _construct_and_upload_job_spec(self, base_image: str, kaniko_shell_script_stage_location: str) -> None:
-        assert kaniko_shell_script_stage_location.startswith(
-            "@"
-        ), f"stage path should start with @, actual: {kaniko_shell_script_stage_location}"
-
-        spec_template = (
-            importlib_resources.files(image_builds)
-            .joinpath(f"templates/{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}")
-            .read_text("utf-8")
-        )
-
-        spec_file_path = os.path.join(self.context_dir, f"{constants.IMAGE_BUILD_JOB_SPEC_TEMPLATE}.yaml")
-
-        with file_utils.open_file(spec_file_path, "w+") as spec_file:
-            assert self.artifact_stage_location.startswith("@")
-            normed_artifact_stage_path = posixpath.normpath(identifier.remove_prefix(self.artifact_stage_location, "@"))
-            (db, schema, stage, path) = identifier.parse_snowflake_stage_path(normed_artifact_stage_path)
-            content = Template(spec_template).safe_substitute(
-                {
-                    "base_image": base_image,
-                    "container_name": constants.KANIKO_CONTAINER_NAME,
-                    "stage": identifier.get_schema_level_object_identifier(db, schema, stage),
-                    # Remove @ in the beginning, append "/" to denote root directory.
-                    "script_path": "/"
-                    + posixpath.normpath(identifier.remove_prefix(kaniko_shell_script_stage_location, "@")),
-                    "mounted_token_path": constants.SPCS_MOUNTED_TOKEN_PATH,
-                }
-            )
-            spec_file.write(content)
-            spec_file.seek(0)
-            logger.debug(f"Kaniko job spec file: \n\n {spec_file.read()}")
-
-        self.session.file.put(
-            local_file_name=spec_file_path,
-            stage_location=self.artifact_stage_location,
-            auto_compress=False,
-            overwrite=True,
-        )
-
-    def _launch_kaniko_job(self, spec_stage_location: str) -> None:
-        logger.debug(f"Submitting job {self.job_name} for building docker image with kaniko")
-        self.client.create_job(
-            job_name=self.job_name,
-            compute_pool=self.compute_pool,
-            spec_stage_location=spec_stage_location,
-            external_access_integrations=self.external_access_integrations,
-        )

snowflake/ml/model/_deploy_client/image_builds/templates/dockerfile_template

@@ -1,53 +0,0 @@
-# Note that base image tag should not be 'latest' as it might cause false positive image cache hit.
-FROM ${base_image} as build
-
-COPY ${model_env_folder}/conda.yml conda.yml
-COPY ${model_env_folder}/requirements.txt requirements.txt
-COPY ${inference_server_dir} ./${inference_server_dir}
-COPY ${entrypoint_script} ./${entrypoint_script}
-
-USER root
-RUN if id mambauser >/dev/null 2>&1; then \
-        echo "mambauser already exists."; \
-    else \
-        # Set environment variables
-        export USER=mambauser && \
-        export UID=1000 && \
-        export HOME=/home/$USER && \
-        echo "Creating $USER user..." && \
-        adduser --disabled-password \
-            --gecos "A non-root user for running inference server" \
-            --uid $UID \
-            --home $HOME \
-            $USER; \
-    fi
-
-RUN chmod +rx conda.yml
-RUN chmod +rx requirements.txt
-RUN chmod +x ./${entrypoint_script}
-
-USER mambauser
-
-# Set MAMBA_DOCKERFILE_ACTIVATE=1 to activate the conda environment during build time.
-ARG MAMBA_DOCKERFILE_ACTIVATE=1
-ARG MAMBA_NO_LOW_SPEED_LIMIT=1
-
-# Bitsandbytes uses this ENVVAR to determine CUDA library location
-ENV CONDA_PREFIX=/opt/conda
-
-# The micromamba image comes with an empty environment named base.
-# CONDA_OVERRIDE_CUDA ref https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-virtual.html
-RUN --mount=type=cache,target=/opt/conda/pkgs CONDA_OVERRIDE_CUDA="${cuda_override_env}" \
-    micromamba install -y -n base -f conda.yml && \
-	python -m pip install "uvicorn[standard]" gunicorn starlette==0.30.0 && \
-	python -m pip install -r requirements.txt && \
-    micromamba clean -afy
-
-${copy_model_statement}
-
-${extra_env_statement}
-
-# Expose the port on which the Starlette app will run.
-EXPOSE 5000
-
-CMD ["./${entrypoint_script}"]

snowflake/ml/model/_deploy_client/image_builds/templates/image_build_job_spec_template

@@ -1,38 +0,0 @@
-spec:
-  container:
-    - name: "${container_name}"
-      image: "${base_image}"
-      command:
-        - sh
-      args:
-        - -c
-        - |
-          wait_for_file() {
-              file_path="$1"
-              timeout="$2"
-              elapsed_time=0
-              while [ ! -f "${file_path}" ]; do
-                if [ "${elapsed_time}" -ge "${timeout}" ]; then
-                  echo "Error: ${file_path} not found within ${timeout} seconds. Exiting."
-                  exit 1
-                fi
-                elapsed_time=$((elapsed_time + 1))
-                remaining_time=$((timeout - elapsed_time))
-                echo "Awaiting the mounting of ${file_path}. Wait time remaining: ${remaining_time} seconds"
-                sleep 1
-              done
-           }
-           wait_for_file "${script_path}" 300
-           wait_for_file "${mounted_token_path}" 300
-           chmod +x "${script_path}"
-           sh "${script_path}"
-      volumeMounts:
-        - name: vol1
-          mountPath: /local/user/vol1
-        - name: stagemount
-          mountPath: "/${stage}"
-  volume:
-    - name: vol1
-      source: local  # only local emptyDir volume is supported
-    - name: stagemount
-      source: "@${stage}"

snowflake/ml/model/_deploy_client/image_builds/templates/kaniko_shell_script_template

@@ -1,105 +0,0 @@
-#!/bin/sh
-
-# Set the file path to monitor
-REGISTRY_CRED_PATH="/kaniko/.docker/config.json"
-SESSION_TOKEN_PATH="/snowflake/session/token"
-
-# Function to gracefully terminate the file monitoring job
-cleanup() {
-  echo "Stopping file monitoring job..."
-  trap - INT TERM # Remove the signal handlers
-  kill -- -$$$ # Kill the entire process group. Extra $ to escape, the generated shell script should have two $.
-}
-
-# SNOW-990976, This is an additional safety check to ensure token file exists, on top of the token file check upon
-# launching SPCS job. This additional check could provide value in cases things go wrong with token refresh that result
-# in token file to disappear.
-wait_till_token_file_exists() {
-  timeout=60  # 1 minute timeout
-  elapsed_time=0
-
-  while [ ! -f "${SESSION_TOKEN_PATH}" ] && [ "$elapsed_time" -lt "$timeout" ]; do
-    sleep 1
-    elapsed_time=$((elapsed_time + 1))
-    remaining_time=$((timeout - elapsed_time))
-    echo "Waiting for token file to exist. Wait time remaining: ${remaining_time} seconds."
-  done
-
-  if [ ! -f "${SESSION_TOKEN_PATH}" ]; then
-    echo "Error: Token file '${SESSION_TOKEN_PATH}' does not show up within the ${timeout} seconds timeout period."
-    exit 1
-  fi
-}
-
-generate_registry_cred() {
-  wait_till_token_file_exists
-  AUTH_TOKEN=$(printf '0auth2accesstoken:%s' "$(cat ${SESSION_TOKEN_PATH})" | base64);
-  echo '{"auths":{"$image_repo":{"auth":"'"$AUTH_TOKEN"'"}}}' | tr -d '\n' > $REGISTRY_CRED_PATH;
-}
-
-on_session_token_change() {
-  wait_till_token_file_exists
-  # Get the initial checksum of the file
-  CHECKSUM=$(md5sum "${SESSION_TOKEN_PATH}" | awk '{ print $1 }')
-  # Run the command once before the loop
-  echo "Monitoring session token changes in the background..."
-  (
-    while true; do
-      wait_till_token_file_exists
-      # Get the current checksum of the file
-      CURRENT_CHECKSUM=$(md5sum "${SESSION_TOKEN_PATH}" | awk '{ print $1 }')
-      if [ "${CURRENT_CHECKSUM}" != "${CHECKSUM}" ]; then
-        # Session token file has changed, regenerate registry credential.
-        echo "Session token has changed. Regenerating registry auth credentials."
-        generate_registry_cred
-        CHECKSUM="${CURRENT_CHECKSUM}"
-      fi
-      # Wait for a short period of time before checking again
-      sleep 1
-    done
-  )
-}
-
-run_kaniko() {
-  # Run the Kaniko command in the foreground
-  echo "Starting Kaniko command..."
-
-  # Set cache ttl to a large value as snowservice registry doesn't support deleting cache anyway.
-  # Compression level set to 1 for fastest compression/decompression speed at the cost of compression ration.
-  /kaniko/executor \
-    --dockerfile Dockerfile \
-    --context ${context_dir} \
-    --destination=${image_destination} \
-    --cache=true \
-    --compressed-caching=false \
-    --cache-copy-layers=false \
-    --use-new-run \
-    --snapshot-mode=redo \
-    --cache-repo=${cache_repo} \
-    --cache-run-layers=true \
-    --cache-ttl=8760h \
-    --push-retry=3 \
-    --image-fs-extract-retry=5 \
-    --compression=zstd \
-    --compression-level=1 \
-    --log-timestamp
-}
-
-setup() {
-  tar -C "${tar_to}" -xf "${tar_from}";
-  generate_registry_cred
-  # Set up the signal handlers
-  trap cleanup TERM
-}
-
-setup
-
-# Running kaniko job on the foreground and session token monitoring on the background. When session token changes,
-# overwrite the existing registry cred file with the new session token.
-on_session_token_change &
-run_kaniko
-
-# Capture the exit code from the previous kaniko command.
-KANIKO_EXIT_CODE=$?
-# Exit with the same exit code as the Kaniko command. This then triggers the cleanup function.
-exit $KANIKO_EXIT_CODE