snowflake-ml-python 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

snowflake/ml/model/_api.py

@@ -2,6 +2,7 @@ from types import ModuleType
 from typing import Any, Dict, List, Literal, Optional, Union, cast, overload
 
 import pandas as pd
+from typing_extensions import deprecated
 
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -23,6 +24,7 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import DataFrame as SnowparkDataFrame, Session, functions as F
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def save_model(
     *,
@@ -61,6 +63,7 @@ def save_model(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def save_model(
     *,
@@ -101,6 +104,7 @@ def save_model(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def save_model(
     *,
@@ -142,6 +146,7 @@ def save_model(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 def save_model(
     *,
     name: str,
@@ -208,6 +213,7 @@ def save_model(
     return m
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def load_model(*, session: Session, stage_path: str) -> model_composer.ModelComposer:
     """Load the model into memory from a zip file in the stage.
@@ -219,6 +225,7 @@ def load_model(*, session: Session, stage_path: str) -> model_composer.ModelComp
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def load_model(*, session: Session, stage_path: str, meta_only: Literal[False]) -> model_composer.ModelComposer:
     """Load the model into memory from a zip file in the stage.
@@ -231,6 +238,7 @@ def load_model(*, session: Session, stage_path: str, meta_only: Literal[False])
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def load_model(*, session: Session, stage_path: str, meta_only: Literal[True]) -> model_composer.ModelComposer:
     """Load the model into memory from a zip file in the stage with metadata only.
@@ -243,6 +251,7 @@ def load_model(*, session: Session, stage_path: str, meta_only: Literal[True]) -
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 def load_model(
     *,
     session: Session,
@@ -261,10 +270,11 @@ def load_model(
         Loaded model.
     """
     m = model_composer.ModelComposer(session=session, stage_path=stage_path)
-    m.load(meta_only=meta_only)
+    m.legacy_load(meta_only=meta_only)
     return m
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def deploy(
     session: Session,
@@ -290,6 +300,7 @@ def deploy(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def deploy(
     session: Session,
@@ -319,6 +330,7 @@ def deploy(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 def deploy(
     session: Session,
     *,
@@ -423,6 +435,7 @@ def deploy(
     return info
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def predict(
     session: Session,
@@ -443,6 +456,7 @@ def predict(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 @overload
 def predict(
     session: Session,
@@ -462,6 +476,7 @@ def predict(
     ...
 
 
+@deprecated("Only used by PrPr model registry.")
 def predict(
     session: Session,
     *,

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

@@ -1,9 +1,9 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Union
 
 import pandas as pd
 
 from snowflake.ml._internal import telemetry
-from snowflake.ml._internal.utils import identifier, sql_identifier
+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
 
@@ -45,7 +45,7 @@ class Model:
     @property
     def fully_qualified_name(self) -> str:
         """Return the fully qualified name of the model that can be used to refer to it in SQL."""
-        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+        return self._model_ops._model_version_client.fully_qualified_object_name(None, None, self._model_name)
 
     @property
     @telemetry.send_api_usage_telemetry(
@@ -76,6 +76,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_comment(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -92,6 +94,8 @@ class Model:
         )
         return self._model_ops.set_comment(
             comment=comment,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -109,7 +113,7 @@ class Model:
             class_name=self.__class__.__name__,
         )
         default_version_name = self._model_ops.get_default_version(
-            model_name=self._model_name, statement_params=statement_params
+            database_name=None, schema_name=None, model_name=self._model_name, statement_params=statement_params
         )
         return self.version(default_version_name)
 
@@ -129,7 +133,11 @@ class Model:
         else:
             version_name = version._version_name
         self._model_ops.set_default_version(
-            model_name=self._model_name, version_name=version_name, statement_params=statement_params
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
 
     @telemetry.send_api_usage_telemetry(
@@ -155,6 +163,8 @@ class Model:
         )
         version_id = sql_identifier.SqlIdentifier(version_name)
         if self._model_ops.validate_existence(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=version_id,
             statement_params=statement_params,
@@ -184,6 +194,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         version_names = self._model_ops.list_models_or_versions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -211,6 +223,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         rows = self._model_ops.show_models_or_versions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -231,6 +245,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         self._model_ops.delete_model_or_version(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=sql_identifier.SqlIdentifier(version_name),
             statement_params=statement_params,
@@ -250,29 +266,9 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        return self._model_ops.show_tags(model_name=self._model_name, statement_params=statement_params)
-
-    def _parse_tag_name(
-        self,
-        tag_name: str,
-    ) -> Tuple[sql_identifier.SqlIdentifier, sql_identifier.SqlIdentifier, sql_identifier.SqlIdentifier]:
-        _tag_db, _tag_schema, _tag_name, _ = identifier.parse_schema_level_object_identifier(tag_name)
-        if _tag_db is None:
-            tag_db_id = self._model_ops._model_client._database_name
-        else:
-            tag_db_id = sql_identifier.SqlIdentifier(_tag_db)
-
-        if _tag_schema is None:
-            tag_schema_id = self._model_ops._model_client._schema_name
-        else:
-            tag_schema_id = sql_identifier.SqlIdentifier(_tag_schema)
-
-        if _tag_name is None:
-            raise ValueError(f"Unable parse the tag name `{tag_name}` you input.")
-
-        tag_name_id = sql_identifier.SqlIdentifier(_tag_name)
-
-        return tag_db_id, tag_schema_id, tag_name_id
+        return self._model_ops.show_tags(
+            database_name=None, schema_name=None, model_name=self._model_name, statement_params=statement_params
+        )
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
@@ -292,8 +288,10 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         return self._model_ops.get_tag_value(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
@@ -317,8 +315,10 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         self._model_ops.set_tag(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
@@ -342,11 +342,45 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         self._model_ops.unset_tag(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
             tag_name=tag_name_id,
             statement_params=statement_params,
         )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def rename(self, model_name: str) -> None:
+        """Rename a model. Can be used to move a model when a fully qualified name is provided.
+
+        Args:
+            model_name: The new model name.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        new_db, new_schema, new_model = sql_identifier.parse_fully_qualified_name(model_name)
+
+        self._model_ops.rename(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            new_model_db=new_db,
+            new_model_schema=new_schema,
+            new_model_name=new_model,
+            statement_params=statement_params,
+        )
+        self._model_ops = model_ops.ModelOperator(
+            self._model_ops._session,
+            database_name=new_db or self._model_ops._model_client._database_name,
+            schema_name=new_schema or self._model_ops._model_client._schema_name,
+        )
+        self._model_name = new_model

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

@@ -1,17 +1,29 @@
+import enum
+import pathlib
+import tempfile
+import warnings
 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 type_hints as model_types
 from snowflake.ml.model._client.ops import metadata_ops, model_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 dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 
 
+class ExportMode(enum.Enum):
+    MODEL = "model"
+    FULL = "full"
+
+
 class ModelVersion:
     """Model Version Object representing a specific version of the model that could be run."""
 
@@ -60,7 +72,7 @@ class ModelVersion:
     @property
     def fully_qualified_model_name(self) -> str:
         """Return the fully qualified name of the model to which the model version belongs."""
-        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+        return self._model_ops._model_version_client.fully_qualified_object_name(None, None, self._model_name)
 
     @property
     @telemetry.send_api_usage_telemetry(
@@ -91,6 +103,8 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_comment(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -108,6 +122,8 @@ class ModelVersion:
         )
         return self._model_ops.set_comment(
             comment=comment,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -128,7 +144,11 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops._metadata_ops.load(
-            model_name=self._model_name, version_name=self._version_name, statement_params=statement_params
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
         )["metrics"]
 
     @telemetry.send_api_usage_telemetry(
@@ -171,6 +191,8 @@ class ModelVersion:
         metrics[metric_name] = value
         self._model_ops._metadata_ops.save(
             metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -199,6 +221,8 @@ class ModelVersion:
         del metrics[metric_name]
         self._model_ops._metadata_ops.save(
             metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -210,6 +234,8 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_functions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -240,6 +266,7 @@ class ModelVersion:
         X: Union[pd.DataFrame, dataframe.DataFrame],
         *,
         function_name: Optional[str] = None,
+        partition_column: Optional[str] = None,
         strict_input_validation: bool = False,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
         """Invoke a method in a model version object.
@@ -248,12 +275,14 @@ class ModelVersion:
             X: The input data, which could be a pandas DataFrame or Snowpark DataFrame.
             function_name: The function name to run. It is the name used to call a function in SQL.
                 Defaults to None. It can only be None if there is only 1 method.
+            partition_column: The partition column name to partition by.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
 
         Raises:
             ValueError: When no method with the corresponding name is available.
             ValueError: When there are more than 1 target methods available in the model but no function name specified.
+            ValueError: When the partition column is not a valid Snowflake identifier.
 
         Returns:
             The prediction data. It would be the same type dataframe as your input.
@@ -263,6 +292,10 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
+        if partition_column is not None:
+            # Partition column must be a valid identifier
+            partition_column = sql_identifier.SqlIdentifier(partition_column)
+
         functions: List[model_manifest_schema.ModelFunctionInfo] = self._functions
         if function_name:
             req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
@@ -287,10 +320,134 @@ class ModelVersion:
             target_function_info = functions[0]
         return self._model_ops.invoke_method(
             method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
+            method_function_type=target_function_info["target_method_function_type"],
             signature=target_function_info["signature"],
             X=X,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             strict_input_validation=strict_input_validation,
+            partition_column=partition_column,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["export_mode"]
+    )
+    def export(self, target_path: str, *, export_mode: ExportMode = ExportMode.MODEL) -> None:
+        """Export model files to a local directory.
+
+        Args:
+            target_path: Path to a local directory to export files to. A directory will be created if does not exist.
+            export_mode: The mode to export the model. Defaults to ExportMode.MODEL.
+                ExportMode.MODEL: All model files including environment to load the model and model weights.
+                ExportMode.FULL: Additional files to run the model in Warehouse, besides all files in MODEL mode,
+
+        Raises:
+            ValueError: Raised when the target path is a file or an non-empty folder.
+        """
+        target_local_path = pathlib.Path(target_path)
+        if target_local_path.is_file() or any(target_local_path.iterdir()):
+            raise ValueError(f"Target path {target_local_path} is a file or an non-empty folder.")
+
+        target_local_path.mkdir(parents=False, exist_ok=True)
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        self._model_ops.download_files(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            target_path=target_local_path,
+            mode=export_mode.value,
             statement_params=statement_params,
         )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["force", "options"]
+    )
+    def load(
+        self,
+        *,
+        force: bool = False,
+        options: Optional[model_types.ModelLoadOption] = None,
+    ) -> model_types.SupportedModelType:
+        """Load the underlying original Python object back from a model.
+            This operation requires to have the exact the same environment as the one when logging the model, otherwise,
+            the model might be not functional or some other problems might occur.
+
+        Args:
+            force: Bypass the best-effort environment validation. Defaults to False.
+            options: Options to specify when loading the model, check `snowflake.ml.model.type_hints` for available
+                options. Defaults to None.
+
+        Raises:
+            ValueError: Raised when the best-effort environment validation fails.
+
+        Returns:
+            The original Python object loaded from the model object.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        if not force:
+            with tempfile.TemporaryDirectory() as tmp_workspace_for_validation:
+                ws_path_for_validation = pathlib.Path(tmp_workspace_for_validation)
+                self._model_ops.download_files(
+                    database_name=None,
+                    schema_name=None,
+                    model_name=self._model_name,
+                    version_name=self._version_name,
+                    target_path=ws_path_for_validation,
+                    mode="minimal",
+                    statement_params=statement_params,
+                )
+                pk_for_validation = model_composer.ModelComposer.load(
+                    ws_path_for_validation, meta_only=True, options=options
+                )
+                assert pk_for_validation.meta, (
+                    "Unable to load model metadata for validation. "
+                    f"model_name={self._model_name}, version_name={self._version_name}"
+                )
+
+                validation_errors = pk_for_validation.meta.env.validate_with_local_env(
+                    check_snowpark_ml_version=(
+                        pk_for_validation.meta.model_type == snowmlmodel.SnowMLModelHandler.HANDLER_TYPE
+                    )
+                )
+                if validation_errors:
+                    raise ValueError(
+                        f"Unable to load this model due to following validation errors: {validation_errors}. "
+                        "Make sure your local environment is the same as that when you logged the model, "
+                        "or if you believe it should work, specify `force=True` to bypass this check."
+                    )
+
+        warnings.warn(
+            "Loading model requires to have the exact the same environment as the one when "
+            "logging the model, otherwise, the model might be not functional or "
+            "some other problems might occur.",
+            category=RuntimeWarning,
+            stacklevel=2,
+        )
+
+        # We need the folder to be existed.
+        workspace = pathlib.Path(tempfile.mkdtemp())
+        self._model_ops.download_files(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            target_path=workspace,
+            mode="model",
+            statement_params=statement_params,
+        )
+        pk = model_composer.ModelComposer.load(workspace, meta_only=False, options=options)
+        assert pk.model, (
+            "Unable to load model. "
+            f"model_name={self._model_name}, version_name={self._version_name}, metadata={pk.meta}"
+        )
+        return pk.model

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

@@ -61,12 +61,18 @@ class MetadataOperator:
     def _get_current_metadata_dict(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         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
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
         metadata_str = version_info_list[0][self._model_client.MODEL_VERSION_METADATA_COL_NAME]
         if not metadata_str:
@@ -79,12 +85,18 @@ class MetadataOperator:
     def load(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         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
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
         return MetadataOperator._parse(metadata_dict)
 
@@ -92,14 +104,25 @@ class MetadataOperator:
         self,
         metadata: ModelVersionMetadataSchema,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         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
+            database_name=database_name,
+            schema_name=schema_name,
+            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
+            metadata_dict,
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )