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

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -4,8 +4,13 @@ from typing import Any, Dict, List, Literal, TypedDict
 
 from typing_extensions import NotRequired, Required
 
+from snowflake.ml.model import model_signature
+
 MODEL_MANIFEST_VERSION = "1.0"
 
+MANIFEST_CLIENT_DATA_KEY_NAME = "snowpark_ml_data"
+MANIFEST_CLIENT_DATA_SCHEMA_VERSION = "2024-02-01"
+
 
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
@@ -38,6 +43,31 @@ class ModelFunctionMethodDict(TypedDict):
 ModelMethodDict = ModelFunctionMethodDict
 
 
+class ModelFunctionInfo(TypedDict):
+    """Function information.
+
+    Attributes:
+        name: Name of the function to be called via SQL.
+        target_method: actual target method name to be called.
+        signature: The signature of the model method.
+    """
+
+    name: Required[str]
+    target_method: Required[str]
+    signature: Required[model_signature.ModelSignature]
+
+
+class ModelFunctionInfoDict(TypedDict):
+    name: Required[str]
+    target_method: Required[str]
+    signature: Required[Dict[str, Any]]
+
+
+class SnowparkMLDataDict(TypedDict):
+    schema_version: Required[str]
+    functions: Required[List[ModelFunctionInfoDict]]
+
+
 class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
     runtimes: Required[Dict[str, ModelRuntimeDict]]

snowflake/ml/model/_model_composer/model_method/function_generator.py

@@ -1,7 +1,6 @@
 import pathlib
 from typing import Optional, TypedDict
 
-import importlib_resources
 from typing_extensions import NotRequired
 
 from snowflake.ml.model import type_hints
@@ -33,6 +32,8 @@ class FunctionGenerator:
         target_method: str,
         options: Optional[FunctionGenerateOptions] = None,
     ) -> None:
