snowflake-ml-python 1.7.0__py3-none-any.whl → 1.7.2__py3-none-any.whl

snowflake/cortex/_sse_client.py

@@ -1,73 +1,125 @@
-from typing import Iterator, cast
+import json
+from typing import Any, Iterator, Optional
 
-import requests
+_FIELD_SEPARATOR = ":"
 
 
 class Event:
-    def __init__(self, event: str = "message", data: str = "") -> None:
+    """Representation of an event from the event stream."""
+
+    def __init__(
+        self,
+        id: Optional[str] = None,
+        event: str = "message",
+        data: str = "",
+        comment: Optional[str] = None,
+        retry: Optional[int] = None,
+    ) -> None:
+        self.id = id
         self.event = event
         self.data = data
+        self.comment = comment
+        self.retry = retry
 
     def __str__(self) -> str:
         s = f"{self.event} event"
+        if self.id:
+            s += f" #{self.id}"
         if self.data:
-            s += f", {len(self.data)} bytes"
+            s += ", {} byte{}".format(len(self.data), "s" if len(self.data) else "")
         else:
             s += ", no data"
+        if self.comment:
+            s += f", comment: {self.comment}"
+        if self.retry:
+            s += f", retry in {self.retry}ms"
         return s
 
 
+# This is copied from the snowpy library:
+# https://github.com/snowflakedb/snowpy/blob/main/libs/snowflake.core/src/snowflake/core/rest.py#L39
+# TODO(SNOW-1750723) - Current there’s code duplication across snowflake-ml-python
+# and snowpy library for Cortex REST API which was done to meet our GA timelines
+# Once snowpy has a release with https://github.com/snowflakedb/snowpy/pull/679, we should
+# remove the class here and directly refer from the snowflake.core package directly
 class SSEClient:
-    def __init__(self, response: requests.Response) -> None:
+    def __init__(self, event_source: Any, char_enc: str = "utf-8") -> None:
+        self._event_source = event_source
+        self._char_enc = char_enc
 
-        self.response = response
-
-    def _read(self) -> Iterator[str]:
-
-        lines = b""
-        for chunk in self.response:
+    def _read(self) -> Iterator[bytes]:
+        data = b""
+        for chunk in self._event_source:
             for line in chunk.splitlines(True):
-                lines += line
-                if lines.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
-                    yield cast(str, lines)
-                    lines = b""
-        if lines:
-            yield cast(str, lines)
+                data += line
+                if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
+                    yield data
+                    data = b""
+        if data:
+            yield data
 
     def events(self) -> Iterator[Event]:
-        for raw_event in self._read():
+        content_type = self._event_source.headers.get("Content-Type")
+        # The check for empty content-type is present because it's being populated after
+        # the change in https://github.com/snowflakedb/snowflake/pull/217654.
+        # This can be removed once the above change makes it to prod or we move to snowpy
+        # for SSEClient implementation.
+        if content_type == "text/event-stream" or not content_type:
+            return self._handle_sse()
+        elif content_type == "application/json":
+            return self._handle_json()
+        else:
+            raise ValueError(f"Unknown Content-Type: {content_type}")
+
+    def _handle_sse(self) -> Iterator[Event]:
+        for chunk in self._read():
             event = Event()
-            # splitlines() only uses \r and \n
-            for line in raw_event.splitlines():
+            # Split before decoding so splitlines() only uses \r and \n
+            for line_bytes in chunk.splitlines():
+                # Decode the line.
+                line = line_bytes.decode(self._char_enc)
 
-                line = cast(bytes, line).decode("utf-8")
+                # Lines starting with a separator are comments and are to be
+                # ignored.
+                if not line.strip() or line.startswith(_FIELD_SEPARATOR):
+                    continue
 
-                data = line.split(":", 1)
+                data = line.split(_FIELD_SEPARATOR, 1)
                 field = data[0]
 
+                # Ignore unknown fields.
+                if not hasattr(event, field):
+                    continue
+
                 if len(data) > 1:
+                    # From the spec:
                     # "If value starts with a single U+0020 SPACE character,
-                    # remove it from value. .strip() would remove all white spaces"
+                    # remove it from value."
                     if data[1].startswith(" "):
                         value = data[1][1:]
                     else:
                         value = data[1]
                 else:
+                    # If no value is present after the separator,
+                    # assume an empty value.
                     value = ""
 
                 # The data field may come over multiple lines and their values
                 # are concatenated with each other.
+                current_value = getattr(event, field, "")
                 if field == "data":
