snowflake-ml-python 1.5.3__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -3,6 +3,7 @@ import os
 import warnings
 from typing import (
     TYPE_CHECKING,
+    Any,
     Callable,
     Dict,
     List,
@@ -250,9 +251,18 @@ class HuggingFacePipelineHandler(
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
     @staticmethod
-    def _get_device_config() -> Dict[str, str]:
-        device_config = {}
-        device_config["device_map"] = "auto"
+    def _get_device_config(**kwargs: Unpack[model_types.HuggingFaceLoadOptions]) -> Dict[str, str]:
+        device_config: Dict[str, Any] = {}
+        if (
+            kwargs.get("use_gpu", False)
+            and kwargs.get("device_map", None) is None
+            and kwargs.get("device", None) is None
+        ):
+            device_config["device_map"] = "auto"
+        elif kwargs.get("device_map", None) is not None:
+            device_config["device_map"] = kwargs["device_map"]
+        elif kwargs.get("device", None) is not None:
+            device_config["device"] = kwargs["device"]
 
         return device_config
 
@@ -262,7 +272,7 @@ class HuggingFacePipelineHandler(
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
     ) -> Union[huggingface_pipeline.HuggingFacePipelineModel, "transformers.Pipeline"]:
         if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             # We need to redirect the some folders to a writable location in the sandbox.
@@ -292,10 +302,7 @@ class HuggingFacePipelineHandler(
             ) as f:
                 pipeline_params = cloudpickle.load(f)
 
-            if kwargs.get("use_gpu", False):
-                device_config = cls._get_device_config()
-            else:
-                device_config = {}
+            device_config = cls._get_device_config(**kwargs)
 
             m = transformers.pipeline(
                 model_blob_options["task"],
@@ -310,12 +317,8 @@ class HuggingFacePipelineHandler(
             with open(model_blob_file_or_dir_path, "rb") as f:
                 m = cloudpickle.load(f)
             assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
-            if (
-                getattr(m, "device", None) is None
-                and getattr(m, "device_map", None) is None
-                and kwargs.get("use_gpu", False)
-            ):
-                m.__dict__.update(cls._get_device_config())
+            if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
+                m.__dict__.update(cls._get_device_config(**kwargs))
 
             if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
                 m.__dict__.update(torch_dtype="auto")
@@ -326,7 +329,7 @@ class HuggingFacePipelineHandler(
         cls,
         raw_model: Union[huggingface_pipeline.HuggingFacePipelineModel, "transformers.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
     ) -> custom_model.CustomModel:
         import transformers
 

snowflake/ml/model/_packager/model_handlers/lightgbm.py

@@ -43,6 +43,45 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
     MODELE_BLOB_FILE_OR_DIR = "model.pkl"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+    _BINARY_CLASSIFICATION_OBJECTIVES = ["binary"]
+    _MULTI_CLASSIFICATION_OBJECTIVES = ["multiclass", "multiclassova"]
+    _RANKING_OBJECTIVES = ["lambdarank", "rank_xendcg"]
+    _REGRESSION_OBJECTIVES = [
+        "regression",
+        "regression_l1",
+        "huber",
+        "fair",
+        "poisson",
+        "quantile",
+        "tweedie",
+        "mape",
+        "gamma",
+    ]
+
+    @classmethod
+    def get_model_objective(cls, model: Union["lightgbm.Booster", "lightgbm.LGBMModel"]) -> _base.ModelObjective:
+        import lightgbm
+
+        # does not account for cross-entropy and custom
+        if isinstance(model, lightgbm.LGBMClassifier):
+            num_classes = handlers_utils.get_num_classes_if_exists(model)
+            if num_classes == 2:
+                return _base.ModelObjective.BINARY_CLASSIFICATION
+            return _base.ModelObjective.MULTI_CLASSIFICATION
+        if isinstance(model, lightgbm.LGBMRanker):
+            return _base.ModelObjective.RANKING
+        if isinstance(model, lightgbm.LGBMRegressor):
+            return _base.ModelObjective.REGRESSION
+        model_objective = model.params["objective"]
+        if model_objective in cls._BINARY_CLASSIFICATION_OBJECTIVES:
+            return _base.ModelObjective.BINARY_CLASSIFICATION
+        if model_objective in cls._MULTI_CLASSIFICATION_OBJECTIVES:
+            return _base.ModelObjective.MULTI_CLASSIFICATION
+        if model_objective in cls._RANKING_OBJECTIVES:
+            return _base.ModelObjective.RANKING
+        if model_objective in cls._REGRESSION_OBJECTIVES:
+            return _base.ModelObjective.REGRESSION
+        return _base.ModelObjective.UNKNOWN
 
     @classmethod
     def can_handle(
@@ -105,6 +144,19 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
+            if kwargs.get("enable_explainability", False):
+                output_type = model_signature.DataType.DOUBLE
+                if cls.get_model_objective(model) in [
+                    _base.ModelObjective.BINARY_CLASSIFICATION,
+                    _base.ModelObjective.MULTI_CLASSIFICATION,
+                ]:
+                    output_type = model_signature.DataType.STRING
+                model_meta = handlers_utils.add_explain_method_signature(
+                    model_meta=model_meta,
+                    explain_method="explain",
+                    target_method="predict",
+                    output_return_type=output_type,
+                )
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
@@ -130,6 +182,11 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
             ],
             check_local_version=True,
         )
+        if kwargs.get("enable_explainability", False):
+            model_meta.env.include_if_absent(
+                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
+                check_local_version=True,
+            )
 
         return None
 
@@ -139,7 +196,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.LGBMModelLoadOptions],
     ) -> Union["lightgbm.Booster", "lightgbm.LGBMModel"]:
         import lightgbm
 
@@ -169,7 +226,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         cls,
         raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.LGBMModelLoadOptions],
     ) -> custom_model.CustomModel:
         import lightgbm
 
