snowflake-ml-python 1.7.5__py3-none-any.whl → 1.8.1__py3-none-any.whl

snowflake/ml/model/model_signature.py

@@ -32,6 +32,7 @@ from snowflake.ml.model._signatures import (
     base_handler,
     builtins_handler as builtins_handler,
     core,
+    dmatrix_handler,
     numpy_handler,
     pandas_handler,
     pytorch_handler,
@@ -52,8 +53,11 @@ _LOCAL_DATA_HANDLERS: List[Type[base_handler.BaseDataHandler[Any]]] = [
     numpy_handler.NumpyArrayHandler,
     builtins_handler.ListOfBuiltinHandler,
     numpy_handler.SeqOfNumpyArrayHandler,
+    pytorch_handler.PyTorchTensorHandler,
     pytorch_handler.SeqOfPyTorchTensorHandler,
+    tensorflow_handler.TensorflowTensorHandler,
     tensorflow_handler.SeqOfTensorflowTensorHandler,
+    dmatrix_handler.XGBoostDMatrixHandler,
 ]
 _ALL_DATA_HANDLERS = _LOCAL_DATA_HANDLERS + [snowpark_handler.SnowparkDataFrameHandler]
 
@@ -218,7 +222,6 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         strict: Enable strict validation, this includes value range based validation
 
     Raises:
-        SnowflakeMLException: NotImplementedError: FeatureGroupSpec is not supported.
         SnowflakeMLException: ValueError: Raised when a feature cannot be found.
         SnowflakeMLException: ValueError: Raised when feature is scalar but confront list element.
         SnowflakeMLException: ValueError: Raised when feature type is not aligned in list element.
@@ -236,7 +239,10 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         except KeyError:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_DATA,
-                original_exception=ValueError(f"Data Validation Error: feature {ft_name} does not exist in data."),
+                original_exception=ValueError(
+                    f"Data Validation Error: feature {ft_name} does not exist in data. "
+                    f"Available columns are {data.columns}."
+                ),
             )
 
         if data_col.isnull().any():
@@ -244,10 +250,15 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         df_col_dtype = data_col.dtype
 
         if isinstance(feature, core.FeatureGroupSpec):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.NOT_IMPLEMENTED,
-                original_exception=NotImplementedError("FeatureGroupSpec is not supported."),
-            )
+            if df_col_dtype != np.dtype("O"):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INVALID_DATA,
+                    original_exception=ValueError(
+                        f"Data Validation Error in feature group {ft_name}: "
+                        + f"It needs to be a dictionary or list of dictionary, but get {df_col_dtype}."
+                    ),
+                )
+            continue
 
         assert isinstance(feature, core.FeatureSpec)  # assert for mypy.
         ft_type = feature._dtype