-                    event.data += value + "\n"
-                elif field == "event":
-                    event.event = value
+                    new_value = current_value + value + "\n"
+                else:
+                    new_value = value
+                setattr(event, field, new_value)
 
+            # Events with no data are not dispatched.
             if not event.data:
                 continue
 
             # If the data field ends with a newline, remove it.
             if event.data.endswith("\n"):
-                event.data = event.data[0:-1]  # Replace trailing newline - rstrip would remove multiple.
+                event.data = event.data[0:-1]
 
             # Empty event names default to 'message'
             event.event = event.event or "message"
@@ -77,5 +129,16 @@ class SSEClient:
 
             yield event
 
+    def _handle_json(self) -> Iterator[Event]:
+        data_list = json.loads(self._event_source.data.decode(self._char_enc))
+        for data in data_list:
+            yield Event(
+                id=data.get("id"),
+                event=data.get("event"),
+                data=data.get("data"),
+                comment=data.get("comment"),
+                retry=data.get("retry"),
+            )
+
     def close(self) -> None:
-        self.response.close()
+        self._event_source.close()

snowflake/cortex/_util.py

@@ -1,6 +1,8 @@
-from typing import Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 from snowflake import snowpark
+from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml._internal.utils import formatting
 from snowflake.snowpark import context, functions
 
 CORTEX_FUNCTIONS_TELEMETRY_PROJECT = "CortexFunctions"
@@ -64,3 +66,30 @@ def _call_sql_function_immediate(
     empty_df = session.create_dataframe([snowpark.Row()])
     df = empty_df.select(functions.builtin(function)(*lit_args))
     return cast(str, df.collect()[0][0])
+
+
+def call_sql_function_literals(function: str, session: Optional[snowpark.Session], *args: Any) -> str:
+    r"""Call a SQL function with only literal arguments.
+
+    This is useful for calling system functions.
+
+    Args:
+        function: The name of the function to be called.
+        session: The Snowpark session to use.
+        *args: The list of arguments
+
+    Returns:
+        String value that corresponds the the first cell in the dataframe.
+
+    Raises:
+        SnowflakeMLException: If no session is given and no active session exists.
+    """
+    if session is None:
+        session = context.get_active_session()
+    if session is None:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_SNOWPARK_SESSION,
+        )
+
+    function_arguments = ",".join(["NULL" if arg is None else formatting.format_value_for_select(arg) for arg in args])
+    return cast(str, session.sql(f"SELECT {function}({function_arguments})").collect()[0][0])

snowflake/ml/_internal/type_utils.py

@@ -1,4 +1,4 @@
-import sys
+import importlib
 from typing import Any, Generic, Type, TypeVar, Union, cast
 
 import numpy as np
@@ -51,8 +51,8 @@ class LazyType(Generic[T]):
     def get_class(self) -> Type[T]:
         if self._runtime_class is None:
             try:
-                m = sys.modules[self.module]
-            except KeyError:
+                m = importlib.import_module(self.module)
+            except ModuleNotFoundError:
                 raise ValueError(f"Module {self.module} not imported.")
 
             self._runtime_class = cast("Type[T]", getattr(m, self.qualname))

snowflake/ml/_internal/utils/jwt_generator.py