+        import importlib_resources
+
         if options is None:
             options = {}
         function_template = (

snowflake/ml/model/_model_composer/model_runtime/_runtime_requirements.py

@@ -1 +1,10 @@
-REQUIREMENTS = ['absl-py>=0.15,<2', 'anyio>=3.5.0,<4', 'numpy>=1.23,<2', 'packaging>=20.9,<24', 'pandas>=1.0.0,<2', 'pyyaml>=6.0,<7', 'snowflake-snowpark-python>=1.8.0,<2', 'typing-extensions>=4.1.0,<5']
+REQUIREMENTS = [
+    "absl-py>=0.15,<2",
+    "anyio>=3.5.0,<4",
+    "numpy>=1.23,<2",
+    "packaging>=20.9,<24",
+    "pandas>=1.0.0,<2",
+    "pyyaml>=6.0,<7",
+    "snowflake-snowpark-python>=1.8.0,<2",
+    "typing-extensions>=4.1.0,<5"
+]

snowflake/ml/model/_model_composer/model_runtime/model_runtime.py

@@ -44,12 +44,17 @@ class ModelRuntime:
         if self.runtime_env._snowpark_ml_version.local:
             self.embed_local_ml_library = True
         else:
-            snowml_server_availability = env_utils.validate_requirements_in_information_schema(
-                session=session,
-                reqs=[requirements.Requirement(snowml_pkg_spec)],
-                python_version=snowml_env.PYTHON_VERSION,
+            snowml_server_availability = (
+                len(
+                    env_utils.get_matched_package_versions_in_information_schema(
+                        session=session,
+                        reqs=[requirements.Requirement(snowml_pkg_spec)],
+                        python_version=snowml_env.PYTHON_VERSION,
+                    ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
+                )
+                >= 1
             )
-            self.embed_local_ml_library = snowml_server_availability is None
+            self.embed_local_ml_library = not snowml_server_availability
 
         if self.embed_local_ml_library:
             self.runtime_env.include_if_absent(
@@ -57,7 +62,6 @@ class ModelRuntime:
                     model_env.ModelDependency(requirement=dep, pip_name=requirements.Requirement(dep).name)
                     for dep in _UDF_INFERENCE_DEPENDENCIES
                 ],
-                check_local_version=True,
             )
         else:
             self.runtime_env.include_if_absent(
@@ -65,7 +69,6 @@ class ModelRuntime:
                     model_env.ModelDependency(requirement=dep, pip_name=requirements.Requirement(dep).name)
                     for dep in _UDF_INFERENCE_DEPENDENCIES + [snowml_pkg_spec]
                 ],
-                check_local_version=True,
             )
 
     def save(self, workspace_path: pathlib.Path) -> model_manifest_schema.ModelRuntimeDict:

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

@@ -59,7 +59,7 @@ def get_requirements_from_task(task: str, spcs_only: bool = False) -> List[model
         return (
             [model_env.ModelDependency(requirement="tokenizers>=0.13.3", pip_name="tokenizers")]
             if spcs_only
-            else [model_env.ModelDependency(requirement="tokenizers<=0.13.2", pip_name="tokenizers")]
+            else [model_env.ModelDependency(requirement="tokenizers", pip_name="tokenizers")]
         )
 
     return []

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

@@ -1,6 +1,16 @@
 # mypy: disable-error-code="import"
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, Union, cast, final
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Type,
+    Union,
+    cast,
+    final,
+)
 
 import numpy as np
 import pandas as pd
@@ -150,6 +160,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         m.load_model(os.path.join(model_blob_path, model_blob_filename))
 
         if kwargs.get("use_gpu", False):
+            assert type(kwargs.get("use_gpu", False)) == bool
             gpu_params = {"tree_method": "gpu_hist", "predictor": "gpu_predictor"}
             if isinstance(m, xgboost.Booster):
                 m.set_param(gpu_params)
@@ -197,7 +208,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
                 return fn
 
-            type_method_dict = {}
+            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

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

@@ -1 +1,11 @@
-REQUIREMENTS = ['absl-py>=0.15,<2', 'anyio>=3.5.0,<4', 'cloudpickle>=2.0.0', 'numpy>=1.23,<2', 'packaging>=20.9,<24', 'pandas>=1.0.0,<2', 'pyyaml>=6.0,<7', 'snowflake-snowpark-python>=1.8.0,<2', 'typing-extensions>=4.1.0,<5']
+REQUIREMENTS = [
+    "absl-py>=0.15,<2",
+    "anyio>=3.5.0,<4",
+    "cloudpickle>=2.0.0",
+    "numpy>=1.23,<2",
+    "packaging>=20.9,<24",
+    "pandas>=1.0.0,<2",
+    "pyyaml>=6.0,<7",
+    "snowflake-snowpark-python>=1.8.0,<2",
+    "typing-extensions>=4.1.0,<5"
+]

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

@@ -0,0 +1,3 @@
+REQUIREMENTS = [
+    "cloudpickle>=2.0.0"
+]

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

@@ -18,6 +18,7 @@ from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_meta import (
     _core_requirements,
+    _packaging_requirements,
     model_blob_meta,
     model_meta_schema,
 )
@@ -26,7 +27,8 @@ from snowflake.ml.model._packager.model_meta_migrator import migrator_plans
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
 
-_PACKAGING_CORE_DEPENDENCIES = _core_requirements.REQUIREMENTS
+_PACKAGING_CORE_DEPENDENCIES = _core_requirements.REQUIREMENTS  # Legacy Model only
+_PACKAGING_REQUIREMENTS = _packaging_requirements.REQUIREMENTS  # New Model only
 _SNOWFLAKE_PKG_NAME = "snowflake"
 _SNOWFLAKE_ML_PKG_NAME = f"{_SNOWFLAKE_PKG_NAME}.ml"
 
@@ -73,6 +75,8 @@ def create_model_metadata(
     model_dir_path = os.path.normpath(model_dir_path)
     embed_local_ml_library = kwargs.pop("embed_local_ml_library", False)
     legacy_save = kwargs.pop("_legacy_save", False)
+    relax_version = kwargs.pop("relax_version", False)
+
     if embed_local_ml_library:
         # Use the last one which is loaded first, that is mean, it is loaded from site-packages.
         # We could make sure that user does not overwrite our library with their code follow the same naming.
@@ -94,6 +98,8 @@ def create_model_metadata(
         pip_requirements=pip_requirements,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
+        legacy_save=legacy_save,
+        relax_version=relax_version,
     )
 
     if embed_local_ml_library:
@@ -146,6 +152,8 @@ def _create_env_for_model_metadata(
     pip_requirements: Optional[List[str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
+    legacy_save: bool = False,
+    relax_version: bool = False,
 ) -> model_env.ModelEnv:
     env = model_env.ModelEnv()
 
@@ -154,11 +162,14 @@ def _create_env_for_model_metadata(
     env.pip_requirements = pip_requirements  # type: ignore[assignment]
     env.python_version = python_version  # type: ignore[assignment]
     env.snowpark_ml_version = snowml_env.VERSION
+
+    requirements_to_add = _PACKAGING_CORE_DEPENDENCIES if legacy_save else _PACKAGING_REQUIREMENTS
+
     if embed_local_ml_library:
         env.include_if_absent(
             [
                 model_env.ModelDependency(requirement=dep, pip_name=requirements.Requirement(dep).name)
-                for dep in _PACKAGING_CORE_DEPENDENCIES
+                for dep in requirements_to_add
             ],
             check_local_version=True,
         )
@@ -166,11 +177,14 @@ def _create_env_for_model_metadata(
         env.include_if_absent(
             [
                 model_env.ModelDependency(requirement=dep, pip_name=requirements.Requirement(dep).name)
-                for dep in _PACKAGING_CORE_DEPENDENCIES + [env_utils.SNOWPARK_ML_PKG_NAME]
+                for dep in requirements_to_add + [env_utils.SNOWPARK_ML_PKG_NAME]
             ],
             check_local_version=True,
         )
 
+    if relax_version:
+        env.relax_version()
+
     return env
 
 

snowflake/ml/model/_signatures/core.py

@@ -146,7 +146,8 @@ class DataType(Enum):
                     " is being automatically converted to INT64 in the Snowpark DataFrame. "
                     "This automatic conversion may lead to potential precision loss and rounding errors. "
                     "If you wish to prevent this conversion, you should manually perform "
-                    "the necessary data type conversion."
+                    "the necessary data type conversion.",
+                    stacklevel=2,
                 )
                 return DataType.INT64
             else:
@@ -155,7 +156,8 @@ class DataType(Enum):
                     " is being automatically converted to DOUBLE in the Snowpark DataFrame. "
                     "This automatic conversion may lead to potential precision loss and rounding errors. "
                     "If you wish to prevent this conversion, you should manually perform "
-                    "the necessary data type conversion."
+                    "the necessary data type conversion.",
+                    stacklevel=2,
                 )
                 return DataType.DOUBLE
         raise snowml_exceptions.SnowflakeMLException(
@@ -202,23 +204,24 @@ class FeatureSpec(BaseFeatureSpec):
         dtype: DataType,
         shape: Optional[Tuple[int, ...]] = None,
     ) -> None:
-        """Initialize a feature.
+        """
+        Initialize a feature.
 
         Args:
             name: Name of the feature.
             dtype: Type of the elements in the feature.
-            shape: Used to represent scalar feature, 1-d feature list or n-d tensor.
-                -1 is used to represent variable length.Defaults to None.
+            shape: Used to represent scalar feature, 1-d feature list,
+                or n-d tensor. Use -1 to represent variable length. Defaults to None.
 
-                E.g.
-                None: scalar
-                (2,): 1d list with fixed len of 2.
-                (-1,): 1d list with variable length. Used for ragged tensor representation.
-                (d1, d2, d3): 3d tensor.
+                Examples:
+                    - None: scalar
+                    - (2,): 1d list with a fixed length of 2.
+                    - (-1,): 1d list with variable length, used for ragged tensor representation.
+                    - (d1, d2, d3): 3d tensor.
 
         Raises:
-            SnowflakeMLException: TypeError: Raised when the dtype input type is incorrect.
-            SnowflakeMLException: TypeError: Raised when the shape input type is incorrect.
+            SnowflakeMLException: TypeError: When the dtype input type is incorrect.
+            SnowflakeMLException: TypeError: When the shape input type is incorrect.
         """
         super().__init__(name=name)
 
@@ -408,13 +411,13 @@ class ModelSignature:
     """Signature of a model that specifies the input and output of a model."""
 
     def __init__(self, inputs: Sequence[BaseFeatureSpec], outputs: Sequence[BaseFeatureSpec]) -> None:
-        """Initialize a model signature
+        """Initialize a model signature.
 
         Args:
-            inputs: A sequence of feature specifications and feature group specifications that will compose the
-                input of the model.
-            outputs: A sequence of feature specifications and feature group specifications that will compose the
-                output of the model.
+            inputs: A sequence of feature specifications and feature group specifications that will compose
+                the input of the model.
+            outputs: A sequence of feature specifications and feature group specifications that will compose
+                the output of the model.
         """
         self._inputs = inputs
         self._outputs = outputs

snowflake/ml/model/custom_model.py

@@ -9,15 +9,16 @@ from snowflake.ml.model import type_hints as model_types
 
 
 class MethodRef:
-    """Represents an method invocation of an instance of `ModelRef`.
+    """Represents a method invocation of an instance of `ModelRef`.
+
+    This allows us to:
+        1) Customize the place of actual execution of the method (inline, thread/process pool, or remote).
+        2) Enrich the way of execution (sync versus async).
 
-    This allows us to
-        1) Customize the place of actual execution of the method(inline, thread/process pool or remote).
-        2) Enrich the way of execution(sync versus async).
     Example:
