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

snowflake/ml/registry/_manager/model_manager.py

@@ -0,0 +1,163 @@
+from types import ModuleType
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from absl.logging import logging
+
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model import model_signature, type_hints as model_types
+from snowflake.ml.model._client.model import model_impl, model_version_impl
+from snowflake.ml.model._client.ops import metadata_ops, model_ops
+from snowflake.ml.model._model_composer import model_composer
+from snowflake.snowpark import session
+
+logger = logging.getLogger(__name__)
+
+
+class ModelManager:
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._database_name = database_name
+        self._schema_name = schema_name
+        self._model_ops = model_ops.ModelOperator(
+            session, database_name=self._database_name, schema_name=self._schema_name
+        )
+
+    def log_model(
+        self,
+        model: model_types.SupportedModelType,
+        *,
+        model_name: str,
+        version_name: str,
+        comment: Optional[str] = None,
+        metrics: Optional[Dict[str, Any]] = None,
+        conda_dependencies: Optional[List[str]] = None,
+        pip_requirements: Optional[List[str]] = None,
+        python_version: Optional[str] = None,
+        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        code_paths: Optional[List[str]] = None,
+        ext_modules: Optional[List[ModuleType]] = None,
+        options: Optional[model_types.ModelSaveOption] = None,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> model_version_impl.ModelVersion:
+        model_name_id = sql_identifier.SqlIdentifier(model_name)
+
+        version_name_id = sql_identifier.SqlIdentifier(version_name)
+
+        if self._model_ops.validate_existence(
+            model_name=model_name_id, statement_params=statement_params
+        ) and self._model_ops.validate_existence(
+            model_name=model_name_id, version_name=version_name_id, statement_params=statement_params
+        ):
+            raise ValueError(f"Model {model_name} version {version_name} already existed.")
+
+        stage_path = self._model_ops.prepare_model_stage_path(
+            statement_params=statement_params,
+        )
+
+        logger.info("Start packaging and uploading your model. It might take some time based on the size of the model.")
+
+        mc = model_composer.ModelComposer(self._model_ops._session, stage_path=stage_path)
+        mc.save(
+            name=model_name_id.resolved(),
+            model=model,
+            signatures=signatures,
+            sample_input=sample_input_data,
+            conda_dependencies=conda_dependencies,
+            pip_requirements=pip_requirements,
+            python_version=python_version,
+            code_paths=code_paths,
+            ext_modules=ext_modules,
+            options=options,
+        )
+
+        logger.info("Start creating MODEL object for you in the Snowflake.")
+
+        self._model_ops.create_from_stage(
+            composed_model=mc,
+            model_name=model_name_id,
+            version_name=version_name_id,
+            statement_params=statement_params,
+        )
+
+        mv = model_version_impl.ModelVersion._ref(
+            self._model_ops,
+            model_name=model_name_id,
+            version_name=version_name_id,
+        )
+
+        if comment:
+            mv.comment = comment
+
+        if metrics:
+            self._model_ops._metadata_ops.save(
+                metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+                model_name=model_name_id,
+                version_name=version_name_id,
+                statement_params=statement_params,
+            )
+
+        return mv
+
+    def get_model(
+        self,
+        model_name: str,
+        *,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> model_impl.Model:
+        model_name_id = sql_identifier.SqlIdentifier(model_name)
+        if self._model_ops.validate_existence(
+            model_name=model_name_id,
+            statement_params=statement_params,
+        ):
+            return model_impl.Model._ref(
+                self._model_ops,
+                model_name=model_name_id,
+            )
+        else:
+            raise ValueError(f"Unable to find model {model_name}")
+
+    def models(
+        self,
+        *,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[model_impl.Model]:
+        model_names = self._model_ops.list_models_or_versions(
+            statement_params=statement_params,
+        )
+        return [
+            model_impl.Model._ref(
+                self._model_ops,
+                model_name=model_name,
+            )
+            for model_name in model_names
+        ]
+
+    def show_models(
+        self,
+        *,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> pd.DataFrame:
+        rows = self._model_ops.show_models_or_versions(
+            statement_params=statement_params,
+        )
+        return pd.DataFrame([row.as_dict() for row in rows])
+
+    def delete_model(
+        self,
+        model_name: str,
+        *,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        model_name_id = sql_identifier.SqlIdentifier(model_name)
+
+        self._model_ops.delete_model_or_version(
+            model_name=model_name_id,
+            statement_params=statement_params,
+        )

snowflake/ml/registry/model_registry.py

@@ -3,6 +3,7 @@ import json
 import sys
 import textwrap
 import types
+import warnings
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -305,6 +306,17 @@ class ModelRegistry:
             schema_name: Desired name of the schema used by this model registry inside the database.
             create_if_not_exists: create model registry if it's not exists already.
         """
+
+        warnings.warn(
+            """
+The `snowflake.ml.registry.model_registry.ModelRegistry` has been deprecated starting from version 1.2.0.
+It will stay in the Private Preview phase. For future implementations, kindly utilize `snowflake.ml.registry.Registry`,
+except when specifically required. The old model registry will be removed once all its primary functionalities are
+fully integrated into the new registry.
+        """,
+            DeprecationWarning,
+            stacklevel=2,
+        )
         if create_if_not_exists:
             create_model_registry(session=session, database_name=database_name, schema_name=schema_name)
 

snowflake/ml/registry/registry.py

@@ -1,12 +1,17 @@
 from types import ModuleType
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
+
+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, type_hints as model_types
-from snowflake.ml.model._client.model import model_impl, model_version_impl
-from snowflake.ml.model._client.ops import model_ops
-from snowflake.ml.model._model_composer import model_composer
+from snowflake.ml.model import (
+    Model,
+    ModelVersion,
+    model_signature,
+    type_hints as model_types,
+)
+from snowflake.ml.registry._manager import model_manager
 from snowflake.snowpark import session
 
 _TELEMETRY_PROJECT = "MLOps"
@@ -21,6 +26,18 @@ class Registry:
         database_name: Optional[str] = None,
         schema_name: Optional[str] = None,
     ) -> None:
+        """Opens a registry within a pre-created Snowflake schema.
+
+        Args:
+            session: The Snowpark Session to connect with Snowflake.
+            database_name: The name of the database. If None, the current database of the session
+                will be used. Defaults to None.
+            schema_name: The name of the schema. If None, the current schema of the session
+                will be used. If there is no active schema, the PUBLIC schema will be used. Defaults to None.
+
+        Raises:
+            ValueError: When there is no specified or active database in the session.
+        """
         if database_name:
             self._database_name = sql_identifier.SqlIdentifier(database_name)
         else:
@@ -42,12 +59,13 @@ class Registry:
                 else sql_identifier.SqlIdentifier("PUBLIC")
             )
 
-        self._model_ops = model_ops.ModelOperator(
+        self._model_manager = model_manager.ModelManager(
             session, database_name=self._database_name, schema_name=self._schema_name
         )
 
     @property
     def location(self) -> str:
+        """Get the location (database.schema) of the registry."""
         return ".".join([self._database_name.identifier(), self._schema_name.identifier()])
 
     @telemetry.send_api_usage_telemetry(
@@ -60,6 +78,8 @@ class Registry:
         *,
         model_name: str,
         version_name: str,
+        comment: Optional[str] = None,
+        metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
         python_version: Optional[str] = None,
@@ -68,148 +88,138 @@ class Registry:
         code_paths: Optional[List[str]] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         options: Optional[model_types.ModelSaveOption] = None,
-    ) -> model_version_impl.ModelVersion:
-        """Log a model.
+    ) -> ModelVersion:
+        """
+        Log a model with various parameters and metadata.
 
         Args:
-            model: Model Python object
-            model_name: A string as name.
-            version_name: A string as version. model_name and version_name combination must be unique.
-            signatures: Model data signatures for inputs and output for every target methods. If it is None,
+            model: Model object of supported types such as Scikit-learn, XGBoost, Snowpark ML,
+                PyTorch, TorchScript, Tensorflow, Tensorflow Keras, MLFlow, HuggingFace Pipeline,
+                Peft-finetuned LLM, or Custom Model.
+            model_name: Name to identify the model.
+            version_name: Version identifier for the model. Combination of model_name and version_name must be unique.
+            comment: Comment associated with the model version. Defaults to None.
+            metrics: A JSON serializable dictionary containing metrics linked to the model version. Defaults to None.
+            signatures: Model data signatures for inputs and outputs for various target methods. If it is None,
                 sample_input_data would be used to infer the signatures for those models that cannot automatically
-                infer the signature. If not None, sample_input should not be specified. Defaults to None.
-            sample_input_data: Sample input data to infer the model signatures from. If it is None, signatures must be
-                specified if the model cannot automatically infer the signature. If not None, signatures should not be
-                specified. Defaults to None.
-            conda_dependencies: List of Conda package specs. Use "[channel::]package [operator version]" syntax to
-                specify a dependency. It is a recommended way to specify your dependencies using conda. When channel is
-                not specified, Snowflake Anaconda Channel will be used.
-            pip_requirements: List of Pip package specs.
-            python_version: A string of python version where model is run. Used for user override. If specified as None,
-                current version would be captured. Defaults to None.
-            code_paths: Directory of code to import.
-            ext_modules: External modules that user might want to get pickled with model object. Defaults to None.
-            options: Model specific kwargs.
+                infer the signature. If not None, sample_input_data should not be specified. Defaults to None.
+            sample_input_data: Sample input data to infer model signatures from. Defaults to None.
+            conda_dependencies: List of Conda package specifications. Use "[channel::]package [operator version]" syntax
+                to specify a dependency. It is a recommended way to specify your dependencies using conda. When channel
+                is not specified, Snowflake Anaconda Channel will be used. Defaults to None.
+            pip_requirements: List of Pip package specifications. Defaults to None.
+            python_version: Python version in which the model is run. Defaults to None.
+            code_paths: List of directories containing code to import. Defaults to None.
+            ext_modules: List of external modules to pickle with the model object.
+                Only supported when logging the following types of model:
+                Scikit-learn, Snowpark ML, PyTorch, TorchScript and Custom Model. Defaults to None.
+            options (Dict[str, Any], optional): Additional model saving options.
+
+                Model Saving Options include:
+
+                - embed_local_ml_library: Embed local Snowpark ML into the code directory or folder.
+                    Override to True if the local Snowpark ML version is not available in the Snowflake Anaconda
+                    Channel. Otherwise, defaults to False
+                - relax_version: Whether or not relax the version constraints of the dependencies.
+                    It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
+                - method_options: Per-method saving options including:
+                    - case_sensitive: Indicates whether the method and its signature should be case sensitive.
+                        This means when you refer the method in the SQL, you need to double quote it.
+                        This will be helpful if you need case to tell apart your methods or features, or you have
+                        non-alphabetic characters in your method or feature name. Defaults to False.
+                    - max_batch_size: Maximum batch size that the method could accept in the Snowflake Warehouse.
+                        Defaults to None, determined automatically by Snowflake.
 
         Returns:
-            A ModelVersion object corresponding to the model just get logged.
+            ModelVersion: ModelVersion object corresponding to the model just logged.
         """
 
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
-        model_name_id = sql_identifier.SqlIdentifier(model_name)
-
-        version_name_id = sql_identifier.SqlIdentifier(version_name)
-
-        stage_path = self._model_ops.prepare_model_stage_path(
-            statement_params=statement_params,
-        )
-
-        mc = model_composer.ModelComposer(self._model_ops._session, stage_path=stage_path)
-        mc.save(
-            name=model_name_id.resolved(),
+        return self._model_manager.log_model(
             model=model,
-            signatures=signatures,
-            sample_input=sample_input_data,
+            model_name=model_name,
+            version_name=version_name,
+            comment=comment,
+            metrics=metrics,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             python_version=python_version,
+            signatures=signatures,
+            sample_input_data=sample_input_data,
             code_paths=code_paths,
             ext_modules=ext_modules,
             options=options,
-        )
-        self._model_ops.create_from_stage(
-            composed_model=mc,
-            model_name=model_name_id,
-            version_name=version_name_id,
             statement_params=statement_params,
         )
 
-        return model_version_impl.ModelVersion._ref(
-            self._model_ops,
-            model_name=model_name_id,
-            version_name=version_name_id,
-        )
-
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_MODEL_TELEMETRY_SUBPROJECT,
     )
-    def get_model(self, model_name: str) -> model_impl.Model:
-        """Get the model object.
+    def get_model(self, model_name: str) -> Model:
+        """Get the model object by its name.
 
         Args:
-            model_name: The model name.
-
-        Raises:
-            ValueError: Raised when the model requested does not exist.
+            model_name: The name of the model.
 
         Returns:
-            The model object.
+            The corresponding model object.
         """
-        model_name_id = sql_identifier.SqlIdentifier(model_name)
-
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
-        if self._model_ops.validate_existence(
-            model_name=model_name_id,
-            statement_params=statement_params,
-        ):
-            return model_impl.Model._ref(
-                self._model_ops,
-                model_name=model_name_id,
-            )
-        else:
-            raise ValueError(f"Unable to find model {model_name}")
+        return self._model_manager.get_model(model_name=model_name, statement_params=statement_params)
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_MODEL_TELEMETRY_SUBPROJECT,
     )
-    def list_models(self) -> List[model_impl.Model]:
-        """List all models in the schema where the registry is opened.
+    def models(self) -> List[Model]:
+        """Get all models in the schema where the registry is opened.
 
         Returns:
-            A List of Model= object representing all models in the schema where the registry is opened.
+            A list of Model objects representing all models in the opened registry.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
-        model_names = self._model_ops.list_models_or_versions(
-            statement_params=statement_params,
+        return self._model_manager.models(statement_params=statement_params)
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_MODEL_TELEMETRY_SUBPROJECT,
+    )
+    def show_models(self) -> pd.DataFrame:
+        """Show information of all models in the schema where the registry is opened.
+
+        Returns:
+            A Pandas DataFrame containing information of all models in the schema.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
-        return [
-            model_impl.Model._ref(
-                self._model_ops,
-                model_name=model_name,
-            )
-            for model_name in model_names
-        ]
+        return self._model_manager.show_models(statement_params=statement_params)
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_MODEL_TELEMETRY_SUBPROJECT,
     )
     def delete_model(self, model_name: str) -> None:
-        """Delete the model.
+        """
+        Delete the model by its name.
 
         Args:
-            model_name: The model name, can be fully qualified one.
-                If not, use database name and schema name of the registry.
+            model_name: The name of the model to be deleted.
         """
-        model_name_id = sql_identifier.SqlIdentifier(model_name)
-
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
 
-        self._model_ops.delete_model_or_version(
-            model_name=model_name_id,
-            statement_params=statement_params,
-        )
+        self._model_manager.delete_model(model_name=model_name, statement_params=statement_params)

snowflake/ml/version.py

@@ -1 +1 @@
-VERSION="1.1.2"
+VERSION="1.2.1"