@@ -0,0 +1,141 @@
+import base64
+import hashlib
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+import jwt
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import types
+
+logger = logging.getLogger(__name__)
+
+ISSUER = "iss"
+EXPIRE_TIME = "exp"
+ISSUE_TIME = "iat"
+SUBJECT = "sub"
+
+
+class JWTGenerator:
+    """
+    Creates and signs a JWT with the specified private key file, username, and account identifier. The JWTGenerator
+    keeps the generated token and only regenerates the token if a specified period of time has passed.
+    """
+
+    _DEFAULT_LIFETIME = timedelta(minutes=59)  # The tokens will have a 59-minute lifetime
+    _DEFAULT_RENEWAL_DELTA = timedelta(minutes=54)  # Tokens will be renewed after 54 minutes
+    ALGORITHM = "RS256"  # Tokens will be generated using RSA with SHA256
+
+    def __init__(
+        self,
+        account: str,
+        user: str,
+        private_key: types.PRIVATE_KEY_TYPES,
+        lifetime: Optional[timedelta] = None,
+        renewal_delay: Optional[timedelta] = None,
+    ) -> None:
+        """
+        Create a new JWTGenerator object.
+
+        Args:
+            account: The account identifier.
+            user: The username.
+            private_key: The private key used to sign the JWT.
+            lifetime: The lifetime of the token.
+            renewal_delay: The time before the token expires to renew it.
+        """
+
+        # Construct the fully qualified name of the user in uppercase.
+        self.account = JWTGenerator._prepare_account_name_for_jwt(account)
+        self.user = user.upper()
+        self.qualified_username = self.account + "." + self.user
+        self.private_key = private_key
+        self.public_key_fp = JWTGenerator._calculate_public_key_fingerprint(self.private_key)
+
+        self.issuer = self.qualified_username + "." + self.public_key_fp
+        self.lifetime = lifetime or JWTGenerator._DEFAULT_LIFETIME
+        self.renewal_delay = renewal_delay or JWTGenerator._DEFAULT_RENEWAL_DELTA
+        self.renew_time = datetime.now(timezone.utc)
+        self.token: Optional[str] = None
+
+        logger.info(
+            """Creating JWTGenerator with arguments
+            account : %s, user : %s, lifetime : %s, renewal_delay : %s""",
+            self.account,
+            self.user,
+            self.lifetime,
+            self.renewal_delay,
+        )
+
+    @staticmethod
+    def _prepare_account_name_for_jwt(raw_account: str) -> str:
+        account = raw_account
+        if ".global" not in account:
+            # Handle the general case.
+            idx = account.find(".")
+            if idx > 0:
+                account = account[0:idx]
+        else:
+            # Handle the replication case.
+            idx = account.find("-")
+            if idx > 0:
+                account = account[0:idx]
+        # Use uppercase for the account identifier.
+        return account.upper()
+
+    def get_token(self) -> str:
+        now = datetime.now(timezone.utc)  # Fetch the current time
+        if self.token is not None and self.renew_time > now:
+            return self.token
+
+        # If the token has expired or doesn't exist, regenerate the token.
+        logger.info(
+            "Generating a new token because the present time (%s) is later than the renewal time (%s)",
+            now,
+            self.renew_time,
+        )
+        # Calculate the next time we need to renew the token.
+        self.renew_time = now + self.renewal_delay
+
+        # Create our payload
+        payload = {
+            # Set the issuer to the fully qualified username concatenated with the public key fingerprint.
+            ISSUER: self.issuer,
+            # Set the subject to the fully qualified username.
+            SUBJECT: self.qualified_username,
+            # Set the issue time to now.
+            ISSUE_TIME: now,
+            # Set the expiration time, based on the lifetime specified for this object.
+            EXPIRE_TIME: now + self.lifetime,
+        }
+
+        # Regenerate the actual token
+        token = jwt.encode(payload, key=self.private_key, algorithm=JWTGenerator.ALGORITHM)
+        # If you are using a version of PyJWT prior to 2.0, jwt.encode returns a byte string instead of a string.
+        # If the token is a byte string, convert it to a string.
+        if isinstance(token, bytes):
+            token = token.decode("utf-8")
+        self.token = token
+        logger.info(
+            "Generated a JWT with the following payload: %s",
+            jwt.decode(self.token, key=self.private_key.public_key(), algorithms=[JWTGenerator.ALGORITHM]),
+        )
+
+        return token
+
+    @staticmethod
+    def _calculate_public_key_fingerprint(private_key: types.PRIVATE_KEY_TYPES) -> str:
+        # Get the raw bytes of public key.
+        public_key_raw = private_key.public_key().public_bytes(
+            serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+
+        # Get the sha256 hash of the raw bytes.
+        sha256hash = hashlib.sha256()
+        sha256hash.update(public_key_raw)
+
+        # Base64-encode the value and prepend the prefix 'SHA256:'.
+        public_key_fp = "SHA256:" + base64.b64encode(sha256hash.digest()).decode("utf-8")
+        logger.info("Public key fingerprint is %s", public_key_fp)
+
+        return public_key_fp

snowflake/ml/data/__init__.py

@@ -0,0 +1,5 @@
+from .data_connector import DataConnector
+from .data_ingestor import DataIngestor, DataIngestorType
+from .data_source import DataFrameInfo, DatasetInfo, DataSource
+
+__all__ = ["DataConnector", "DataSource", "DataFrameInfo", "DatasetInfo", "DataIngestor", "DataIngestorType"]

snowflake/ml/model/_client/model/model_version_impl.py

@@ -14,7 +14,7 @@ from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._packager.model_handlers import snowmlmodel
-from snowflake.snowpark import Session, dataframe
+from snowflake.snowpark import Session, async_job, dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
@@ -631,7 +631,8 @@ class ModelVersion(lineage_node.LineageNode):
         max_batch_rows: Optional[int] = None,
         force_rebuild: bool = False,
         build_external_access_integration: Optional[str] = None,
-    ) -> str:
+        block: bool = True,
+    ) -> Union[str, async_job.AsyncJob]:
         """Create an inference service with the given spec.
 
         Args:
@@ -659,6 +660,9 @@ class ModelVersion(lineage_node.LineageNode):
             force_rebuild: Whether to force a model inference image rebuild.
             build_external_access_integration: (Deprecated) The external access integration for image build. This is
                 usually permitting access to conda & PyPI repositories.
+            block: A bool value indicating whether this function will wait until the service is available.
+                When it is ``False``, this function executes the underlying service creation asynchronously
+                and returns an :class:`AsyncJob`.
         """
         ...
 
