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

snowflake/ml/_internal/container_services/image_registry/registry_client.py

@@ -1,212 +0,0 @@
-import http
-import logging
-from typing import Dict, Optional, cast
-from urllib.parse import urlunparse
-
-from snowflake.ml._internal.container_services.image_registry import (
-    http_client as image_registry_http_client,
-    imagelib,
-)
-from snowflake.ml._internal.exceptions import (
-    error_codes,
-    exceptions as snowml_exceptions,
-)
-from snowflake.snowpark import Session
-from snowflake.snowpark._internal import utils as snowpark_utils
-
-_MANIFEST_V1_HEADER = "application/vnd.oci.image.manifest.v1+json"
-_MANIFEST_V2_HEADER = "application/vnd.docker.distribution.manifest.v2+json"
-_SUPPORTED_MANIFEST_HEADERS = [_MANIFEST_V1_HEADER, _MANIFEST_V2_HEADER]
-
-logger = logging.getLogger(__name__)
-
-
-class ImageRegistryClient:
-    """
-    A partial implementation of an SPCS image registry client. The client utilizes the ImageRegistryHttpClient under
-    the hood, incorporating a retry mechanism to handle intermittent 401 errors from the SPCS image registry.
-    """
-
-    def __init__(self, session: Session, full_dest_image_name: str) -> None:
-        """Initialization
-
-        Args:
-            session: Snowpark session
-            full_dest_image_name: Based on dest image name, repo url can be inferred.
-        """
-        self.image_registry_http_client = image_registry_http_client.ImageRegistryHttpClient(
-            session=session,
-            repo_url=self._convert_to_v2_manifests_url(full_image_name=full_dest_image_name),
-        )
-
-    def _convert_to_v2_manifests_url(self, full_image_name: str) -> str:
-        """Converts a full image name to a Docker Registry HTTP API V2 URL:
-        https://docs.docker.com/registry/spec/api/#existing-manifests
-
-        org-account.registry-dev.snowflakecomputing.com/db/schema/repo/image_name:tag becomes
-        https://org-account.registry-dev.snowflakecomputing.com/v2/db/schema/repo/image_name/manifests/tag
-
-        Args:
-            full_image_name: a string consists of image name and image tag.
-
-        Returns:
-            Docker HTTP V2 URL for checking manifest existence.
-        """
-        scheme = "https"
-        full_image_name_parts = full_image_name.split(":")
-        assert len(full_image_name_parts) == 2, "full image name should include both image name and tag"
-
-        image_name = full_image_name_parts[0]
-        tag = full_image_name_parts[1]
-        image_name_parts = image_name.split("/")
-        domain = image_name_parts[0]
-        rest = "/".join(image_name_parts[1:])
-        path = f"/v2/{rest}/manifests/{tag}"
-        url_tuple = (scheme, domain, path, "", "", "")
-        return urlunparse(url_tuple)
-
-    def _get_accept_headers(self) -> Dict[str, str]:
-        # Depending on the built image, the media type of the image manifest might be either
-        # application/vnd.oci.image.manifest.v1+json or application/vnd.docker.distribution.manifest.v2+json
-        # Hence we need to check for both, otherwise it could result in false negative.
-        return {"Accept": ",".join(_SUPPORTED_MANIFEST_HEADERS)}
-
-    def image_exists(self, full_image_name: str) -> bool:
-        """Check whether image already exists in the registry.
-
-        Args:
-            full_image_name: Full image name consists of image name and image tag.
-
-        Returns:
-            Boolean value. True when image already exists, else False.
-
-        """
-        # When running in SPROC, the Sproc session connection will not have _rest object associated, which makes it
-        # unable to fetch session token needed to authenticate to SPCS image registry.
-        if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-            return False
-        v2_api_url = self._convert_to_v2_manifests_url(full_image_name)
-        headers = self._get_accept_headers()
-        status = self.image_registry_http_client.head(v2_api_url, headers=headers).status_code
-        return status == http.HTTPStatus.OK
-
-    def _get_manifest(self, full_image_name: str) -> Dict[str, str]:
-        """Retrieve image manifest file. Given Docker manifest comes with two versions, and for each version the
-        corresponding request header is required for a successful HTTP response.
-
-        Args:
-            full_image_name: Full image name.
-
-        Returns:
-            Full manifest content as a python dict.
-
-        Raises:
-            SnowflakeMLException: when failed to retrieve manifest.
-        """
-
-        v2_api_url = self._convert_to_v2_manifests_url(full_image_name)
-        res = self.image_registry_http_client.get(v2_api_url, headers=self._get_accept_headers())
-        if res.status_code == http.HTTPStatus.OK:
-            return cast(Dict[str, str], res.json())
-        raise snowml_exceptions.SnowflakeMLException(
-            error_code=error_codes.INTERNAL_SNOWFLAKE_IMAGE_REGISTRY_ERROR,
-            original_exception=ValueError(
-                f"Failed to retrieve manifest for {full_image_name}. \n"
-                f"HTTP status code: {res.status_code}. Full response: {res.text}."
-            ),
-        )
-
-    def add_tag_to_remote_image(self, original_full_image_name: str, new_tag: str) -> None:
-        """Add a tag to an image in the registry.
-
-        Args:
-            original_full_image_name: The full image name is required to fetch manifest.
-            new_tag:  New tag to be added to the image.
-
-        Returns:
-            None
-
-        Raises:
-            SnowflakeMLException: when failed to push the newly updated manifest.
-        """
-
-        if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-            return None
-
-        full_image_name_parts = original_full_image_name.split(":")
-        assert len(full_image_name_parts) == 2, "full image name should include both image name and tag"
-        new_full_image_name = ":".join([full_image_name_parts[0], new_tag])
-        if self.image_exists(new_full_image_name):
-            # Early return if image with the associated tag already exists.
-            return
-        api_url = self._convert_to_v2_manifests_url(new_full_image_name)
-        manifest = self._get_manifest(full_image_name=original_full_image_name)
-        manifest_copy = manifest.copy()
-        manifest_copy["tag"] = new_tag
-        headers = self._get_accept_headers()
-        # Http Content-Type does not support multi-value, hence need to construct separate header.
-        put_header_v1 = {
-            **headers,
-            "Content-Type": _MANIFEST_V1_HEADER,
-        }
-        put_header_v2 = {
-            **headers,
-            "Content-Type": _MANIFEST_V2_HEADER,
-        }
-
-        res1 = self.image_registry_http_client.put(api_url, headers=put_header_v1, json=manifest_copy)
-        if res1.status_code != http.HTTPStatus.CREATED:
-            res2 = self.image_registry_http_client.put(api_url, headers=put_header_v2, json=manifest_copy)
-            if res2.status_code != http.HTTPStatus.CREATED:
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INTERNAL_SNOWFLAKE_IMAGE_REGISTRY_ERROR,
-                    original_exception=ValueError(
-                        f"Failed to push manifest for {new_full_image_name}. Two requests filed: \n"
-                        f"HTTP status code 1: {res1.status_code}. Full response 1: {res1.text}. \n"
-                        f"HTTP status code 2: {res2.status_code}. Full response 2: {res2.text}"
-                    ),
-                )
-        assert self.image_exists(
-            new_full_image_name
-        ), f"{new_full_image_name} should exist in image repo after a successful manifest update"
-
-    def copy_image(
-        self,
-        source_image_with_digest: str,
-        dest_image_with_tag: str,
-        arch: Optional[imagelib._Arch] = None,
-    ) -> None:
-        """Util function to copy image across registry. Currently supported pulling from public docker image repo to
-        SPCS image registry.
-
-        Args:
-            source_image_with_digest: source image with digest, e.g. gcr.io/kaniko-project/executor@sha256:b8c0977
-            dest_image_with_tag: destination image with tag.
-            arch: architecture of source image.
-
-        Returns:
-            None
-        """
-        logger.info(f"Copying image from {source_image_with_digest} to {dest_image_with_tag}")
-        if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-            logger.warning(f"Running inside Sproc. Please ensure image already exists at {dest_image_with_tag}")
-            return None
-
-        arch = arch or imagelib._Arch("amd64", "linux")
-
-        src_image = imagelib.convert_to_image_descriptor(source_image_with_digest, with_digest=True)
-        dest_image = imagelib.convert_to_image_descriptor(
-            dest_image_with_tag,
-            with_tag=True,
-        )
-        # TODO[shchen]: Remove the imagelib, instead rely on the copy image system function later.
-        imagelib.copy_image(
-            src_image=src_image,
-            dest_image=dest_image,
-            arch=arch,
-            src_retryable_http=image_registry_http_client.ImageRegistryHttpClient(
-                repo_url=src_image.registry_name, no_cred=True
-            ),
-            dest_retryable_http=self.image_registry_http_client,
-        )
-        logger.info("Image copy completed successfully")