@@ -198,6 +255,17 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
+                @custom_model.inference_api
+                def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    import shap
+
+                    explainer = shap.TreeExplainer(raw_model)
+                    df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                if target_method == "explain":
+                    return explain_fn
+
                 return fn
 
             type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}

snowflake/ml/model/_packager/model_handlers/llm.py

@@ -118,7 +118,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.LLMLoadOptions],
     ) -> llm.LLM:
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         if not hasattr(model_meta, "models"):
@@ -143,7 +143,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
         cls,
         raw_model: llm.LLM,
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.LLMLoadOptions],
     ) -> custom_model.CustomModel:
         import gc
         import tempfile

snowflake/ml/model/_packager/model_handlers/mlflow.py

@@ -160,7 +160,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.MLFlowLoadOptions],
     ) -> "mlflow.pyfunc.PyFuncModel":
         import mlflow
 
@@ -194,7 +194,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         cls,
         raw_model: "mlflow.pyfunc.PyFuncModel",
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.MLFlowLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model
 

snowflake/ml/model/_packager/model_handlers/pytorch.py

@@ -137,7 +137,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.PyTorchLoadOptions],
     ) -> "torch.nn.Module":
         import torch
 
@@ -156,7 +156,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         cls,
         raw_model: "torch.nn.Module",
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.PyTorchLoadOptions],
     ) -> custom_model.CustomModel:
         import torch
 

snowflake/ml/model/_packager/model_handlers/sentence_transformers.py

@@ -126,7 +126,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],  # use_gpu
+        **kwargs: Unpack[model_types.SentenceTransformersLoadOptions],  # use_gpu
     ) -> "sentence_transformers.SentenceTransformer":
         import sentence_transformers
 
@@ -154,7 +154,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         cls,
         raw_model: "sentence_transformers.SentenceTransformer",
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.SentenceTransformersLoadOptions],
     ) -> custom_model.CustomModel:
         import sentence_transformers
 

snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -133,7 +133,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.SKLModelLoadOptions],
     ) -> Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]:
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         model_blobs_metadata = model_meta.models
@@ -153,7 +153,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         cls,
         raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.SKLModelLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model
 

snowflake/ml/model/_packager/model_handlers/snowmlmodel.py

@@ -127,7 +127,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.SNOWModelLoadOptions],
     ) -> "BaseEstimator":
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         model_blobs_metadata = model_meta.models
@@ -146,7 +146,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         cls,
         raw_model: "BaseEstimator",
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.SNOWModelLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model
 

snowflake/ml/model/_packager/model_handlers/tensorflow.py

@@ -138,7 +138,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
     ) -> "tensorflow.Module":
         import tensorflow
 
@@ -156,7 +156,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         cls,
         raw_model: "tensorflow.Module",
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
     ) -> custom_model.CustomModel:
         import tensorflow
 

snowflake/ml/model/_packager/model_handlers/torchscript.py

@@ -128,7 +128,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.TorchScriptLoadOptions],
     ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
         import torch
 
@@ -152,7 +152,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         cls,
         raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.TorchScriptLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model
 

snowflake/ml/model/_packager/model_handlers/xgboost.py

@@ -1,4 +1,5 @@
 # mypy: disable-error-code="import"
+import json
 import os
 from typing import (
     TYPE_CHECKING,
@@ -46,6 +47,39 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
     MODELE_BLOB_FILE_OR_DIR = "model.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+    _BINARY_CLASSIFICATION_OBJECTIVE_PREFIX = ["binary:"]
+    _MULTI_CLASSIFICATION_OBJECTIVE_PREFIX = ["multi:"]
+    _RANKING_OBJECTIVE_PREFIX = ["rank:"]
+    _REGRESSION_OBJECTIVE_PREFIX = ["reg:"]
+
+    @classmethod
+    def get_model_objective(cls, model: Union["xgboost.Booster", "xgboost.XGBModel"]) -> _base.ModelObjective:
+        import xgboost
+
+        if isinstance(model, xgboost.XGBClassifier) or isinstance(model, xgboost.XGBRFClassifier):
+            num_classes = handlers_utils.get_num_classes_if_exists(model)
+            if num_classes == 2:
+                return _base.ModelObjective.BINARY_CLASSIFICATION
+            return _base.ModelObjective.MULTI_CLASSIFICATION
+        if isinstance(model, xgboost.XGBRegressor) or isinstance(model, xgboost.XGBRFRegressor):
+            return _base.ModelObjective.REGRESSION
+        if isinstance(model, xgboost.XGBRanker):
+            return _base.ModelObjective.RANKING
+        model_params = json.loads(model.save_config())
+        model_objective = model_params["learner"]["objective"]
+        for classification_objective in cls._BINARY_CLASSIFICATION_OBJECTIVE_PREFIX:
+            if classification_objective in model_objective:
+                return _base.ModelObjective.BINARY_CLASSIFICATION
+        for classification_objective in cls._MULTI_CLASSIFICATION_OBJECTIVE_PREFIX:
+            if classification_objective in model_objective:
+                return _base.ModelObjective.MULTI_CLASSIFICATION
+        for ranking_objective in cls._RANKING_OBJECTIVE_PREFIX:
+            if ranking_objective in model_objective:
+                return _base.ModelObjective.RANKING
+        for regression_objective in cls._REGRESSION_OBJECTIVE_PREFIX:
+            if regression_objective in model_objective:
+                return _base.ModelObjective.REGRESSION
+        return _base.ModelObjective.UNKNOWN
 
     @classmethod
     def can_handle(
@@ -112,6 +146,16 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
+            if kwargs.get("enable_explainability", False):
+                output_type = model_signature.DataType.DOUBLE
+                if cls.get_model_objective(model) == _base.ModelObjective.MULTI_CLASSIFICATION:
+                    output_type = model_signature.DataType.STRING
+                model_meta = handlers_utils.add_explain_method_signature(
+                    model_meta=model_meta,
+                    explain_method="explain",
+                    target_method="predict",
+                    output_return_type=output_type,
+                )
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
@@ -133,6 +177,11 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
             ],
             check_local_version=True,
         )
+        if kwargs.get("enable_explainability", False):
+            model_meta.env.include_if_absent(
+                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
+                check_local_version=True,
+            )
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
     @classmethod
@@ -141,7 +190,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.XGBModelLoadOptions],
     ) -> Union["xgboost.Booster", "xgboost.XGBModel"]:
         import xgboost
 
@@ -175,7 +224,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         cls,
         raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.XGBModelLoadOptions],
     ) -> custom_model.CustomModel:
         import xgboost
 
@@ -206,6 +255,16 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
+                @custom_model.inference_api
+                def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    import shap
+
+                    explainer = shap.TreeExplainer(raw_model)
+                    df = pd.DataFrame(explainer(X).values)
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                if target_method == "explain":
+                    return explain_fn
                 return fn
 
             type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}

snowflake/ml/model/_packager/model_meta/_core_requirements.py

@@ -6,6 +6,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.15.0,<2",
+    "snowflake-snowpark-python>=1.17.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_packager/model_meta/model_blob_meta.py

@@ -23,6 +23,7 @@ class ModelBlobMeta:
         self.model_type = kwargs["model_type"]
         self.path = kwargs["path"]
         self.handler_version = kwargs["handler_version"]
+        self.function_properties = kwargs.get("function_properties", {})
 
         self.artifacts: Dict[str, str] = {}
         artifacts = kwargs.get("artifacts", None)
@@ -39,6 +40,7 @@ class ModelBlobMeta:
             model_type=self.model_type,
             path=self.path,
             handler_version=self.handler_version,
+            function_properties=self.function_properties,
             artifacts=self.artifacts,
             options=self.options,
         )

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -7,11 +7,12 @@ import zipfile
 from contextlib import contextmanager
 from datetime import datetime
 from types import ModuleType
-from typing import Any, Dict, Generator, List, Optional
+from typing import Any, Dict, Generator, List, Optional, TypedDict
 
 import cloudpickle
 import yaml
 from packaging import requirements, version