@@ -679,7 +683,8 @@ class ModelVersion(lineage_node.LineageNode):
         max_batch_rows: Optional[int] = None,
         force_rebuild: bool = False,
         build_external_access_integrations: Optional[List[str]] = None,
-    ) -> str:
+        block: bool = True,
+    ) -> Union[str, async_job.AsyncJob]:
         """Create an inference service with the given spec.
 
         Args:
@@ -707,6 +712,9 @@ class ModelVersion(lineage_node.LineageNode):
             force_rebuild: Whether to force a model inference image rebuild.
             build_external_access_integrations: The external access integrations for image build. This is usually
                 permitting access to conda & PyPI repositories.
+            block: A bool value indicating whether this function will wait until the service is available.
+                When it is ``False``, this function executes the underlying service creation asynchronously
+                and returns an :class:`AsyncJob`.
         """
         ...
 
@@ -742,7 +750,8 @@ class ModelVersion(lineage_node.LineageNode):
         force_rebuild: bool = False,
         build_external_access_integration: Optional[str] = None,
         build_external_access_integrations: Optional[List[str]] = None,
-    ) -> str:
+        block: bool = True,
+    ) -> Union[str, async_job.AsyncJob]:
         """Create an inference service with the given spec.
 
         Args:
@@ -772,12 +781,16 @@ class ModelVersion(lineage_node.LineageNode):
                 usually permitting access to conda & PyPI repositories.
             build_external_access_integrations: The external access integrations for image build. This is usually
                 permitting access to conda & PyPI repositories.
+            block: A bool value indicating whether this function will wait until the service is available.
+                When it is False, this function executes the underlying service creation asynchronously
+                and returns an AsyncJob.
 
         Raises:
             ValueError: Illegal external access integration arguments.
 
         Returns:
-            Result information about service creation from server.
+            If `block=True`, return result information about service creation from server.
+            Otherwise, return the service creation AsyncJob.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
@@ -829,6 +842,7 @@ class ModelVersion(lineage_node.LineageNode):
                 if build_external_access_integrations is None
                 else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
             ),
+            block=block,
             statement_params=statement_params,
         )
 
@@ -851,17 +865,13 @@ class ModelVersion(lineage_node.LineageNode):
         )
 
         return pd.DataFrame(
-            self._model_ops.list_inference_services(
+            self._model_ops.show_services(
                 database_name=None,
                 schema_name=None,
                 model_name=self._model_name,
                 version_name=self._version_name,
                 statement_params=statement_params,
-            ),
-            columns=[
-                self._model_ops.INFERENCE_SERVICE_NAME_COL_NAME,
-                self._model_ops.INFERENCE_SERVICE_ENDPOINT_COL_NAME,
-            ],
+            )
         )
 
     @telemetry.send_api_usage_telemetry(
@@ -889,12 +899,16 @@ class ModelVersion(lineage_node.LineageNode):
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
+
+        database_name_id, schema_name_id, service_name_id = sql_identifier.parse_fully_qualified_name(service_name)
         self._model_ops.delete_service(
             database_name=None,
             schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
-            service_name=service_name,
+            service_database_name=database_name_id,
+            service_schema_name=schema_name_id,
+            service_name=service_name_id,
             statement_params=statement_params,
         )
 

snowflake/ml/model/_client/ops/model_ops.py

@@ -3,7 +3,7 @@ import os
 import pathlib
 import tempfile
 import warnings
-from typing import Any, Dict, List, Literal, Optional, Union, cast, overload
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union, cast, overload
 
 import yaml
 
@@ -31,9 +31,14 @@ from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
+class ServiceInfo(TypedDict):
+    name: str
+    inference_endpoint: Optional[str]
+
+
 class ModelOperator:
-    INFERENCE_SERVICE_NAME_COL_NAME = "service_name"
-    INFERENCE_SERVICE_ENDPOINT_COL_NAME = "endpoints"
+    INFERENCE_SERVICE_ENDPOINT_NAME = "inference"
+    INGRESS_ENDPOINT_URL_SUFFIX = "snowflakecomputing.app"
 
     def __init__(
         self,
@@ -517,7 +522,7 @@ class ModelOperator:
             statement_params=statement_params,
         )
 
-    def list_inference_services(
+    def show_services(
         self,
         *,
         database_name: Optional[sql_identifier.SqlIdentifier],
@@ -525,7 +530,7 @@ class ModelOperator:
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, List[str]]:
+    ) -> List[ServiceInfo]:
         res = self._model_client.show_versions(
             database_name=database_name,
             schema_name=schema_name,
@@ -546,21 +551,28 @@ class ModelOperator:
 
         json_array = json.loads(res[0][service_col_name])
         # TODO(sdas): Figure out a better way to filter out MODEL_BUILD_ services server side.
-        services = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
-        endpoint_col_name = self._model_client.MODEL_INFERENCE_SERVICE_ENDPOINT_COL_NAME
-
-        services_col, endpoints_col = [], []
-        for service in services:
-            res = self._model_client.show_endpoints(service_name=service)
-            endpoints = [endpoint[endpoint_col_name] for endpoint in res]
-            for endpoint in endpoints:
-                services_col.append(service)
-                endpoints_col.append(endpoint)
-
-        return {
-            self.INFERENCE_SERVICE_NAME_COL_NAME: services_col,
-            self.INFERENCE_SERVICE_ENDPOINT_COL_NAME: endpoints_col,
-        }
+        fully_qualified_service_names = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
+
+        result = []
+        ingress_url: Optional[str] = None
+        for fully_qualified_service_name in fully_qualified_service_names:
+            db, schema, service_name = sql_identifier.parse_fully_qualified_name(fully_qualified_service_name)
+            for res_row in self._service_client.show_endpoints(
+                database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
+            ):
+                if (
+                    res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME]
+                    == self.INFERENCE_SERVICE_ENDPOINT_NAME
+                    and res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME] is not None
+                ):
+                    ingress_url = str(
+                        res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME]
+                    )
+                    if not ingress_url.endswith(ModelOperator.INGRESS_ENDPOINT_URL_SUFFIX):
+                        ingress_url = None
+            result.append(ServiceInfo(name=fully_qualified_service_name, inference_endpoint=ingress_url))
+
+        return result
 
     def delete_service(
         self,
@@ -569,33 +581,42 @@ class ModelOperator:
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        service_name: str,
+        service_database_name: Optional[sql_identifier.SqlIdentifier],
+        service_schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        services = self.list_inference_services(
+        services = self.show_services(
             database_name=database_name,
             schema_name=schema_name,
             model_name=model_name,
             version_name=version_name,
             statement_params=statement_params,
         )
-        db, schema, service_name = sql_identifier.parse_fully_qualified_name(service_name)
+
+        # Fall back to the model's database and schema.
+        # database_name or schema_name are set if the model is created or get using fully qualified name
+        # Otherwise, the model's database and schema are same as registry's database and schema, which are set in the
+        # self._model_client.
+
+        service_database_name = service_database_name or database_name or self._model_client._database_name
+        service_schema_name = service_schema_name or schema_name or self._model_client._schema_name
         fully_qualified_service_name = sql_identifier.get_fully_qualified_name(
-            db, schema, service_name, self._session.get_current_database(), self._session.get_current_schema()
+            service_database_name, service_schema_name, service_name
         )
 
-        service_col_name = self.INFERENCE_SERVICE_NAME_COL_NAME
-        for service in services[service_col_name]:
-            if service == fully_qualified_service_name:
+        for service_info in services:
+            if service_info["name"] == fully_qualified_service_name:
                 self._service_client.drop_service(
-                    database_name=db,
-                    schema_name=schema,
+                    database_name=service_database_name,
+                    schema_name=service_schema_name,
                     service_name=service_name,
                     statement_params=statement_params,
                 )
                 return
         raise ValueError(
-            f"Service '{service_name}' does not exist or unauthorized or not associated with this model version."
+            f"Service '{fully_qualified_service_name}' does not exist "
+            "or unauthorized or not associated with this model version."
         )
 
     def get_model_version_manifest(