snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

snowflake/ml/fileset/stage_fs.py

@@ -8,7 +8,7 @@ import fsspec
 from fsspec.implementations import http as httpfs
 
 from snowflake import snowpark
-from snowflake.connector import connection
+from snowflake.connector import connection, errorcode
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -17,6 +17,7 @@ from snowflake.ml._internal.exceptions import (
     fileset_errors,
 )
 from snowflake.snowpark import exceptions as snowpark_exceptions
+from snowflake.snowpark._internal import utils as snowpark_utils
 
 # The default length of how long a presigned url stays active in seconds.
 # Presigned url here is used to fetch file objects from Snowflake when SFStageFileSystem.open() is called.
@@ -79,7 +80,9 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     #   None -> Try pre-signed URL access, fall back to file download
     #   True -> Use file download path without trying pre-signed URL access
     #   False -> Use pre-signed URL access, skip download fallback on failure
-    _USE_FALLBACK_FILE_ACCESS = None
+    _USE_FALLBACK_FILE_ACCESS = (
+        True if snowpark_utils.is_in_stored_procedure() else None  # type: ignore[no-untyped-call]
+    )
 
     def __init__(
         self,
@@ -141,7 +144,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         project=_PROJECT,
         func_params_to_log=["detail"],
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def ls(self, path: str, detail: bool = False) -> Union[List[str], List[Dict[str, Any]]]:
         """Override fsspec `ls` method. List single "directory" with or without details.
 
@@ -164,7 +166,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         """
         try:
             loc = self.stage_name
-            objects = self._session.sql(f"LIST {loc}/{path}").collect()
+            path = path.lstrip("/")
+            objects = self._session.sql(f"LIST '{loc}/{path}'").collect()
         except snowpark_exceptions.SnowparkClientException as e:
             if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST):
                 raise snowml_exceptions.SnowflakeMLException(
@@ -187,7 +190,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def optimize_read(self, files: Optional[List[str]] = None) -> None:
         """Prefetch and cache the presigned urls for all the given files to speed up the read performance.
 
@@ -214,7 +216,6 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
-    @snowpark._internal.utils.private_preview(version="0.2.0")
     def _open(self, path: str, mode: str = "rb", **kwargs: Any) -> fsspec.spec.AbstractBufferedFile:
         """Override fsspec `_open` method. Open a file for reading.
 
@@ -234,7 +235,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         """
         path = path.lstrip("/")
         if self._USE_FALLBACK_FILE_ACCESS:
-            return self._session.file.get_stream(f"{self.stage_name}/{path}")
+            return self._open_with_snowpark(path)
         cached_presigned_url = self._url_cache.get(path, None)
         if not cached_presigned_url:
             res = self._fetch_presigned_urls([path])
@@ -252,19 +253,42 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         except FileNotFoundError:
             # Enable fallback if _USE_FALLBACK_FILE_ACCESS is True or None; set to False to disable
             if self._USE_FALLBACK_FILE_ACCESS != False:  # noqa: E712
-                try:
-                    # Try falling back to Snowpark file read if presigned URL access failed
-                    # If fallback is successful, most likely in sproc without External Access Integration
-                    content = self._session.file.get_stream(f"{self.stage_name}/{path}")
-                    self._USE_FALLBACK_FILE_ACCESS = True
-                    return content
-                except Exception:
-                    pass
+                content = self._open_with_snowpark(path)
+                self._USE_FALLBACK_FILE_ACCESS = True
+                return content
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.SNOWML_NOT_FOUND,
                 original_exception=fileset_errors.StageFileNotFoundError(f"Stage file {path} doesn't exist."),
             )
 
+    def _open_with_snowpark(self, path: str, **kwargs: Dict[str, Any]) -> fsspec.spec.AbstractBufferedFile:
+        """Open the a file for reading using snowflake.snowpark.file_operation
+
+        Args:
+            path: Path of file in Snowflake stage.
+            **kwargs: Extra options to pass to snowflake.snowpark.file_operation.get_stream
+
+        Returns:
+            A fsspec file-like object.
+
+        Raises:
+            SnowflakeMLException: An error occurred when the given path points to a file that cannot be found.
+            SnowflakeMLException: An unknown Snowpark error occurred during file read.
+        """
+        try:
+            return self._session.file.get_stream(f"{self.stage_name}/{path}", **kwargs)
+        except snowpark_exceptions.SnowparkSQLException as e:
+            if _match_error_code(e, errorcode.ER_FILE_NOT_EXISTS):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.SNOWML_NOT_FOUND,
+                    original_exception=fileset_errors.StageFileNotFoundError(f"Stage file {path} doesn't exist."),
+                )
+            else:
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
+                    original_exception=e,
+                )
+
     def _parse_list_result(
         self, list_result: List[Tuple[str, int, str, str]], search_path: str
     ) -> List[Dict[str, Any]]:
@@ -352,7 +376,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         file_df = self._session.create_dataframe(files).to_df("name")
         try:
             presigned_urls: List[Tuple[str, str]] = file_df.select_expr(
-                f"name, get_presigned_url({self.stage_name}, name, {url_lifetime}) as url"
+                f"name, get_presigned_url('{self.stage_name}', name, {url_lifetime}) as url"
             ).collect(
                 statement_params=telemetry.get_function_usage_statement_params(
                     project=_PROJECT,
@@ -363,7 +387,9 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                 ),
             )
         except snowpark_exceptions.SnowparkClientException as e:
-            if e.message.startswith(fileset_errors.ERRNO_STAGE_NOT_EXIST):
+            if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST) or e.message.startswith(
+                fileset_errors.ERRNO_STAGE_NOT_EXIST
+            ):
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_NOT_FOUND,
                     original_exception=fileset_errors.StageNotFoundError(
@@ -376,3 +402,9 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     original_exception=fileset_errors.FileSetError(str(e)),
                 )
         return presigned_urls
+
+
+def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code: int) -> bool:
+    # Snowpark writes error code to message instead of populating e.error_code
+    error_code_str = str(error_code)
+    return ex.error_code == error_code_str or error_code_str in ex.message

snowflake/ml/model/__init__.py

@@ -1,6 +1,6 @@
 from snowflake.ml.model._client.model.model_impl import Model
-from snowflake.ml.model._client.model.model_version_impl import ModelVersion
+from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
 from snowflake.ml.model.models.llm import LLM, LLMOptions
 
-__all__ = ["Model", "ModelVersion", "HuggingFacePipelineModel", "LLM", "LLMOptions"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "LLM", "LLMOptions"]

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

@@ -350,3 +350,30 @@ class Model:
             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,
+        )
+        db, schema, model, _ = identifier.parse_schema_level_object_identifier(model_name)
+        new_model_db = sql_identifier.SqlIdentifier(db) if db else None
+        new_model_schema = sql_identifier.SqlIdentifier(schema) if schema else None
+        new_model_id = sql_identifier.SqlIdentifier(model)
+        self._model_ops.rename(
+            model_name=self._model_name,
+            new_model_db=new_model_db,
+            new_model_schema=new_model_schema,
+            new_model_name=new_model_id,
+            statement_params=statement_params,
+        )
+        self._model_name = new_model_id

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

@@ -1,20 +1,29 @@
-import re
+import enum
+import pathlib
+import tempfile
+import warnings
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import pandas as pd
 
-from snowflake import connector
 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 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."""
 
@@ -207,61 +216,16 @@ class ModelVersion:
             statement_params=statement_params,
         )
 
