snowflake-ml-python 1.14.0__py3-none-any.whl → 1.16.0__py3-none-any.whl

snowflake/ml/model/_client/sql/service.py

@@ -1,8 +1,9 @@
+import contextlib
 import dataclasses
 import enum
 import logging
 import textwrap
-from typing import Any, Optional
+from typing import Any, Generator, Optional
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import (
@@ -17,6 +18,11 @@ from snowflake.snowpark._internal import utils as snowpark_utils
 
 logger = logging.getLogger(__name__)
 
+# Using this token instead of '?' to avoid escaping issues
+# After quotes are escaped, we replace this token with '|| ? ||'
+QMARK_RESERVED_TOKEN = "<QMARK_RESERVED_TOKEN>"
+QMARK_PARAMETER_TOKEN = "'|| ? ||'"
+
 
 class ServiceStatus(enum.Enum):
     PENDING = "PENDING"
@@ -70,12 +76,26 @@ class ServiceSQLClient(_base._BaseSQLClient):
     CONTAINER_STATUS = "status"
     MESSAGE = "message"
 
+    @contextlib.contextmanager
+    def _qmark_paramstyle(self) -> Generator[None, None, None]:
+        """Context manager that temporarily changes paramstyle to qmark and restores original value on exit."""
+        if not hasattr(self._session, "_options"):
+            yield
+        else:
+            original_paramstyle = self._session._options["paramstyle"]
+            try:
+                self._session._options["paramstyle"] = "qmark"
+                yield
+            finally:
+                self._session._options["paramstyle"] = original_paramstyle
+
     def deploy_model(
         self,
         *,
         stage_path: Optional[str] = None,
         model_deployment_spec_yaml_str: Optional[str] = None,
         model_deployment_spec_file_rel_path: Optional[str] = None,
+        query_params: Optional[list[Any]] = None,
         statement_params: Optional[dict[str, Any]] = None,
     ) -> tuple[str, snowpark.AsyncJob]:
         assert model_deployment_spec_yaml_str or model_deployment_spec_file_rel_path
@@ -83,11 +103,18 @@ class ServiceSQLClient(_base._BaseSQLClient):
             model_deployment_spec_yaml_str = snowpark_utils.escape_single_quotes(
                 model_deployment_spec_yaml_str
             )  # type: ignore[no-untyped-call]
+            model_deployment_spec_yaml_str = model_deployment_spec_yaml_str.replace(  # type: ignore[union-attr]
+                QMARK_RESERVED_TOKEN, QMARK_PARAMETER_TOKEN
+            )
             logger.info(f"Deploying model with spec={model_deployment_spec_yaml_str}")
             sql_str = f"CALL SYSTEM$DEPLOY_MODEL('{model_deployment_spec_yaml_str}')"
         else:
             sql_str = f"CALL SYSTEM$DEPLOY_MODEL('@{stage_path}/{model_deployment_spec_file_rel_path}')"
-        async_job = self._session.sql(sql_str).collect(block=False, statement_params=statement_params)
+        with self._qmark_paramstyle():
+            async_job = self._session.sql(
+                sql_str,
+                params=query_params if query_params else None,
+            ).collect(block=False, statement_params=statement_params)
         assert isinstance(async_job, snowpark.AsyncJob)
         return async_job.query_id, async_job
 

snowflake/ml/model/_client/sql/stage.py

@@ -2,6 +2,7 @@ from typing import Any, Optional
 
 from snowflake.ml._internal.utils import query_result_checker, sql_identifier
 from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import Row
 
 
 class StageSQLClient(_base._BaseSQLClient):
@@ -21,3 +22,10 @@ class StageSQLClient(_base._BaseSQLClient):
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
         return fq_stage_name
+
+    def list_stage(self, stage_name: str) -> list[Row]:
+        try:
+            list_results = self._session.sql(f"LIST {stage_name}").collect()
+        except Exception as e:
+            raise RuntimeError(f"Failed to check stage location '{stage_name}': {e}")
+        return list_results

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

@@ -46,6 +46,7 @@ class ModelFunctionMethodDict(TypedDict):
     handler: Required[str]
     inputs: Required[list[ModelMethodSignatureFieldWithName]]
     outputs: Required[Union[list[ModelMethodSignatureField], list[ModelMethodSignatureFieldWithName]]]
+    volatility: NotRequired[str]
 
 
 ModelMethodDict = ModelFunctionMethodDict

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

@@ -4,6 +4,7 @@ from typing import Optional, TypedDict, Union
 
 from typing_extensions import NotRequired
 
+from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model import model_signature, type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -12,6 +13,7 @@ from snowflake.ml.model._model_composer.model_method import (
     function_generator,
 )
 from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
+from snowflake.ml.model.volatility import Volatility
 from snowflake.snowpark._internal import type_utils
 
 
@@ -20,10 +22,12 @@ class ModelMethodOptions(TypedDict):
 
     case_sensitive: Specify when the name of the method should be considered as case sensitive when registered to SQL.
     function_type: One of `ModelMethodFunctionTypes` specifying function type.
+    volatility: One of `Volatility` enum values specifying function volatility.
     """
 
     case_sensitive: NotRequired[bool]
     function_type: NotRequired[str]
+    volatility: NotRequired[Volatility]
 
 
 def get_model_method_options_from_options(
@@ -38,10 +42,19 @@ def get_model_method_options_from_options(
     if function_type not in [function_type.value for function_type in model_manifest_schema.ModelMethodFunctionTypes]:
         raise NotImplementedError(f"Function type {function_type} is not supported.")
 
-    return ModelMethodOptions(
+    default_volatility = options.get("volatility")
+    method_volatility = method_option.get("volatility")
+    resolved_volatility = method_volatility or default_volatility
+
+    # Only include volatility if explicitly provided in method options
+    result: ModelMethodOptions = ModelMethodOptions(
         case_sensitive=method_option.get("case_sensitive", False),
         function_type=function_type,
     )
+    if resolved_volatility:
+        result["volatility"] = resolved_volatility
+
+    return result
 
 
 class ModelMethod:
@@ -94,6 +107,9 @@ class ModelMethod:
             "function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
         )
 
+        # Volatility is optional; when not provided, we omit it from the manifest
+        self.volatility = self.options.get("volatility")
+
     @staticmethod
     def _get_method_arg_from_feature(
         feature: model_signature.BaseFeatureSpec, case_sensitive: bool = False
@@ -148,7 +164,7 @@ class ModelMethod:
         else:
             outputs = [model_manifest_schema.ModelMethodSignatureField(type="OBJECT")]
 
-        return model_manifest_schema.ModelFunctionMethodDict(
+        method_dict = model_manifest_schema.ModelFunctionMethodDict(
             name=self.method_name.resolved(),
             runtime=self.runtime_name,
             type=self.function_type,
@@ -158,3 +174,10 @@ class ModelMethod:
             inputs=input_list,
             outputs=outputs,
         )
+        should_set_volatility = (
+            platform_capabilities.PlatformCapabilities.get_instance().is_set_module_functions_volatility_from_manifest()
+        )
+        if should_set_volatility and self.volatility is not None:
+            method_dict["volatility"] = self.volatility.name
+
+        return method_dict

snowflake/ml/model/_packager/model_env/model_env.py

@@ -145,11 +145,12 @@ class ModelEnv:
         """
         if (self.pip_requirements or self.prefer_pip) and not self.conda_dependencies and pkgs:
             pip_pkg_reqs: list[str] = []
-            if self.targets_warehouse:
+            if self.targets_warehouse and not self.artifact_repository_map:
                 self._warn_once(
                     (
                         "Dependencies specified from pip requirements."
                         " This may prevent model deploying to Snowflake Warehouse."
+                        " Use 'artifact_repository_map' to deploy the model to Warehouse."
                     ),
                     stacklevel=2,
                 )
@@ -177,7 +178,11 @@ class ModelEnv:
                 req_to_add.name = conda_req.name
             else:
                 req_to_add = conda_req
-            show_warning_message = conda_req_channel == env_utils.DEFAULT_CHANNEL_NAME and self.targets_warehouse
+            show_warning_message = (
+                conda_req_channel == env_utils.DEFAULT_CHANNEL_NAME
+                and self.targets_warehouse
+                and not self.artifact_repository_map
+            )
 
             if any(added_pip_req.name == pip_name for added_pip_req in self._pip_requirements):
                 if show_warning_message:
@@ -185,6 +190,7 @@ class ModelEnv:
                         (
                             f"Basic dependency {req_to_add.name} specified from pip requirements."
                             " This may prevent model deploying to Snowflake Warehouse."
+                            " Use 'artifact_repository_map' to deploy the model to Warehouse."
                         ),
                         stacklevel=2,
                     )
