snowflake-ml-python 1.5.2__py3-none-any.whl → 1.5.4__py3-none-any.whl

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

@@ -17,6 +17,7 @@ from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
+    model_meta_schema,
 )
 
 
@@ -68,6 +69,11 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 predictions_df = target_method(model, sample_input_data)
             return predictions_df
 
+        for func_name in model._get_partitioned_infer_methods():
+            function_properties = model_meta.function_properties.get(func_name, {})
+            function_properties[model_meta_schema.FunctionProperties.PARTITIONED.value] = True
+            model_meta.function_properties[func_name] = function_properties
+
         if not is_sub_model:
             model_meta = handlers_utils.validate_signature(
                 model=model,
@@ -101,14 +107,16 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
 
         # Make sure that the module where the model is defined get pickled by value as well.
         cloudpickle.register_pickle_by_value(sys.modules[model.__module__])
-        picked_obj = (model.__class__, model.context)
+        pickled_obj = (model.__class__, model.context)
         with open(os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR), "wb") as f:
-            cloudpickle.dump(picked_obj, f)
+            cloudpickle.dump(pickled_obj, f)
+        # model meta will be saved by the context manager
         model_meta.models[name] = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             path=cls.MODELE_BLOB_FILE_OR_DIR,
             handler_version=cls.HANDLER_VERSION,
+            function_properties=model_meta.function_properties,
             artifacts={
                 name: pathlib.Path(
                     os.path.join(cls.MODEL_ARTIFACTS_DIR, os.path.basename(os.path.normpath(path=uri)))
@@ -128,7 +136,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         name: str,
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.CustomModelLoadOption],
     ) -> "custom_model.CustomModel":
         model_blob_path = os.path.join(model_blobs_dir_path, name)
 
@@ -175,6 +183,6 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         cls,
         raw_model: custom_model.CustomModel,
         model_meta: model_meta_api.ModelMetadata,
-        **kwargs: Unpack[model_types.ModelLoadOption],
+        **kwargs: Unpack[model_types.CustomModelLoadOption],
     ) -> custom_model.CustomModel:
         return raw_model

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

@@ -139,7 +139,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 +169,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
 

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

@@ -141,7 +141,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 +175,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
 

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.11.1,<2,!=1.12.0",
+    "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.11.1,<2,!=1.12.0",
+    "snowflake-snowpark-python>=1.17.0,<2",
     "typing-extensions>=4.1.0,<5"
 ]

snowflake/ml/model/_signatures/builtins_handler.py

@@ -1,3 +1,4 @@
+import datetime
 from collections import abc
 from typing import Literal, Sequence
 