snowflake/ml/_internal/utils/log_stream_processor.py

@@ -1,30 +0,0 @@
-import logging
-from typing import Optional
-
-logger = logging.getLogger(__name__)
-
-
-class LogStreamProcessor:
-    def __init__(self) -> None:
-        self.last_line_seen = 0
-
-    def process_new_logs(self, job_logs: Optional[str], *, log_level: int = logging.INFO) -> None:
-        if not job_logs:
-            return
-        log_entries = job_logs.split("\n")
-        start_index = self.last_line_seen
-        log_length = len(log_entries)
-        for i in range(start_index, log_length):
-            log_entry = log_entries[i]
-            if log_level == logging.DEBUG:
-                logger.debug(log_entry)
-            elif log_level == logging.INFO:
-                logger.info(log_entry)
-            elif log_level == logging.WARNING:
-                logger.warning(log_entry)
-            elif log_level == logging.ERROR:
-                logger.error(log_entry)
-            elif log_level == logging.CRITICAL:
-                logger.critical(log_entry)
-
-        self.last_line_seen = log_length

snowflake/ml/_internal/utils/session_token_manager.py

@@ -1,46 +0,0 @@
-import base64
-import json
-from typing import TypedDict
-
-from snowflake import snowpark
-
-
-class SessionToken(TypedDict):
-    token: str
-    expires_in: str
-
-
-class SessionTokenManager:
-    def __init__(self, session: snowpark.Session) -> None:
-        self._session = session
-
-    def get_session_token(self) -> SessionToken:
-        """
-        This function retrieves the session token from Snowpark session object.
-
-        Returns:
-            The session token string value.
-        """
-        ctx = self._session._conn._conn
-        assert ctx._rest, "SnowflakeRestful is not set in session"
-        token_data = ctx._rest._token_request("ISSUE")
-        session_token = token_data["data"]["sessionToken"]
-        validity_in_seconds = token_data["data"]["validityInSecondsST"]
-        assert session_token, "session_token is not obtained successfully from the session object"
-        assert validity_in_seconds, "validityInSecondsST is not obtained successfully from the session object"
-        return {"token": session_token, "expires_in": validity_in_seconds}
-
-    def get_base64_encoded_token(self, username: str = "0sessiontoken") -> str:
-        """This function returns the base64 encoded username:password, which is compatible with registry, such as
-        SnowService image registry, that uses Docker credential helper. In this case, password will be session token.
-
-        Args:
-            username: username for authentication.
-
-        Returns:
-            base64 encoded credential string.
-
-        """
-        credentials = f"{username}:{json.dumps(self.get_session_token())}"
-        encoded_credentials = base64.b64encode(credentials.encode("utf-8")).decode("utf-8")
-        return encoded_credentials