-        If you have a SKL model, you would normally invoke by `skl_ref.predict(df)` which has sync API.
-        Within inference graph, you could invoke `await skl_ref.predict.async_run(df)` which automatically
-        will be run on thread with async interface.
+        If you have an SKL model, you would normally invoke it by `skl_ref.predict(df)`, which has a synchronous API.
+        Within the inference graph, you could invoke `await skl_ref.predict.async_run(df)`, which will automatically
+        run on a thread with an asynchronous interface.
     """
 
     def __init__(self, model_ref: "ModelRef", method_name: str) -> None:
@@ -27,11 +28,11 @@ class MethodRef:
         return self._func(*args, **kwargs)
 
     async def async_run(self, *args: Any, **kwargs: Any) -> Any:
-        """Run the method in a async way. If the method is defined as async, this will simply run it. If not, this will
-        be run in a separate thread.
+        """Run the method in an asynchronous way. If the method is defined as async, this will simply run it.
+        If not, this will be run in a separate thread.
 
         Args:
-            *args: Arguments of the original method,
+            *args: Arguments of the original method.
             **kwargs: Keyword arguments of the original method.
 
         Returns:
@@ -43,19 +44,20 @@ class MethodRef:
 
 
 class ModelRef:
-    """Represents an model in the inference graph. Method could be directly called using this reference object as if
-    with the original model object.
+    """
+    Represents a model in the inference graph. Methods can be directly called using this reference object
+    as if with the original model object.
 
-    This enables us to separate physical and logical representation of a model which
-    will allows us to deeply understand the graph and perform optimization at entire
-    graph level.
+    This enables us to separate the physical and logical representation of a model, allowing for a deep understanding
+    of the graph and enabling optimization at the entire graph level.
     """
 
     def __init__(self, name: str, model: model_types.SupportedModelType) -> None:
-        """Initialize the ModelRef.
+        """
+        Initialize the ModelRef.
 
         Args:
-            name: The name of a model to refer it.
+            name: The name of the model to refer to.
             model: The model object.
         """
         self._model = model
@@ -91,11 +93,12 @@ class ModelRef:
 
 
 class ModelContext:
-    """Context for a custom model showing path to artifacts and mapping between model name and object reference.
+    """
+    Context for a custom model showing paths to artifacts and mapping between model name and object reference.
 
     Attributes:
-        artifacts: A dict mapping name of the artifact to its path.
-        model_refs: A dict mapping name of the sub-model to its ModelRef object.
+        artifacts: A dictionary mapping the name of the artifact to its path.
+        model_refs: A dictionary mapping the name of the sub-model to its ModelRef object.
     """
 
     def __init__(
@@ -104,11 +107,11 @@ class ModelContext:
         artifacts: Optional[Dict[str, str]] = None,
         models: Optional[Dict[str, model_types.SupportedModelType]] = None,
     ) -> None:
-        """Initialize the model context
+        """Initialize the model context.
 
         Args:
-            artifacts: A dict mapping name of the artifact to its currently available path. Defaults to None.
-            models: A dict mapping name of the sub-model to the corresponding model object. Defaults to None.
+            artifacts: A dictionary mapping the name of the artifact to its currently available path. Defaults to None.
+            models: A dictionary mapping the name of the sub-model to the corresponding model object. Defaults to None.
         """
         self.artifacts: Dict[str, str] = artifacts if artifacts else dict()
         self.model_refs: Dict[str, ModelRef] = (
@@ -116,7 +119,8 @@ class ModelContext:
         )
 
     def path(self, key: str) -> str:
-        """Get the actual path to a specific artifact.
+        """Get the actual path to a specific artifact. This could be used when defining a Custom Model to retrieve
+            artifacts.
 
         Args:
             key: The name of the artifact.
@@ -127,14 +131,13 @@ class ModelContext:
         return self.artifacts[key]
 
     def model_ref(self, name: str) -> ModelRef:
-        """Get a ModelRef object of a sub-model containing the name and model object, while able to call its method
-        directly as well.
+        """Get a ModelRef object of a sub-model containing the name and model object, allowing direct method calls.
 
         Args:
             name: The name of the sub-model.
 
         Returns:
-            The ModelRef object to the sub-model.
+            The ModelRef object representing the sub-model.
         """
         return self.model_refs[name]
 