@@ -24,7 +25,7 @@ class ListOfBuiltinHandler(base_handler.BaseDataHandler[model_types._SupportedBu
             # String is a Sequence but we take them as an whole
             if isinstance(element, abc.Sequence) and not isinstance(element, str):
                 can_handle = ListOfBuiltinHandler.can_handle(element)
-            elif not isinstance(element, (int, float, bool, str)):
+            elif not isinstance(element, (int, float, bool, str, datetime.datetime)):
                 can_handle = False
                 break
         return can_handle

snowflake/ml/model/_signatures/core.py

@@ -53,6 +53,8 @@ class DataType(Enum):
     STRING = ("string", spt.StringType, np.str_)
     BYTES = ("bytes", spt.BinaryType, np.bytes_)
 
+    TIMESTAMP_NTZ = ("datetime64[ns]", spt.TimestampType, "datetime64[ns]")
+
     def as_snowpark_type(self) -> spt.DataType:
         """Convert to corresponding Snowpark Type.
 
@@ -78,6 +80,13 @@ class DataType(Enum):
             Corresponding DataType.
         """
         np_to_snowml_type_mapping = {i._numpy_type: i for i in DataType}
+
+        # Add datetime types:
+        datetime_res = ["Y", "M", "W", "D", "h", "m", "s", "ms", "us", "ns"]
+
+        for res in datetime_res:
+            np_to_snowml_type_mapping[f"datetime64[{res}]"] = DataType.TIMESTAMP_NTZ
+
         for potential_type in np_to_snowml_type_mapping.keys():
             if np.can_cast(np_type, potential_type, casting="no"):
                 # This is used since the same dtype might represented in different ways.
@@ -247,9 +256,12 @@ class FeatureSpec(BaseFeatureSpec):
             result_type = spt.ArrayType(result_type)
         return result_type
 
-    def as_dtype(self) -> npt.DTypeLike:
+    def as_dtype(self) -> Union[npt.DTypeLike, str]:
         """Convert to corresponding local Type."""
         if not self._shape:
+            # scalar dtype: use keys from `np.sctypeDict` to prevent unit-less dtype 'datetime64'
+            if "datetime64" in self._dtype._value:
+                return self._dtype._value
             return self._dtype._numpy_type
         return np.object_
 

snowflake/ml/model/_signatures/pandas_handler.py

@@ -147,6 +147,8 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
                     specs.append(core.FeatureSpec(dtype=core.DataType.STRING, name=ft_name))
                 elif isinstance(data[df_col].iloc[0], bytes):
                     specs.append(core.FeatureSpec(dtype=core.DataType.BYTES, name=ft_name))
+            elif isinstance(data[df_col].iloc[0], np.datetime64):
+                specs.append(core.FeatureSpec(dtype=core.DataType.TIMESTAMP_NTZ, name=ft_name))
             else:
                 specs.append(core.FeatureSpec(dtype=core.DataType.from_numpy_type(df_col_dtype), name=ft_name))
         return specs

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -107,6 +107,9 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         if not features:
             features = pandas_handler.PandasDataFrameHandler.infer_signature(df, role="input")
         # Role will be no effect on the column index. That is to say, the feature name is the actual column name.
+        if keep_order:
+            df = df.reset_index(drop=True)
+            df[infer_template._KEEP_ORDER_COL_NAME] = df.index
         sp_df = session.create_dataframe(df)
         column_names = []
         columns = []
@@ -122,7 +125,4 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
 
         sp_df = sp_df.with_columns(column_names, columns)
 
-        if keep_order:
-            sp_df = sp_df.with_column(infer_template._KEEP_ORDER_COL_NAME, F.monotonically_increasing_id())
-
         return sp_df

snowflake/ml/model/custom_model.py

@@ -1,6 +1,6 @@
 import functools
 import inspect
-from typing import Any, Callable, Coroutine, Dict, Generator, Optional
+from typing import Any, Callable, Coroutine, Dict, Generator, List, Optional
 
 import anyio
 import pandas as pd
@@ -168,7 +168,7 @@ class CustomModel:
     def _get_infer_methods(
         self,
     ) -> Generator[Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame], None, None]:
-        """Returns all methods in CLS with DECORATOR as the outermost decorator."""
+        """Returns all methods in CLS with `inference_api` decorator as the outermost decorator."""
         for cls_method_str in dir(self):
             cls_method = getattr(self, cls_method_str)
             if getattr(cls_method, "_is_inference_api", False):
@@ -177,6 +177,18 @@ class CustomModel:
                 else:
                     raise TypeError("A non-method inference API function is not supported.")
 
+    def _get_partitioned_infer_methods(self) -> List[str]:
+        """Returns all methods in CLS with `partitioned_inference_api` as the outermost decorator."""
+        rv = []
+        for cls_method_str in dir(self):
+            cls_method = getattr(self, cls_method_str)
+            if getattr(cls_method, "_is_partitioned_inference_api", False):
+                if inspect.ismethod(cls_method):
+                    rv.append(cls_method_str)
+                else:
+                    raise TypeError("A non-method inference API function is not supported.")
+        return rv
+
 
 def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]) -> None:
     """Validate the user provided predict method.
@@ -219,3 +231,11 @@ def inference_api(
 ) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     return func
+
+
+def partitioned_inference_api(
+    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]
+) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func.__dict__["_is_inference_api"] = True
+    func.__dict__["_is_partitioned_inference_api"] = True
+    return func

snowflake/ml/model/model_signature.py

@@ -168,6 +168,8 @@ def _validate_numpy_array(
             max_v <= np.finfo(feature_type._numpy_type).max  # type: ignore[arg-type]
             and min_v >= np.finfo(feature_type._numpy_type).min  # type: ignore[arg-type]
         )
+    elif feature_type in [core.DataType.TIMESTAMP_NTZ]:
+        return np.issubdtype(arr.dtype, np.datetime64)
     else:
         return np.can_cast(arr.dtype, feature_type._numpy_type, casting="no")