snowflake-ml-python 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

snowflake/ml/model/_api.py

@@ -7,7 +7,6 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import (
     deploy_platforms,
     model_signature,
@@ -188,6 +187,10 @@ def save_model(
     Returns:
         Model
     """
+    if options is None:
+        options = {}
+    options["_legacy_save"] = True
+
     m = model_composer.ModelComposer(session=session, stage_path=stage_path)
     m.save(
         name=name,
@@ -481,6 +484,7 @@ def predict(
     # Get options
     INTERMEDIATE_OBJ_NAME = "tmp_result"
     sig = deployment["signature"]
+    identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
     # Validate and prepare input
     if not isinstance(X, SnowparkDataFrame):
@@ -491,7 +495,7 @@ def predict(
     else:
         keep_order = False
         output_with_input_features = True
-        model_signature._validate_snowpark_data(X, sig.inputs)
+        identifier_rule = model_signature._validate_snowpark_data(X, sig.inputs)
         s_df = X
 
     if statement_params:
@@ -500,10 +504,14 @@ def predict(
         else:
             s_df._statement_params = statement_params  # type: ignore[assignment]
 
+    original_cols = s_df.columns
+
     # Infer and get intermediate result
     input_cols = []
-    for col_name in s_df.columns:
-        literal_col_name = identifier.get_unescaped_names(col_name)
+    for input_feature in sig.inputs:
+        literal_col_name = input_feature.name
+        col_name = identifier_rule.get_identifier_from_feature(input_feature.name)
+
         input_cols.extend(
             [
                 F.lit(literal_col_name),
@@ -511,29 +519,28 @@ def predict(
             ]
         )
 
-    # TODO[shchen]: SNOW-870032, For SnowService, external function name cannot be double quoted, else it results in
-    # external function no found.
     udf_name = deployment["name"]
-    output_obj = F.call_udf(udf_name, F.object_construct(*input_cols))
-
-    if output_with_input_features:
-        df_res = s_df.with_column(INTERMEDIATE_OBJ_NAME, output_obj)
-    else:
-        df_res = s_df.select(output_obj.alias(INTERMEDIATE_OBJ_NAME))
+    output_obj = F.call_udf(udf_name, F.object_construct_keep_null(*input_cols))
+    df_res = s_df.with_column(INTERMEDIATE_OBJ_NAME, output_obj)
 
     if keep_order:
         df_res = df_res.order_by(
-            F.col(INTERMEDIATE_OBJ_NAME)[infer_template._KEEP_ORDER_COL_NAME],
+            F.col(infer_template._KEEP_ORDER_COL_NAME),
             ascending=True,
         )
 
+    if not output_with_input_features:
+        df_res = df_res.drop(*original_cols)
+
     # Prepare the output
     output_cols = []
+    output_col_names = []
     for output_feature in sig.outputs:
         output_cols.append(F.col(INTERMEDIATE_OBJ_NAME)[output_feature.name].astype(output_feature.as_snowpark_type()))
+        output_col_names.append(identifier_rule.get_identifier_from_feature(output_feature.name))
 
     df_res = df_res.with_columns(
-        [identifier.get_inferred_name(output_feature.name) for output_feature in sig.outputs],
+        output_col_names,
         output_cols,
     ).drop(INTERMEDIATE_OBJ_NAME)
 

snowflake/ml/model/_client/model/model_impl.py

@@ -0,0 +1,176 @@
+from typing import List, Union
+
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._client.model import model_version_impl
+from snowflake.ml.model._client.ops import model_ops
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
+
+class Model:
+    """Model Object containing multiple versions. Mapping to SQL's MODEL object."""
+
+    _model_ops: model_ops.ModelOperator
+    _model_name: sql_identifier.SqlIdentifier
+
+    def __init__(self) -> None:
+        raise RuntimeError("Model's initializer is not meant to be used. Use `get_model` from registry instead.")
+
+    @classmethod
+    def _ref(
+        cls,
+        model_ops: model_ops.ModelOperator,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+    ) -> "Model":
+        self: "Model" = object.__new__(cls)
+        self._model_ops = model_ops
+        self._model_name = model_name
+        return self
+
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, Model):
+            return False
+        return self._model_ops == __value._model_ops and self._model_name == __value._model_name
+
+    @property
+    def name(self) -> str:
+        return self._model_name.identifier()
+
+    @property
+    def fully_qualified_name(self) -> str:
+        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+
+    @property
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def description(self) -> str:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.get_comment(
+            model_name=self._model_name,
+            statement_params=statement_params,
+        )
+
+    @description.setter
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def description(self, description: str) -> None:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.set_comment(
+            comment=description,
+            model_name=self._model_name,
+            statement_params=statement_params,
+        )
+
+    @property
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def default(self) -> model_version_impl.ModelVersion:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+            class_name=self.__class__.__name__,
+        )
+        default_version_name = self._model_ops._model_version_client.get_default_version(
+            model_name=self._model_name, statement_params=statement_params
+        )
+        return self.version(default_version_name)
+
+    @default.setter
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def default(self, version: Union[str, model_version_impl.ModelVersion]) -> None:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+            class_name=self.__class__.__name__,
+        )
+        if isinstance(version, str):
+            version_name = sql_identifier.SqlIdentifier(version)
+        else:
+            version_name = version._version_name
+        self._model_ops._model_version_client.set_default_version(
+            model_name=self._model_name, version_name=version_name, statement_params=statement_params
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def version(self, version_name: str) -> model_version_impl.ModelVersion:
+        """Get a model version object given a version name in the model.
+
+        Args:
+            version_name: The name of version
+
+        Raises:
+            ValueError: Raised when the version requested does not exist.
+
+        Returns:
+            The model version object.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        version_id = sql_identifier.SqlIdentifier(version_name)
+        if self._model_ops.validate_existence(
+            model_name=self._model_name,
+            version_name=version_id,
+            statement_params=statement_params,
+        ):
+            return model_version_impl.ModelVersion._ref(
+                self._model_ops,
+                model_name=self._model_name,
+                version_name=version_id,
+            )
+        else:
+            raise ValueError(
+                f"Unable to find version with name {version_id.identifier()} in model {self.fully_qualified_name}"
+            )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def list_versions(self) -> List[model_version_impl.ModelVersion]:
+        """List all versions in the model.
+
+        Returns:
+            A List of ModelVersion object representing all versions in the model.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        version_names = self._model_ops.list_models_or_versions(
+            model_name=self._model_name,
+            statement_params=statement_params,
+        )
+        return [
+            model_version_impl.ModelVersion._ref(
+                self._model_ops,
+                model_name=self._model_name,
+                version_name=version_name,
+            )
+            for version_name in version_names
+        ]
+
+    def delete_version(self, version_name: str) -> None:
+        raise NotImplementedError("Deleting version has not been supported yet.")

snowflake/ml/model/_client/model/model_method_info.py

@@ -0,0 +1,19 @@
+from typing import TypedDict
+
+from typing_extensions import Required
+
+from snowflake.ml.model import model_signature
+
+
+class ModelMethodInfo(TypedDict):
+    """Method information.
+
+    Attributes:
+        name: Name of the method to be called via SQL.
+        target_method: actual target method name to be called.
+        signature: The signature of the model method.
+    """
+
+    name: Required[str]
+    target_method: Required[str]
+    signature: Required[model_signature.ModelSignature]

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

@@ -0,0 +1,291 @@
+import re
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import pandas as pd
+
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model import model_signature
+from snowflake.ml.model._client.model import model_method_info
+from snowflake.ml.model._client.ops import metadata_ops, model_ops
+from snowflake.snowpark import dataframe
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
+
+class ModelVersion:
+    """Model Version Object representing a specific version of the model that could be run."""
+
+    _model_ops: model_ops.ModelOperator
+    _model_name: sql_identifier.SqlIdentifier
+    _version_name: sql_identifier.SqlIdentifier
+
+    def __init__(self) -> None:
+        raise RuntimeError("ModelVersion's initializer is not meant to be used. Use `version` from model instead.")
+
+    @classmethod
+    def _ref(
+        cls,
+        model_ops: model_ops.ModelOperator,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+    ) -> "ModelVersion":
+        self: "ModelVersion" = object.__new__(cls)
+        self._model_ops = model_ops
+        self._model_name = model_name
+        self._version_name = version_name
+        return self
+
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, ModelVersion):
+            return False
+        return (
+            self._model_ops == __value._model_ops
+            and self._model_name == __value._model_name
+            and self._version_name == __value._version_name
+        )
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name.identifier()
+
+    @property
+    def version_name(self) -> str:
+        return self._version_name.identifier()
+
+    @property
+    def fully_qualified_model_name(self) -> str:
+        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+
+    @property
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def description(self) -> str:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.get_comment(
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @description.setter
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def description(self, description: str) -> None:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.set_comment(
+            comment=description,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def list_metrics(self) -> Dict[str, Any]:
+        """Show all metrics logged with the model version.
+
+        Returns:
+            A dictionary showing the metrics
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops._metadata_ops.load(
+            model_name=self._model_name, version_name=self._version_name, statement_params=statement_params
+        )["metrics"]
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def get_metric(self, metric_name: str) -> Any:
+        """Get the value of a specific metric.
+
+        Args:
+            metric_name: The name of the metric
+
+        Raises:
+            KeyError: Raised when the requested metric name does not exist.
+
+        Returns:
+            The value of the metric.
+        """
+        metrics = self.list_metrics()
+        if metric_name not in metrics:
+            raise KeyError(f"Cannot find metric with name {metric_name}.")
+        return metrics[metric_name]
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def set_metric(self, metric_name: str, value: Any) -> None:
+        """Set the value of a specific metric name
+
+        Args:
+            metric_name: The name of the metric
+            value: The value of the metric.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        metrics = self.list_metrics()
+        metrics[metric_name] = value
+        self._model_ops._metadata_ops.save(
+            metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def delete_metric(self, metric_name: str) -> None:
+        """Delete a metric from metric storage.
+
+        Args:
+            metric_name: The name of the metric to be deleted.
+
+        Raises:
+            KeyError: Raised when the requested metric name does not exist.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        metrics = self.list_metrics()
+        if metric_name not in metrics:
+            raise KeyError(f"Cannot find metric with name {metric_name}.")
+        del metrics[metric_name]
+        self._model_ops._metadata_ops.save(
+            metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def list_methods(self) -> List[model_method_info.ModelMethodInfo]:
+        """List all method information in a model version that is callable.
+
+        Returns:
+            A list of ModelMethodInfo object containing the following information:
+            - name: The name of the method to be called (both in SQL and in Python SDK).
+            - target_method: The original method name in the logged Python object.
+            - Signature: Python signature of the original method.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        # TODO(SNOW-986673, SNOW-986675): Avoid parsing manifest and meta file and put Python signature into user_data.
+        manifest = self._model_ops.get_model_version_manifest(
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+        model_meta = self._model_ops.get_model_version_native_packing_meta(
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+        return_methods_info: List[model_method_info.ModelMethodInfo] = []
+        for method in manifest["methods"]:
+            # Method's name is resolved so we need to use case_sensitive as True to get the user-facing identifier.
+            method_name = sql_identifier.SqlIdentifier(method["name"], case_sensitive=True).identifier()
+            # Method's handler is `functions.<target_method>.infer`
+            assert re.match(
+                r"^functions\.([^\d\W]\w*)\.infer$", method["handler"]
+            ), f"Get unexpected handler name {method['handler']}"
+            target_method = method["handler"].split(".")[1]
+            signature_dict = model_meta["signatures"][target_method]
+            method_info = model_method_info.ModelMethodInfo(
+                name=method_name,
+                target_method=target_method,
+                signature=model_signature.ModelSignature.from_dict(signature_dict),
+            )
+            return_methods_info.append(method_info)
+
+        return return_methods_info
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def run(
+        self,
+        X: Union[pd.DataFrame, dataframe.DataFrame],
+        *,
+        method_name: Optional[str] = None,
+    ) -> Union[pd.DataFrame, dataframe.DataFrame]:
+        """Invoke a method in a model version object
+
+        Args:
+            X: The input data. Could be pandas DataFrame or Snowpark DataFrame
+            method_name: The method name to run. It is the name you will use to call a method in SQL. Defaults to None.
+                It can only be None if there is only 1 method.
+
+        Raises:
+            ValueError: No method with the corresponding name is available.
+            ValueError: There are more than 1 target methods available in the model but no method name specified.
+
+        Returns:
+            The prediction data.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+
+        methods: List[model_method_info.ModelMethodInfo] = self.list_methods()
+        if method_name:
+            req_method_name = sql_identifier.SqlIdentifier(method_name).identifier()
+            find_method: Callable[[model_method_info.ModelMethodInfo], bool] = (
+                lambda method: method["name"] == req_method_name
+            )
+            target_method_info = next(
+                filter(find_method, methods),
+                None,
+            )
+            if target_method_info is None:
+                raise ValueError(
+                    f"There is no method with name {method_name} available in the model"
+                    f" {self.fully_qualified_model_name} version {self.version_name}"
+                )
+        elif len(methods) != 1:
+            raise ValueError(
+                f"There are more than 1 target methods available in the model {self.fully_qualified_model_name}"
+                f" version {self.version_name}. Please specify a `method_name` when calling the `run` method."
+            )
+        else:
+            target_method_info = methods[0]
+        return self._model_ops.invoke_method(
+            method_name=sql_identifier.SqlIdentifier(target_method_info["name"]),
+            signature=target_method_info["signature"],
+            X=X,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )

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

@@ -0,0 +1,107 @@
+import json
+from typing import Any, Dict, Optional, TypedDict
+
+from typing_extensions import NotRequired
+
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._client.sql import (
+    model as model_sql,
+    model_version as model_version_sql,
+)
+from snowflake.snowpark import session
+
+MODEL_VERSION_METADATA_SCHEMA_VERSION = "2024-01-01"
+
+
+class ModelVersionMetadataSchema(TypedDict):
+    metrics: NotRequired[Dict[str, Any]]
+
+
+class MetadataOperator:
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._model_client = model_sql.ModelSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
+        self._model_version_client = model_version_sql.ModelVersionSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
+
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, MetadataOperator):
+            return False
+        return (
+            self._model_client == __value._model_client and self._model_version_client == __value._model_version_client
+        )
+
+    @staticmethod
+    def _parse(metadata_dict: Dict[str, Any]) -> ModelVersionMetadataSchema:
+        loaded_metadata_schema_version = metadata_dict.get("snowpark_ml_schema_version", None)
+        if loaded_metadata_schema_version is None:
+            return ModelVersionMetadataSchema(metrics={})
+        elif (
+            not isinstance(loaded_metadata_schema_version, str)
+            or loaded_metadata_schema_version != MODEL_VERSION_METADATA_SCHEMA_VERSION
+        ):
+            raise ValueError(f"Unsupported model metadata schema version {loaded_metadata_schema_version} confronted.")
+        loaded_metrics = metadata_dict.get("metrics", {})
+        if not isinstance(loaded_metrics, dict):
+            raise ValueError(f"Metrics in the metadata is expected to be a dictionary, getting {loaded_metrics}")
+        return ModelVersionMetadataSchema(metrics=loaded_metrics)
+
+    def _get_current_metadata_dict(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        version_info_list = self._model_client.show_versions(
+            model_name=model_name, version_name=version_name, statement_params=statement_params
+        )
+        assert len(version_info_list) == 1
+        version_info = version_info_list[0]
+        metadata_str = version_info.metadata
+        if not metadata_str:
+            return {}
+        res = json.loads(metadata_str)
+        if not isinstance(res, dict):
+            raise ValueError(f"Metadata is expected to be a dictionary, getting {res}")
+        return res
+
+    def load(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> ModelVersionMetadataSchema:
+        metadata_dict = self._get_current_metadata_dict(
+            model_name=model_name, version_name=version_name, statement_params=statement_params
+        )
+        return MetadataOperator._parse(metadata_dict)
+
+    def save(
+        self,
+        metadata: ModelVersionMetadataSchema,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        metadata_dict = self._get_current_metadata_dict(
+            model_name=model_name, version_name=version_name, statement_params=statement_params
+        )
+        metadata_dict.update({**metadata, "snowpark_ml_schema_version": MODEL_VERSION_METADATA_SCHEMA_VERSION})
+        self._model_version_client.set_metadata(
+            metadata_dict, model_name=model_name, version_name=version_name, statement_params=statement_params
+        )