snowflake/ml/model/model_signature.py

@@ -570,32 +570,31 @@ def infer_signature(
     input_feature_names: Optional[List[str]] = None,
     output_feature_names: Optional[List[str]] = None,
 ) -> core.ModelSignature:
-    """Infer model signature from given input and output sample data.
+    """
+    Infer model signature from given input and output sample data.
+
+    Currently supports inferring model signatures from the following data types:
 
-    Currently, we support infer the model signature from example input/output data in the following cases:
-        - Pandas data frame whose column could have types of supported data types,
-            list (including list of supported data types, list of numpy array of supported data types, and nested list),
-            and numpy array of supported data types.
+        - Pandas DataFrame with columns of supported data types, lists (including nested lists) of supported data types,
+            and NumPy arrays of supported data types.
             - Does not support DataFrame with CategoricalIndex column index.
-            - Does not support DataFrame with column of variant length list or numpy array.
-        - Numpy array of supported data types.
-        - List of Numpy array of supported data types.
-        - List of supported data types, or nested list of supported data types.
-            - Does not support list of list of variant length list.
+        - NumPy arrays of supported data types.
+        - Lists of NumPy arrays of supported data types.
+        - Lists of supported data types or nested lists of supported data types.
+
+    When inferring the signature, a ValueError indicates that the data is insufficient or invalid.
 
-    When a ValueError is raised when inferring the signature, it indicates that the data is ill and it is impossible to
-    create a signature reflecting that.
-    When a NotImplementedError is raised, it indicates that it might be possible to create a signature reflecting the
-    provided data, however, we could not infer it.
+    When it might be possible to create a signature reflecting the provided data, but it could not be inferred,
+    a NotImplementedError is raised
 
     Args:
         input_data: Sample input data for the model.
         output_data: Sample output data for the model.
-        input_feature_names: Name for input features. Defaults to None.
-        output_feature_names: Name for output features. Defaults to None.
+        input_feature_names: Names for input features. Defaults to None.
+        output_feature_names: Names for output features. Defaults to None.
 
     Returns:
-        A model signature.
+        A model signature inferred from the given input and output sample data.
     """
     inputs = _infer_signature(input_data, role="input")
     inputs = utils.rename_features(inputs, input_feature_names)

snowflake/ml/model/type_hints.py

@@ -198,9 +198,12 @@ class BaseModelSaveOption(TypedDict):
     """Options for saving the model.
 
     embed_local_ml_library: Embedding local SnowML into the code directory of the folder.
+    relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. It detects any
+        ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to False.
     """
 
     embed_local_ml_library: NotRequired[bool]
+    relax_version: NotRequired[bool]
     _legacy_save: NotRequired[bool]
     method_options: NotRequired[Dict[str, ModelMethodSaveOptions]]