@@ -437,7 +448,6 @@ def _validate_snowpark_data(
         strict: Enable strict validation, this includes value range based validation.
 
     Raises:
-        SnowflakeMLException: NotImplementedError: FeatureGroupSpec is not supported.
         SnowflakeMLException: ValueError: Raised when confronting invalid feature.
         SnowflakeMLException: ValueError: Raised when a feature cannot be found.
 
@@ -467,10 +477,15 @@ def _validate_snowpark_data(
                 if field.name == ft_name:
                     found = True
                     if isinstance(feature, core.FeatureGroupSpec):
-                        raise snowml_exceptions.SnowflakeMLException(
-                            error_code=error_codes.NOT_IMPLEMENTED,
-                            original_exception=NotImplementedError("FeatureGroupSpec is not supported."),
-                        )
+                        if not isinstance(field.datatype, (spt.ArrayType, spt.StructType, spt.VariantType)):
+                            errors[identifier_rule].append(
+                                ValueError(
+                                    f"Data Validation Error in feature group {feature.name}: "
+                                    + f"Feature expects {feature.as_snowpark_type()},"
+                                    + f" while {field.name} has type {field.datatype}."
+                                ),
+                            )
+                        continue
                     assert isinstance(feature, core.FeatureSpec)  # mypy
                     ft_type = feature._dtype
                     field_data_type = field.datatype
@@ -644,11 +659,14 @@ def _validate_snowpark_type_feature(
             )
 
 
-def _convert_local_data_to_df(data: model_types.SupportedLocalDataType) -> pd.DataFrame:
+def _convert_local_data_to_df(
+    data: model_types.SupportedLocalDataType, ensure_serializable: bool = False
+) -> pd.DataFrame:
     """Convert local data to pandas DataFrame or Snowpark DataFrame
 
     Args:
         data: The provided data.
+        ensure_serializable: Ensure the data is serializable. Defaults to False.
 
     Raises:
         SnowflakeMLException: NotImplementedError: Raised when data cannot be handled by any data handler.
@@ -660,7 +678,7 @@ def _convert_local_data_to_df(data: model_types.SupportedLocalDataType) -> pd.Da
     for handler in _LOCAL_DATA_HANDLERS:
         if handler.can_handle(data):
             handler.validate(data)
-            df = handler.convert_to_df(data, ensure_serializable=False)
+            df = handler.convert_to_df(data, ensure_serializable=ensure_serializable)
             break
     if df is None:
         raise snowml_exceptions.SnowflakeMLException(

snowflake/ml/model/type_hints.py

@@ -26,7 +26,15 @@ if TYPE_CHECKING:
     from snowflake.ml.modeling.framework import base  # noqa: F401
 
 
-_SupportedBuiltins = Union[int, float, bool, str, bytes, "_SupportedBuiltinsList"]
+_SupportedBuiltins = Union[
+    int,
+    float,
+    bool,
+    str,
+    bytes,
+    Dict[str, Union["_SupportedBuiltins", "_SupportedBuiltinsList"]],
+    "_SupportedBuiltinsList",
+]
 _SupportedNumpyDtype = Union[
     "np.int8",
     "np.int16",
@@ -48,7 +56,7 @@ _SupportedBuiltinsList = Sequence[_SupportedBuiltins]
 _SupportedArrayLike = Union[_SupportedNumpyArray, "torch.Tensor", "tensorflow.Tensor", "tensorflow.Variable"]
 
 SupportedLocalDataType = Union[
-    "pd.DataFrame", _SupportedNumpyArray, Sequence[_SupportedArrayLike], _SupportedBuiltinsList
+    "pd.DataFrame", _SupportedArrayLike, Sequence[_SupportedArrayLike], _SupportedBuiltinsList
 ]
 
 SupportedDataType = Union[SupportedLocalDataType, "snowflake.snowpark.DataFrame"]
@@ -177,16 +185,19 @@ class SNOWModelSaveOptions(BaseModelSaveOption):
 class PyTorchSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class TorchScriptSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class TensorflowSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class MLFlowSaveOptions(BaseModelSaveOption):

snowflake/ml/modeling/_internal/estimator_utils.py

@@ -130,7 +130,11 @@ def is_single_node(session: Session) -> bool:
     warehouse_name = session.get_current_warehouse()
     if warehouse_name:
         warehouse_name = warehouse_name.replace('"', "")
-        df = session.sql(f"SHOW WAREHOUSES like '{warehouse_name}';")['"type"', '"size"'].collect()[0]
+        df_list = session.sql(f"SHOW WAREHOUSES like '{warehouse_name}';")['"type"', '"size"'].collect()
+        # If no warehouse data is found, default to True (single node)
+        if not df_list:
+            return True
+        df = df_list[0]
         # filter out the conditions when it is single node
         single_node: bool = (df[0] == "SNOWPARK-OPTIMIZED" and df[1] == "Medium") or (
             df[0] == "STANDARD" and df[1] == "X-Small"

snowflake/ml/modeling/metrics/ranking.py

@@ -98,6 +98,7 @@ def precision_recall_curve(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -245,6 +246,7 @@ def roc_auc_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -348,6 +350,7 @@ def roc_curve(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,

snowflake/ml/modeling/metrics/regression.py

@@ -83,6 +83,7 @@ def d2_absolute_error_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -180,6 +181,7 @@ def d2_pinball_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -295,6 +297,7 @@ def explained_variance_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,

snowflake/ml/modeling/preprocessing/k_bins_discretizer.py

@@ -341,7 +341,7 @@ class KBinsDiscretizer(base.BaseTransformer):
             is_permanent=False,
             name=udf_name,
             replace=True,
-            packages=["numpy"],
+            packages=[f"numpy=={np.__version__}"],
             session=dataset._session,
             statement_params=telemetry.get_statement_params(base.PROJECT, base.SUBPROJECT, self.__class__.__name__),
         )

snowflake/ml/registry/_manager/model_manager.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import pandas as pd
 from absl.logging import logging
 
-from snowflake.ml._internal import telemetry
+from snowflake.ml._internal import platform_capabilities, telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import sql_identifier
@@ -13,7 +13,7 @@ from snowflake.ml.model._client.model import model_impl, model_version_impl
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._packager.model_meta import model_meta
-from snowflake.snowpark import session
+from snowflake.snowpark import exceptions as snowpark_exceptions, session
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +46,7 @@ class ModelManager:
         metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         target_platforms: Optional[List[model_types.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
@@ -127,6 +128,7 @@ class ModelManager:
             metrics=metrics,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             target_platforms=target_platforms,
             python_version=python_version,
             signatures=signatures,
@@ -149,6 +151,7 @@ class ModelManager:
         metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         target_platforms: Optional[List[model_types.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
@@ -163,11 +166,42 @@ class ModelManager:
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
         version_name_id = sql_identifier.SqlIdentifier(version_name)
 
-        stage_path = self._model_ops.prepare_model_stage_path(
-            database_name=database_name_id,
-            schema_name=schema_name_id,
-            statement_params=statement_params,
-        )
+        use_live_commit = platform_capabilities.PlatformCapabilities.get_instance().is_live_commit_enabled()
+        if use_live_commit:
+            logger.info("Using live commit model version")
+        else:
+            logger.info("Using non-live commit model version")
+
+        if use_live_commit:
+            # This step creates the live model version, and the files can be written directly to the stage
+            # after this.
+            try:
+                self._model_ops.add_or_create_live_version(
+                    database_name=database_name_id,
+                    schema_name=schema_name_id,
+                    model_name=model_name_id,
+                    version_name=version_name_id,
+                    statement_params=statement_params,
+                )
+            except (AssertionError, snowpark_exceptions.SnowparkSQLException) as e:
+                logger.info(f"Failed to create live model version: {e}, falling back to regular model version creation")
+                use_live_commit = False
+
+        if use_live_commit:
+            # using model version's stage path to write files directly to the stage
+            stage_path = self._model_ops.get_model_version_stage_path(
+                database_name=database_name_id,
+                schema_name=schema_name_id,
+                model_name=model_name_id,
+                version_name=version_name_id,
+            )
+        else:
+            # using a temp path to write files and then upload to the model version's stage
+            stage_path = self._model_ops.prepare_model_temp_stage_path(
+                database_name=database_name_id,
+                schema_name=schema_name_id,
+                statement_params=statement_params,
+            )
 
         platforms = None
         # User specified target platforms are defaulted to None and will not show up in the generated manifest.
@@ -175,6 +209,18 @@ class ModelManager:
             # Convert any string target platforms to TargetPlatform objects
             platforms = [model_types.TargetPlatform(platform) for platform in target_platforms]
 
+        if artifact_repository_map:
+            for channel, artifact_repository_name in artifact_repository_map.items():
+                db_id, schema_id, repo_id = sql_identifier.parse_fully_qualified_name(artifact_repository_name)
+
+                artifact_repository_map[channel] = sql_identifier.get_fully_qualified_name(
+                    db_id,
+                    schema_id,
+                    repo_id,
+                    self._database_name,
+                    self._schema_name,
+                )
+
         logger.info("Start packaging and uploading your model. It might take some time based on the size of the model.")
 
         mc = model_composer.ModelComposer(
@@ -187,6 +233,7 @@ class ModelManager:
             sample_input_data=sample_input_data,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             target_platforms=platforms,
             python_version=python_version,
             user_files=user_files,
@@ -211,6 +258,7 @@ class ModelManager:
             model_name=model_name_id,
             version_name=version_name_id,
             statement_params=statement_params,
+            use_live_commit=use_live_commit,
         )
 
         mv = model_version_impl.ModelVersion._ref(

snowflake/ml/registry/registry.py

@@ -108,12 +108,15 @@ class Registry:
         metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         target_platforms: Optional[List[model_types.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
+        user_files: Optional[Dict[str, List[str]]] = None,
         code_paths: Optional[List[str]] = None,
         ext_modules: Optional[List[ModuleType]] = None,
+        task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> ModelVersion:
         """
@@ -140,6 +143,13 @@ class Registry:
                 See https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/container for more.
                 Models with pip requirements specified will not be executable in Snowflake Warehouse where all
                 dependencies must be retrieved from Snowflake Anaconda Channel.
+            artifact_repository_map: Specifies a mapping of package channels or platforms to custom artifact
+                repositories. Defaults to None. Currently, the mapping applies only to warehouse execution.
+                Note : This feature is currently in Private Preview; please contact your Snowflake account team
+                to enable it.
+                Format: {channel_name: artifact_repository_name}, where:
+                   - channel_name: The name of the Conda package channel (e.g., 'condaforge') or 'pip' for pip packages.
+                   - artifact_repository_name: The name or URL of the repository to fetch packages from.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
@@ -148,10 +158,15 @@ class Registry:
                 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.
                 It would also be used as background data in explanation and to capture data lineage. Defaults to None.
+            user_files: Dictionary where the keys are subdirectories, and values are lists of local file name
+                strings. The local file name strings can include wildcards (? or *) for matching multiple files.
             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.
+            task: The task of the Model Version. It is an enum class Task with values TABULAR_REGRESSION,
+                TABULAR_BINARY_CLASSIFICATION, TABULAR_MULTI_CLASSIFICATION, TABULAR_RANKING, or UNKNOWN. By default,
+                it is set to Task.UNKNOWN and may be overridden by inferring from the Model Object.
             options (Dict[str, Any], optional): Additional model saving options.
 
                 Model Saving Options include:
@@ -163,6 +178,9 @@ class Registry:
                     Warehouse. It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
                 - function_type: Set the method function type globally. To set method function types individually see
                     function_type in model_options.
+                - target_methods: List of target methods to register when logging the model.
+                  This option is not used in MLFlow models. Defaults to None, in which case the model handler's
+                  default target methods will be used.
                 - method_options: Per-method saving options. This dictionary has method names as keys and dictionary
                     values with the desired options.
 
@@ -210,6 +228,7 @@ class Registry:
             "metrics",
             "conda_dependencies",
             "pip_requirements",
+            "artifact_repository_map",
             "target_platforms",
             "python_version",
             "signatures",
@@ -225,6 +244,7 @@ class Registry:
         metrics: Optional[Dict[str, Any]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         target_platforms: Optional[List[model_types.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
@@ -259,6 +279,13 @@ class Registry:
                 See https://docs.snowflake.com/en/developer-guide/snowflake-ml/model-registry/container for more.
                 Models with pip requirements specified will not be executable in Snowflake Warehouse where all
                 dependencies must be retrieved from Snowflake Anaconda Channel.
+            artifact_repository_map: Specifies a mapping of package channels or platforms to custom artifact
+                repositories. Defaults to None. Currently, the mapping applies only to warehouse execution.
+                Note : This feature is currently in Private Preview; please contact your Snowflake account team to
+                enable it.
+                Format: {channel_name: artifact_repository_name}, where:
+                   - channel_name: The name of the Conda package channel (e.g., 'condaforge') or 'pip' for pip packages.
+                   - artifact_repository_name: The name or URL of the repository to fetch packages from.
             target_platforms: List of target platforms to run the model. The only acceptable inputs are a combination of
                 {"WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"}. Defaults to None.
             python_version: Python version in which the model is run. Defaults to None.
@@ -287,6 +314,9 @@ class Registry:
                     Warehouse. It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
                 - function_type: Set the method function type globally. To set method function types individually see
                   function_type in model_options.
+                - target_methods: List of target methods to register when logging the model.
+                  This option is not used in MLFlow models. Defaults to None, in which case the model handler's
+                  default target methods will be used.
                 - method_options: Per-method saving options. This dictionary has method names as keys and dictionary
                     values with the desired options. See the example below.
 
@@ -300,6 +330,9 @@ class Registry:
                         Defaults to None, determined automatically by Snowflake.
                     - function_type: One of supported model method function types (FUNCTION or TABLE_FUNCTION).
 
+        Raises:
+            ValueError: If extra arguments are specified ModelVersion is provided.
+
         Returns:
             ModelVersion: ModelVersion object corresponding to the model just logged.
 
@@ -322,13 +355,37 @@ class Registry:
             registry.log_model(
               model=model,
               model_name="my_model",
-              method_options=method_options,
+              options={"method_options": method_options},
             )
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_MODEL_TELEMETRY_SUBPROJECT,
         )
+        if isinstance(model, ModelVersion):
+            # check that no arguments are provided other than the ones for copy model.
+            invalid_args = [
+                comment,
+                conda_dependencies,
+                pip_requirements,
+                artifact_repository_map,
+                target_platforms,
+                python_version,
+                signatures,
+                sample_input_data,
+                user_files,
+                code_paths,
+                ext_modules,
+                options,
+            ]
+            for arg in invalid_args:
+                if arg is not None:
+                    raise ValueError(
+                        "When calling log_model with a ModelVersion, only model_name and version_name may be specified."
+                    )
+            if task is not model_types.Task.UNKNOWN:
+                raise ValueError("`task` cannot be specified when calling log_model with a ModelVersion.")
+
         if pip_requirements:
             warnings.warn(
                 "Models logged specifying `pip_requirements` can not be executed "
@@ -345,6 +402,7 @@ class Registry:
             metrics=metrics,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             target_platforms=target_platforms,
             python_version=python_version,
             signatures=signatures,

snowflake/ml/version.py

@@ -1 +1 @@
-VERSION="1.7.5"
+VERSION="1.8.1"