snowflake/ml/_internal/utils/spcs_attribution_utils.py

@@ -1,122 +0,0 @@
-import logging
-from datetime import datetime
-from typing import Any, Dict, Optional
-
-from snowflake import snowpark
-from snowflake.ml._internal import telemetry
-from snowflake.ml._internal.utils import query_result_checker
-
-logger = logging.getLogger(__name__)
-
-_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S.%f %z"
-_COMPUTE_POOL = "compute_pool"
-_CREATED_ON = "created_on"
-_INSTANCE_FAMILY = "instance_family"
-_NAME = "name"
-_TELEMETRY_PROJECT = "MLOps"
-_TELEMETRY_SUBPROJECT = "SpcsDeployment"
-_SERVICE_START = "SPCS_SERVICE_START"
-_SERVICE_END = "SPCS_SERVICE_END"
-
-
-def _desc_compute_pool(session: snowpark.Session, compute_pool_name: str) -> Dict[str, Any]:
-    sql = f"DESC COMPUTE POOL {compute_pool_name}"
-    result = (
-        query_result_checker.SqlResultValidator(
-            session=session,
-            query=sql,
-        )
-        .has_column(_INSTANCE_FAMILY)
-        .has_column(_NAME)
-        .has_dimensions(expected_rows=1)
-        .validate()
-    )
-    return result[0].as_dict()
-
-
-def _desc_service(session: snowpark.Session, fully_qualified_name: str) -> Dict[str, Any]:
-    sql = f"DESC SERVICE {fully_qualified_name}"
-    result = (
-        query_result_checker.SqlResultValidator(
-            session=session,
-            query=sql,
-        )
-        .has_column(_COMPUTE_POOL)
-        .has_dimensions(expected_rows=1)
-        .validate()
-    )
-    return result[0].as_dict()
-
-
-def _get_current_time() -> datetime:
-    """
-    This method exists to make it easier to mock datetime in test.
-
-    Returns:
-        current datetime
-    """
-    return datetime.now()
-
-
-def _send_service_telemetry(
-    fully_qualified_name: Optional[str] = None,
-    compute_pool_name: Optional[str] = None,
-    service_details: Optional[Dict[str, Any]] = None,
-    compute_pool_details: Optional[Dict[str, Any]] = None,
-    duration_in_seconds: Optional[int] = None,
-    kwargs: Optional[Dict[str, Any]] = None,
-) -> None:
-    try:
-        telemetry.send_custom_usage(
-            project=_TELEMETRY_PROJECT,
-            subproject=_TELEMETRY_SUBPROJECT,
-            telemetry_type=telemetry.TelemetryField.TYPE_SNOWML_SPCS_USAGE.value,
-            data={
-                "service_name": fully_qualified_name,
-                "compute_pool_name": compute_pool_name,
-                "service_details": service_details,
-                "compute_pool_details": compute_pool_details,
-                "duration_in_seconds": duration_in_seconds,
-            },
-            kwargs=kwargs,
-        )
-    except Exception as e:
-        logger.error(f"Failed to send service telemetry: {e}")
-
-
-def record_service_start(session: snowpark.Session, fully_qualified_name: str) -> None:
-    service_details = _desc_service(session, fully_qualified_name)
-    compute_pool_name = service_details[_COMPUTE_POOL]
-    compute_pool_details = _desc_compute_pool(session, compute_pool_name)
-
-    _send_service_telemetry(
-        fully_qualified_name=fully_qualified_name,
-        compute_pool_name=compute_pool_name,
-        service_details=service_details,
-        compute_pool_details=compute_pool_details,
-        kwargs={telemetry.TelemetryField.KEY_CUSTOM_TAGS.value: _SERVICE_START},
-    )
-
-    logger.info(f"Service {fully_qualified_name} created with compute pool {compute_pool_name}.")
-
-
-def record_service_end(session: snowpark.Session, fully_qualified_name: str) -> None:
-    service_details = _desc_service(session, fully_qualified_name)
-    compute_pool_details = _desc_compute_pool(session, service_details[_COMPUTE_POOL])
-    compute_pool_name = service_details[_COMPUTE_POOL]
-
-    created_on_datetime: datetime = service_details[_CREATED_ON]
-    current_time: datetime = _get_current_time()
-    current_time = current_time.replace(tzinfo=created_on_datetime.tzinfo)
-    duration_in_seconds = int((current_time - created_on_datetime).total_seconds())
-
-    _send_service_telemetry(
-        fully_qualified_name=fully_qualified_name,
-        compute_pool_name=compute_pool_name,
-        service_details=service_details,
-        compute_pool_details=compute_pool_details,
-        duration_in_seconds=duration_in_seconds,
-        kwargs={telemetry.TelemetryField.KEY_CUSTOM_TAGS.value: _SERVICE_END},
-    )
-
-    logger.info(f"Service {fully_qualified_name} deleted from compute pool {compute_pool_name}")