-    # Only used when the model does not contains user_data with client SDK information.
-    def _legacy_show_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
+    def _get_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        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(
+        return self._model_ops.get_functions(
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
         )
-        return_functions_info: List[model_manifest_schema.ModelFunctionInfo] = []
-        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]
-            fi = model_manifest_schema.ModelFunctionInfo(
-                name=method_name,
-                target_method=target_method,
-                signature=model_signature.ModelSignature.from_dict(signature_dict),
-            )
-            return_functions_info.append(fi)
-        return return_functions_info
-
-    def _get_functions(self) -> List[model_manifest_schema.ModelFunctionInfo]:
-        statement_params = telemetry.get_statement_params(
-            project=_TELEMETRY_PROJECT,
-            subproject=_TELEMETRY_SUBPROJECT,
-        )
-        try:
-            client_data = self._model_ops.get_client_data_in_user_data(
-                model_name=self._model_name,
-                version_name=self._version_name,
-                statement_params=statement_params,
-            )
-            return [
-                model_manifest_schema.ModelFunctionInfo(
-                    name=fi["name"],
-                    target_method=fi["target_method"],
-                    signature=model_signature.ModelSignature.from_dict(fi["signature"]),
-                )
-                for fi in client_data["functions"]
-            ]
-        except (NotImplementedError, ValueError, connector.DataError):
-            return self._legacy_show_functions()
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
@@ -288,6 +252,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.
@@ -296,12 +261,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.
@@ -311,6 +278,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()
@@ -335,10 +306,126 @@ 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,
             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(
+            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(
+                    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(
+            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