@@ -318,13 +324,15 @@ class ModelEnv:
                     )
 
         if pip_requirements_list and self.targets_warehouse:
-            self._warn_once(
-                (
-                    "Found dependencies specified as pip requirements."
-                    " This may prevent model deploying to Snowflake Warehouse."
-                ),
-                stacklevel=2,
-            )
+            if not self.artifact_repository_map:
+                self._warn_once(
+                    (
+                        "Found dependencies specified as pip requirements."
+                        " This may prevent model deploying to Snowflake Warehouse."
+                        " Use 'artifact_repository_map' to deploy the model to Warehouse."
+                    ),
+                    stacklevel=2,
+                )
             for pip_dependency in pip_requirements_list:
                 if any(
                     channel_dependency.name == pip_dependency.name
@@ -343,13 +351,15 @@ class ModelEnv:
         pip_requirements_list = env_utils.load_requirements_file(pip_requirements_path)
 
         if pip_requirements_list and self.targets_warehouse:
-            self._warn_once(
-                (
-                    "Found dependencies specified as pip requirements."
-                    " This may prevent model deploying to Snowflake Warehouse."
-                ),
-                stacklevel=2,
-            )
+            if not self.artifact_repository_map:
+                self._warn_once(
+                    (
+                        "Found dependencies specified as pip requirements."
+                        " This may prevent model deploying to Snowflake Warehouse."
+                        " Use 'artifact_repository_map' to deploy the model to Warehouse."
+                    ),
+                    stacklevel=2,
+                )
             for pip_dependency in pip_requirements_list:
                 if any(
                     channel_dependency.name == pip_dependency.name

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

@@ -1,5 +1,6 @@
 import importlib
 import json
+import logging
 import os
 import pathlib
 import warnings
@@ -8,7 +9,6 @@ from typing import Any, Callable, Iterable, Optional, Sequence, cast
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
-from absl import logging
 
 import snowflake.snowpark.dataframe as sp_df
 from snowflake.ml._internal import env
@@ -23,6 +23,8 @@ from snowflake.ml.model._signatures import (
 )
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
 
+logger = logging.getLogger(__name__)
+
 EXPLAIN_BACKGROUND_DATA_ROWS_COUNT_LIMIT = 1000
 
 
@@ -257,7 +259,7 @@ def validate_model_task(passed_model_task: model_types.Task, inferred_model_task
             )
         return inferred_model_task
     elif inferred_model_task != model_types.Task.UNKNOWN:
-        logging.info(f"Inferred Task: {inferred_model_task.name} is used as task for this model " f"version")
+        logger.info(f"Inferred Task: {inferred_model_task.name} is used as task for this model " f"version")
         return inferred_model_task
     return passed_model_task
 

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

@@ -43,7 +43,6 @@ DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message
 def get_requirements_from_task(task: str, spcs_only: bool = False) -> list[model_env.ModelDependency]:
     # Text
     if task in [
-        "conversational",
         "fill-mask",
         "ner",
         "token-classification",
@@ -521,6 +520,7 @@ class HuggingFacePipelineHandler(
                                 input_data = X[signature.inputs[0].name].to_list()
                             temp_res = getattr(raw_model, target_method)(input_data)
                     else:
+                        # TODO: remove conversational pipeline code
                         # For others, we could offer the whole dataframe as a list.
                         # Some of them may need some conversion
                         if hasattr(transformers, "ConversationalPipeline") and isinstance(
@@ -759,11 +759,13 @@ class HuggingFaceOpenAICompatibleModel:
             eos_token_id=self.tokenizer.eos_token_id,
             stop_strings=stop_strings,
             stream=stream,
-            repetition_penalty=frequency_penalty,
-            diversity_penalty=presence_penalty if n > 1 else None,
             num_return_sequences=n,
-            num_beams=max(2, n),  # must be >1
-            num_beam_groups=max(2, n) if presence_penalty else 1,
+            num_beams=max(1, n),  # must be >1
+            repetition_penalty=frequency_penalty,
+            # TODO: Handle diversity_penalty and num_beam_groups
+            # not all models support them making it hard to support any huggingface model
+            # diversity_penalty=presence_penalty if n > 1 else None,
+            # num_beam_groups=max(2, n) if presence_penalty else 1,
             do_sample=False,
         )
 

snowflake/ml/model/_packager/model_packager.py

@@ -1,9 +1,8 @@
+import logging
 import os
 from types import ModuleType
 from typing import Optional
 
-from absl import logging
-
 from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
@@ -12,6 +11,8 @@ from snowflake.ml.model import custom_model, model_signature, type_hints as mode
 from snowflake.ml.model._packager import model_handler
 from snowflake.ml.model._packager.model_meta import model_meta
 
+logger = logging.getLogger(__name__)
+
 
 class ModelPackager:
     """Top-level class to save/load and manage a Snowflake Native formatted model.
@@ -96,7 +97,7 @@ class ModelPackager:
                 **options,
             )
             if signatures is None:
-                logging.info(f"Model signatures are auto inferred as:\n\n{meta.signatures}")
+                logger.info(f"Model signatures are auto inferred as:\n\n{meta.signatures}")
 
         self.model = model
         self.meta = meta

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

@@ -2,7 +2,6 @@
 # Generate by running 'bazel run --config=pre_build //bazel/requirements:sync_requirements'
 
 REQUIREMENTS = [
-    "absl-py>=0.15,<2",
     "aiohttp!=4.0.0a0, !=4.0.0a1",
     "anyio>=3.5.0,<5",
     "cachetools>=3.1.1,<6",
@@ -22,7 +21,7 @@ REQUIREMENTS = [
     "requests",
     "retrying>=1.3.3,<2",
     "s3fs>=2024.6.1,<2026",
-    "scikit-learn<1.7",
+    "scikit-learn<1.8",
     "scipy>=1.9,<2",
     "shap>=0.46.0,<1",
     "snowflake-connector-python>=3.16.0,<4",

snowflake/ml/model/_signatures/utils.py

@@ -110,27 +110,6 @@ def huggingface_pipeline_signature_auto_infer(
 ) -> Optional[core.ModelSignature]:
     # Text
 
-    # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ConversationalPipeline
-    # Needs to convert to conversation object.
-    if task == "conversational":
-        warnings.warn(
-            (
-                "Conversational pipeline is removed from transformers since 4.42.0. "
-                "Support will be removed from snowflake-ml-python soon."
-            ),
-            category=DeprecationWarning,
-            stacklevel=1,
-        )
-        return core.ModelSignature(
-            inputs=[
-                core.FeatureSpec(name="user_inputs", dtype=core.DataType.STRING, shape=(-1,)),
-                core.FeatureSpec(name="generated_responses", dtype=core.DataType.STRING, shape=(-1,)),
-            ],
-            outputs=[
-                core.FeatureSpec(name="generated_responses", dtype=core.DataType.STRING, shape=(-1,)),
-            ],
-        )
-
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.TokenClassificationPipeline
     if task == "fill-mask":
         return core.ModelSignature(

snowflake/ml/model/models/huggingface_pipeline.py

@@ -8,6 +8,7 @@ from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._client.model import inference_engine_utils
 from snowflake.ml.model._client.ops import service_ops
 from snowflake.snowpark import async_job, session
 
@@ -77,6 +78,15 @@ class HuggingFacePipelineModel:
         framework = kwargs.get("framework", None)
         feature_extractor = kwargs.get("feature_extractor", None)
 
+        _can_download_snapshot = False
+        if download_snapshot:
+            try:
+                import huggingface_hub as hf_hub
+
+                _can_download_snapshot = True
+            except ImportError:
+                pass
+
         # ==== Start pipeline logic from transformers ====
         if model_kwargs is None:
             model_kwargs = {}
@@ -141,22 +151,23 @@ class HuggingFacePipelineModel:
         # Instantiate config if needed
         config_obj = None
 
-        if isinstance(config, str):
-            config_obj = transformers.AutoConfig.from_pretrained(
-                config, _from_pipeline=task, **hub_kwargs, **model_kwargs
-            )
-            hub_kwargs["_commit_hash"] = config_obj._commit_hash
-        elif config is None and isinstance(model, str):
-            config_obj = transformers.AutoConfig.from_pretrained(
-                model, _from_pipeline=task, **hub_kwargs, **model_kwargs
-            )
-            hub_kwargs["_commit_hash"] = config_obj._commit_hash
-        # We only support string as config argument.
-        elif config is not None and not isinstance(config, str):
-            raise RuntimeError(
-                "Impossible to use non-string config as input for HuggingFacePipelineModel. Use transformers.Pipeline"
-                " object if required."
-            )
+        if not _can_download_snapshot:
+            if isinstance(config, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    config, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            elif config is None and isinstance(model, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            # We only support string as config argument.
+            elif config is not None and not isinstance(config, str):
+                raise RuntimeError(
+                    "Impossible to use non-string config as input for HuggingFacePipelineModel. "
+                    "Use transformers.Pipeline object if required."
+                )
 
         # ==== Start pipeline logic (Task) from transformers ====
 
@@ -208,7 +219,7 @@ class HuggingFacePipelineModel:
                 "Using a pipeline without specifying a model name and revision in production is not recommended.",
                 stacklevel=2,
             )
-            if config is None and isinstance(model, str):
+            if not _can_download_snapshot and config is None and isinstance(model, str):
                 config_obj = transformers.AutoConfig.from_pretrained(
                     model, _from_pipeline=task, **hub_kwargs, **model_kwargs
                 )
@@ -228,11 +239,10 @@ class HuggingFacePipelineModel:
                 )
 
         repo_snapshot_dir: Optional[str] = None
-        if download_snapshot:
+        if _can_download_snapshot:
             try:
-                from huggingface_hub import snapshot_download
 
-                repo_snapshot_dir = snapshot_download(
+                repo_snapshot_dir = hf_hub.snapshot_download(
                     repo_id=model,
                     revision=revision,
                     token=token,
@@ -268,7 +278,7 @@ class HuggingFacePipelineModel:
         ],
     )
     @snowpark._internal.utils.private_preview(version="1.9.1")
-    def create_service(
+    def log_model_and_create_service(
         self,
         *,
         session: session.Session,
@@ -293,6 +303,7 @@ class HuggingFacePipelineModel:
         force_rebuild: bool = False,
         build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
+        experimental_options: Optional[dict[str, Any]] = None,
     ) -> Union[str, async_job.AsyncJob]:
         """Logs a Hugging Face model and creates a service in Snowflake.
 
@@ -319,6 +330,10 @@ class HuggingFacePipelineModel:
             force_rebuild: Whether to force rebuild the image. Defaults to False.
             build_external_access_integrations: External access integrations for building the image. Defaults to None.
             block: Whether to block the operation. Defaults to True.
+            experimental_options: Experimental options for the service creation with custom inference engine.
+                Currently, only `inference_engine` and `inference_engine_args_override` are supported.
+                `inference_engine` is the name of the inference engine to use.
+                `inference_engine_args_override` is a list of string arguments to pass to the inference engine.
 
         Raises:
             ValueError: if database and schema name is not provided and session doesn't have a
@@ -360,6 +375,24 @@ class HuggingFacePipelineModel:
         )
         logger.info(f"A service job is going to register the hf model as: {model_name}.{version_name}")
 
+        # Check if model is HuggingFace text-generation before doing inference engine checks
+        inference_engine_args = None
+        if experimental_options:
+            if self.task != "text-generation":
+                raise ValueError(
+                    "Currently, InferenceEngine using experimental_options is only supported for "
+                    "HuggingFace text-generation models."
+                )
+
+            inference_engine_args = inference_engine_utils._get_inference_engine_args(experimental_options)
+
+            # Enrich inference engine args if inference engine is specified
+            if inference_engine_args is not None:
+                inference_engine_args = inference_engine_utils._enrich_inference_engine_args(
+                    inference_engine_args,
+                    gpu_requests,
+                )
+
         from snowflake.ml.model import event_handler
         from snowflake.snowpark import exceptions
 
@@ -412,6 +445,8 @@ class HuggingFacePipelineModel:
                         # TODO: remove warehouse in the next release
                         warehouse=session.get_current_warehouse(),
                     ),
+                    # inference engine
+                    inference_engine_args=inference_engine_args,
                 )
                 status.update(label="HuggingFace model service created successfully", state="complete", expanded=False)
                 return result

snowflake/ml/model/type_hints.py

@@ -15,6 +15,7 @@ from typing_extensions import NotRequired
 
 from snowflake.ml.model.target_platform import TargetPlatform
 from snowflake.ml.model.task import Task
+from snowflake.ml.model.volatility import Volatility
 
 if TYPE_CHECKING:
     import catboost
@@ -150,6 +151,7 @@ class ModelMethodSaveOptions(TypedDict):
     case_sensitive: NotRequired[bool]
     max_batch_size: NotRequired[int]
     function_type: NotRequired[Literal["FUNCTION", "TABLE_FUNCTION"]]
+    volatility: NotRequired[Volatility]
 
 
 class BaseModelSaveOption(TypedDict):
@@ -158,12 +160,23 @@ class BaseModelSaveOption(TypedDict):
     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 in 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 method_options.
+    volatility: Set the volatility for all model methods globally. To set volatility for individual methods
+        see volatility in method_options. Defaults are set automatically based on model type: supported
+        models (sklearn, xgboost, pytorch, huggingface_pipeline, mlflow, etc.) default to IMMUTABLE, while
+        custom models default to VOLATILE. When both global volatility and per-method volatility are specified,
+        the per-method volatility takes precedence.
+    method_options: Per-method saving options. This dictionary has method names as keys and dictionary
+        values with the desired options.
+    enable_explainability: Whether to enable explainability features for the model.
     save_location: Local directory path to save the model and metadata.
     """
 
     embed_local_ml_library: NotRequired[bool]
     relax_version: NotRequired[bool]
     function_type: NotRequired[Literal["FUNCTION", "TABLE_FUNCTION"]]
+    volatility: NotRequired[Volatility]
     method_options: NotRequired[dict[str, ModelMethodSaveOptions]]
     enable_explainability: NotRequired[bool]
     save_location: NotRequired[str]

snowflake/ml/model/volatility.py

@@ -0,0 +1,34 @@
+"""Volatility definitions for model functions."""
+
+from enum import Enum, auto
+
+
+class Volatility(Enum):
+    """Volatility levels for model functions.
+
+    Attributes:
+        VOLATILE: Function results may change between calls with the same arguments.
+            Use this for functions that depend on external data or have non-deterministic behavior.
+        IMMUTABLE: Function results are guaranteed to be the same for the same arguments.
+            Use this for pure functions that always return the same output for the same input.
+    """
+
+    VOLATILE = auto()
+    IMMUTABLE = auto()
+
+
+DEFAULT_VOLATILITY_BY_MODEL_TYPE = {
+    "catboost": Volatility.IMMUTABLE,
+    "custom": Volatility.VOLATILE,
+    "huggingface_pipeline": Volatility.IMMUTABLE,
+    "keras": Volatility.IMMUTABLE,
+    "lightgbm": Volatility.IMMUTABLE,
+    "mlflow": Volatility.IMMUTABLE,
+    "pytorch": Volatility.IMMUTABLE,
+    "sentence_transformers": Volatility.IMMUTABLE,
+    "sklearn": Volatility.IMMUTABLE,
+    "snowml": Volatility.IMMUTABLE,
+    "tensorflow": Volatility.IMMUTABLE,
+    "torchscript": Volatility.IMMUTABLE,
+    "xgboost": Volatility.IMMUTABLE,
+}

snowflake/ml/modeling/calibration/calibrated_classifier_cv.py

@@ -60,7 +60,7 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
-SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.7')
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.8')
 # Modeling library estimators require a smaller sklearn version range.
 if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
     raise Exception(

snowflake/ml/modeling/cluster/affinity_propagation.py

@@ -60,7 +60,7 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
-SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.7')
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.8')
 # Modeling library estimators require a smaller sklearn version range.
 if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
     raise Exception(

snowflake/ml/modeling/cluster/agglomerative_clustering.py

@@ -60,7 +60,7 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
-SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.7')
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.8')
 # Modeling library estimators require a smaller sklearn version range.
 if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
     raise Exception(

snowflake/ml/modeling/cluster/birch.py

@@ -60,7 +60,7 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
-SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.7')
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.8')
 # Modeling library estimators require a smaller sklearn version range.
 if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
     raise Exception(