snowflake/ml/_internal/utils/uri.py

@@ -1,77 +0,0 @@
-import posixpath
-from typing import Optional, Tuple
-from urllib.parse import ParseResult, urlparse, urlunparse
-
-from snowflake.ml._internal.utils import identifier
-
-_LOCAL_URI_SCHEMES = ["", "file"]
-_HTTP_URI_SCHEMES = ["http", "https"]
-_SNOWFLAKE_STAGE_URI_SCHEMES = ["sfc", "sfstage"]
-
-
-def is_local_uri(uri: str) -> bool:
-    """Returns true if the URI has a scheme that indicates a local file."""
-    return urlparse(uri).scheme in _LOCAL_URI_SCHEMES
-
-
-def is_http_uri(uri: str) -> bool:
-    """Returns true if the URI has a scheme that indicates a web (http,https) address."""
-    return urlparse(uri).scheme in _HTTP_URI_SCHEMES
-
-
-def is_snowflake_stage_uri(uri: str) -> bool:
-    """Returns true if the URI is a scheme that indicates a Snowflake stage location."""
-    return urlparse(uri).scheme in _SNOWFLAKE_STAGE_URI_SCHEMES
-
-
-def get_snowflake_stage_path_from_uri(uri: str) -> Optional[str]:
-    """Returns the stage path pointed by the URI.
-
-    Args:
-        uri: URI for which stage file is needed.
-
-    Returns:
-        The Snowflake stage location encoded by the given URI. Returns None if the URI is not pointing to a Snowflake
-            stage.
-    """
-    if not is_snowflake_stage_uri(uri):
-        return None
-    uri_components = urlparse(uri)
-    # 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:
-    """Returns the scheme for the given URI."""
-    return urlparse(uri).scheme
-
-
-def get_uri_from_snowflake_stage_path(stage_path: str) -> str:
-    """Generates a URI from Snowflake stage path."""
-    assert stage_path.startswith("@")
-    (db, schema, stage, path) = identifier.parse_snowflake_stage_path(
-        posixpath.normpath(identifier.remove_prefix(stage_path, "@"))
-    )
-    return urlunparse(
-        ParseResult(
-            scheme=_SNOWFLAKE_STAGE_URI_SCHEMES[0],
-            netloc=identifier.get_schema_level_object_identifier(db, schema, stage),
-            path=path,
-            params="",
-            query="",
-            fragment="",
-        )
-    )
-
-
-def get_stage_and_path(stage_path: str) -> Tuple[str, str]:
-    assert stage_path.startswith("@"), f"stage path should start with @, actual: {stage_path}"
-    (db, schema, stage, path) = identifier.parse_snowflake_stage_path(
-        posixpath.normpath(identifier.remove_prefix(stage_path, "@"))
-    )
-    full_qualified_stage = "@" + identifier.get_schema_level_object_identifier(db, schema, stage)
-    return full_qualified_stage, path.lstrip("/")