+from typing_extensions import Required
 
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
@@ -47,6 +48,7 @@ def create_model_metadata(
     name: str,
     model_type: model_types.SupportedModelHandlerType,
     signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+    function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
     metadata: Optional[Dict[str, str]] = None,
     code_paths: Optional[List[str]] = None,
     ext_modules: Optional[List[ModuleType]] = None,
@@ -64,6 +66,7 @@ def create_model_metadata(
         model_type: Type of the model.
         signatures: Signatures of the model. If None, it will be inferred after the model meta is created.
             Defaults to None.
+        function_properties: Dict mapping function names to a dict of properties, mapping property key to value.
         metadata: User provided key-value metadata of the model. Defaults to None.
         code_paths: List of paths to additional codes that needs to be packed with. Defaults to None.
         ext_modules: List of names of modules that need to be pickled with the model. Defaults to None.
@@ -127,6 +130,7 @@ def create_model_metadata(
         metadata=metadata,
         model_type=model_type,
         signatures=signatures,
+        function_properties=function_properties,
     )
 
     code_dir_path = os.path.join(model_dir_path, MODEL_CODE_DIR)
@@ -215,6 +219,12 @@ def load_code_path(model_dir_path: str) -> None:
         sys.path.insert(0, code_path)
 
 
+class ModelMetadataTelemetryDict(TypedDict):
+    model_name: Required[str]
+    framework_type: Required[model_types.SupportedModelHandlerType]
+    number_of_functions: Required[int]
+
+
 class ModelMetadata:
     """Model metadata for Snowflake native model packaged model.
 
@@ -224,10 +234,18 @@ class ModelMetadata:
         env: ModelEnv object containing all environment related object
         models: Dict of model blob metadata
         signatures: A dict mapping from target function name to input and output signatures.
+        function_properties: A dict mapping function names to dict mapping function property key to value.
         metadata: User provided key-value metadata of the model. Defaults to None.
         creation_timestamp: Unix timestamp when the model metadata is created.
     """
 
+    def telemetry_metadata(self) -> ModelMetadataTelemetryDict:
+        return ModelMetadataTelemetryDict(
+            model_name=self.name,
+            framework_type=self.model_type,
+            number_of_functions=len(self.signatures.keys()),
+        )
+
     def __init__(
         self,
         *,
@@ -236,6 +254,7 @@ class ModelMetadata:
         model_type: model_types.SupportedModelHandlerType,
         runtimes: Optional[Dict[str, model_runtime.ModelRuntime]] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
         metadata: Optional[Dict[str, str]] = None,
         creation_timestamp: Optional[str] = None,
         min_snowpark_ml_version: Optional[str] = None,
@@ -246,6 +265,7 @@ class ModelMetadata:
         self.signatures: Dict[str, model_signature.ModelSignature] = dict()
         if signatures:
             self.signatures = signatures
+        self.function_properties = function_properties or {}
         self.metadata = metadata
         self.model_type = model_type
         self.env = env

snowflake/ml/model/_packager/model_meta/model_meta_schema.py

@@ -1,6 +1,6 @@
 # This files contains schema definition of what will be written into model.yml
 # Changing this file should lead to a change of the schema version.
-
+from enum import Enum
 from typing import Any, Dict, List, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
@@ -11,6 +11,10 @@ MODEL_METADATA_VERSION = "2023-12-01"
 MODEL_METADATA_MIN_SNOWPARK_ML_VERSION = "1.0.12"
 
 
+class FunctionProperties(Enum):
+    PARTITIONED = "PARTITIONED"
+
+
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
@@ -72,6 +76,7 @@ class ModelBlobMetadataDict(TypedDict):
     model_type: Required[type_hints.SupportedModelHandlerType]
     path: Required[str]
     handler_version: Required[str]
+    function_properties: NotRequired[Dict[str, Dict[str, Any]]]
     artifacts: NotRequired[Dict[str, str]]
     options: NotRequired[ModelBlobOptions]
 

snowflake/ml/model/_packager/model_packager.py

@@ -47,12 +47,12 @@ class ModelPackager:
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
         options: Optional[model_types.ModelSaveOption] = None,
-    ) -> None:
+    ) -> model_meta.ModelMetadata:
         if (signatures is None) and (sample_input_data is None) and not model_handler.is_auto_signature_model(model):
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
                 original_exception=ValueError(
-                    "Signatures and sample_input_data both cannot be None at the same time for this kind of model."
+                    "Either of `signatures` or `sample_input_data` must be provided for this kind of model."
                 ),
             )
 
@@ -103,6 +103,7 @@ class ModelPackager:
 
         self.model = model
         self.meta = meta
+        return meta
 
     def load(
         self,
@@ -110,7 +111,7 @@ class ModelPackager:
         meta_only: bool = False,
         as_custom_model: bool = False,
         options: Optional[model_types.ModelLoadOption] = None,
-    ) -> None:
+    ) -> model_meta.ModelMetadata:
         """Load the model into memory from directory. Used internal only.
 
         Args:
@@ -120,11 +121,14 @@ class ModelPackager:
 
         Raises:
             SnowflakeMLException: Raised if model is not native format.
+
+        Returns:
+            Metadata of loaded model.
         """
 
         self.meta = model_meta.ModelMetadata.load(self.local_dir_path)
         if meta_only:
-            return
+            return self.meta
 
         model_meta.load_code_path(self.local_dir_path)
 
@@ -146,3 +150,4 @@ class ModelPackager:
             assert isinstance(m, custom_model.CustomModel)
 
         self.model = m
+        return self.meta

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -5,6 +5,6 @@ REQUIREMENTS = [
     "packaging>=20.9,<24",
     "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.15.0,<2",
+    "snowflake-snowpark-python>=1.17.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]