snowflake-ml-python 1.7.0__py3-none-any.whl → 1.7.1__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/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

@@ -851,17 +851,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 +885,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(

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

@@ -109,6 +109,17 @@ class ServiceOperator:
         build_external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> str:
+
+        # Fall back to the registry's database and schema if not provided
+        database_name = database_name or self._database_name
+        schema_name = schema_name or self._schema_name
+
+        # Fall back to the model's database and schema if not provided then to the registry's database and schema
+        service_database_name = service_database_name or database_name or self._database_name
+        service_schema_name = service_schema_name or schema_name or self._schema_name
+
+        image_repo_database_name = image_repo_database_name or database_name or self._database_name
+        image_repo_schema_name = image_repo_schema_name or schema_name or self._schema_name
         # create a temp stage
         stage_name = sql_identifier.SqlIdentifier(
             snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)
@@ -130,8 +141,8 @@ class ServiceOperator:
             raise ValueError("External access integrations are required in Snowflake < 8.40.0.")
 
         self._model_deployment_spec.save(
-            database_name=database_name or self._database_name,
-            schema_name=schema_name or self._schema_name,
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             version_name=version_name,
             service_database_name=service_database_name,

snowflake/ml/model/_client/sql/model.py

@@ -17,8 +17,6 @@ class ModelSQLClient(_base._BaseSQLClient):
     MODEL_VERSION_ALIASES_COL_NAME = "aliases"
     MODEL_VERSION_INFERENCE_SERVICES_COL_NAME = "inference_services"
 
-    MODEL_INFERENCE_SERVICE_ENDPOINT_COL_NAME = "name"
-
     def show_models(
         self,
         *,
@@ -85,18 +83,6 @@ class ModelSQLClient(_base._BaseSQLClient):
 
         return res.validate()
 
-    def show_endpoints(
-        self,
-        *,
-        service_name: str,
-    ) -> List[row.Row]:
-        res = query_result_checker.SqlResultValidator(
-            self._session,
-            (f"SHOW ENDPOINTS IN SERVICE {service_name}"),
-        ).has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
-
-        return res.validate()
-
     def set_comment(
         self,
         *,

snowflake/ml/model/_client/sql/service.py

@@ -10,7 +10,7 @@ from snowflake.ml._internal.utils import (
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
-from snowflake.snowpark import dataframe, functions as F, types as spt
+from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -26,6 +26,9 @@ class ServiceStatus(enum.Enum):
 
 
 class ServiceSQLClient(_base._BaseSQLClient):
+    MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME = "name"
+    MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME = "ingress_url"
+
     def build_model_container(
         self,
         *,
@@ -216,3 +219,24 @@ class ServiceSQLClient(_base._BaseSQLClient):
             f"DROP SERVICE {self.fully_qualified_object_name(database_name, schema_name, service_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def show_endpoints(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[row.Row]:
+        fully_qualified_service_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
+        res = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                (f"SHOW ENDPOINTS IN SERVICE {fully_qualified_service_name}"),
+                statement_params=statement_params,
+            )
+            .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME, allow_empty=True)
+            .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME, allow_empty=True)
+        )
+
+        return res.validate()

snowflake/ml/model/_model_composer/model_method/infer_function.py_template

@@ -5,6 +5,7 @@ import sys
 
 import anyio
 import pandas as pd
+import numpy as np
 from _snowflake import vectorized
 
 from snowflake.ml.model._packager import model_packager
@@ -47,4 +48,4 @@ def {function_name}(df: pd.DataFrame) -> dict:
     df.columns = input_cols
     input_df = df.astype(dtype=dtype_map)
     predictions_df = runner(input_df[input_cols])
-    return predictions_df.to_dict("records")
+    return predictions_df.replace({{pd.NA: None, np.nan: None}}).to_dict("records")

snowflake/ml/model/_packager/model_env/model_env.py

@@ -174,6 +174,18 @@ class ModelEnv:
             except env_utils.DuplicateDependencyError:
                 pass
 
+    def remove_if_present_conda(self, conda_pkgs: List[str]) -> None:
+        """Remove conda requirements from model env if present.
+
+        Args:
+            conda_pkgs: A list of package name to be removed from conda requirements.
+        """
+        for pkg_name in conda_pkgs:
+            spec_conda = env_utils._find_conda_dep_spec(self._conda_dependencies, pkg_name)
+            if spec_conda:
+                channel, spec = spec_conda
+                self._conda_dependencies[channel].remove(spec)
+
     def generate_env_for_cuda(self) -> None:
         if self.cuda_version is None:
             return

snowflake/ml/model/_packager/model_handlers/_utils.py

@@ -179,7 +179,7 @@ def convert_explanations_to_2D_df(
         return pd.DataFrame(explanations)
 
     if hasattr(model, "classes_"):
-        classes_list = [str(cl) for cl in model.classes_]  # type:ignore[union-attr]
+        classes_list = [str(cl) for cl in model.classes_]
         len_classes = len(classes_list)
         if explanations.shape[2] != len_classes:
             raise ValueError(f"Model has {len_classes} classes but explanations have {explanations.shape[2]}")

snowflake/ml/model/_packager/model_handlers/catboost.py

@@ -95,7 +95,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 get_prediction_fn=get_prediction,
             )
             model_task_and_output = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output.task
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
             if enable_explainability:
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(

snowflake/ml/model/_packager/model_handlers/custom.py

@@ -2,7 +2,7 @@ import inspect
 import os
 import pathlib
 import sys
-from typing import Dict, Optional, Type, final
+from typing import Dict, Optional, Type, cast, final
 
 import anyio
 import cloudpickle
@@ -108,6 +108,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                     model_meta=model_meta,
                     model_blobs_dir_path=model_blobs_dir_path,
                     is_sub_model=True,
+                    **cast(model_types.BaseModelSaveOption, kwargs),
                 )
 
         # Make sure that the module where the model is defined get pickled by value as well.
@@ -175,6 +176,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 name=sub_model_name,
                 model_meta=model_meta,
                 model_blobs_dir_path=model_blobs_dir_path,
+                **cast(model_types.BaseModelLoadOption, kwargs),
             )
             models[sub_model_name] = sub_model
         reconstructed_context = custom_model.ModelContext(artifacts=artifacts, models=models)

snowflake/ml/model/_packager/model_handlers/lightgbm.py

@@ -196,13 +196,14 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         with open(model_blob_file_path, "rb") as f:
             model = cloudpickle.load(f)
         assert isinstance(model, getattr(lightgbm, lightgbm_estimator_type))
+        assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
 
         return model
 
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
+        raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.LGBMModelLoadOptions],