snowflake-ml-python 1.15.0__py3-none-any.whl → 1.17.0__py3-none-any.whl

snowflake/ml/model/_client/ops/model_ops.py

@@ -952,6 +952,7 @@ class ModelOperator:
         partition_column: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
+        explain_case_sensitive: bool = False,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -967,6 +968,7 @@ class ModelOperator:
         service_name: sql_identifier.SqlIdentifier,
         strict_input_validation: bool = False,
         statement_params: Optional[dict[str, str]] = None,
+        explain_case_sensitive: bool = False,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -986,6 +988,7 @@ class ModelOperator:
         partition_column: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
+        explain_case_sensitive: bool = False,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
@@ -1068,6 +1071,7 @@ class ModelOperator:
                     version_name=version_name,
                     statement_params=statement_params,
                     is_partitioned=is_partitioned or False,
+                    explain_case_sensitive=explain_case_sensitive,
                 )
 
         if keep_order:

snowflake/ml/model/_client/ops/service_ops.py

@@ -7,6 +7,7 @@ import re
 import tempfile
 import threading
 import time
+import warnings
 from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
@@ -14,6 +15,7 @@ from snowflake.ml import jobs
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
 from snowflake.ml._internal.utils import identifier, service_logger, sql_identifier
 from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
+from snowflake.ml.model._client.model import batch_inference_specs
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
 from snowflake.snowpark import async_job, exceptions, row, session
@@ -155,17 +157,17 @@ class ServiceOperator:
             database_name=database_name,
             schema_name=schema_name,
         )
+        self._stage_client = stage_sql.StageSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
         self._use_inlined_deployment_spec = pc.PlatformCapabilities.get_instance().is_inlined_deployment_spec_enabled()
         if self._use_inlined_deployment_spec:
             self._workspace = None
             self._model_deployment_spec = model_deployment_spec.ModelDeploymentSpec()
         else:
             self._workspace = tempfile.TemporaryDirectory()
-            self._stage_client = stage_sql.StageSQLClient(
-                session,
-                database_name=database_name,
-                schema_name=schema_name,
-            )
             self._model_deployment_spec = model_deployment_spec.ModelDeploymentSpec(
                 workspace_path=pathlib.Path(self._workspace.name)
             )
@@ -651,6 +653,47 @@ class ServiceOperator:
             else:
                 module_logger.warning(f"Service {service.display_service_name} is done, but not transitioning.")
 
+    def _enforce_save_mode(self, output_mode: batch_inference_specs.SaveMode, output_stage_location: str) -> None:
+        """Enforce the save mode for the output stage location.
+
+        Args:
+            output_mode: The output mode
+            output_stage_location: The output stage location to check/clean.
+
+        Raises:
+            FileExistsError: When ERROR mode is specified and files exist in the output location.
+            RuntimeError: When operations fail (checking files or removing files).
+            ValueError: When an invalid SaveMode is specified.
+        """
+        list_results = self._stage_client.list_stage(output_stage_location)
+
+        if output_mode == batch_inference_specs.SaveMode.ERROR:
+            if len(list_results) > 0:
+                raise FileExistsError(
+                    f"Output stage location '{output_stage_location}' is not empty. "
+                    f"Found {len(list_results)} existing files. When using ERROR mode, the output location "
+                    f"must be empty. Please clear the existing files or use OVERWRITE mode."
+                )
+        elif output_mode == batch_inference_specs.SaveMode.OVERWRITE:
+            if len(list_results) > 0:
+                warnings.warn(
+                    f"Output stage location '{output_stage_location}' is not empty. "
+                    f"Found {len(list_results)} existing files. OVERWRITE mode will remove all existing files "
+                    f"in the output location before running the batch inference job.",
+                    stacklevel=2,
+                )
+                try:
+                    self._session.sql(f"REMOVE {output_stage_location}").collect()
+                except Exception as e:
+                    raise RuntimeError(
+                        f"OVERWRITE was specified. However, failed to remove existing files in output stage "
+                        f"{output_stage_location}: {e}. Please clear up the existing files manually and retry "
+                        f"the operation."
+                    )
+        else:
+            valid_modes = list(batch_inference_specs.SaveMode)
+            raise ValueError(f"Invalid SaveMode: {output_mode}. Must be one of {valid_modes}")
+
     def _stream_service_logs(
         self,
         async_job: snowpark.AsyncJob,
@@ -927,6 +970,7 @@ class ServiceOperator:
         max_batch_rows: Optional[int],
         cpu_requests: Optional[str],
         memory_requests: Optional[str],
+        gpu_requests: Optional[str],
         replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
     ) -> jobs.MLJob[Any]:
@@ -961,6 +1005,7 @@ class ServiceOperator:
             warehouse=warehouse,
             cpu=cpu_requests,
             memory=memory_requests,
+            gpu=gpu_requests,
             replicas=replicas,
         )
 

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

@@ -204,7 +204,7 @@ class ModelDeploymentSpec:
         job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
         cpu: Optional[str] = None,
         memory: Optional[str] = None,
-        gpu: Optional[Union[str, int]] = None,
+        gpu: Optional[str] = None,
         num_workers: Optional[int] = None,
         max_batch_rows: Optional[int] = None,
         replicas: Optional[int] = None,

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

@@ -438,6 +438,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         partition_column: Optional[sql_identifier.SqlIdentifier],
         statement_params: Optional[dict[str, Any]] = None,
         is_partitioned: bool = True,
+        explain_case_sensitive: bool = False,
     ) -> dataframe.DataFrame:
         with_statements = []
         if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
@@ -505,7 +506,8 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         cols_to_drop = []
 
         for output_name, output_type, output_col_name in returns:
-            output_identifier = sql_identifier.SqlIdentifier(output_name).identifier()
+            case_sensitive = "explain" in method_name.resolved().lower() and explain_case_sensitive
+            output_identifier = sql_identifier.SqlIdentifier(output_name, case_sensitive=case_sensitive).identifier()
             if output_identifier != output_col_name:
                 cols_to_drop.append(output_identifier)
             output_cols.append(F.col(output_identifier).astype(output_type))

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,14 +4,17 @@ 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
 from snowflake.ml.model._model_composer.model_method import (
     constants,
     function_generator,
+    utils,
 )
 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,28 +23,43 @@ 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(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     default_function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+    method_option = options.get("method_options", {}).get(target_method, {})
+    case_sensitive = method_option.get("case_sensitive", False)
     if target_method == "explain":
         default_function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
-    method_option = options.get("method_options", {}).get(target_method, {})
+        case_sensitive = utils.determine_explain_case_sensitive_from_method_options(
+            options.get("method_options", {}), target_method
+        )
     global_function_type = options.get("function_type", default_function_type)
     function_type = method_option.get("function_type", global_function_type)
     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(
-        case_sensitive=method_option.get("case_sensitive", False),
+    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=case_sensitive,
         function_type=function_type,
     )
+    if resolved_volatility:
+        result["volatility"] = resolved_volatility
+
+    return result
 
 
 class ModelMethod:
@@ -94,6 +112,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 +169,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 +179,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/_model_composer/model_method/utils.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional
+
+
+def determine_explain_case_sensitive_from_method_options(
+    method_options: Mapping[str, Optional[Mapping[str, Any]]],
+    target_method: str,
+) -> bool:
+    """Determine explain method case sensitivity from related predict methods.
+
+    Args:
+        method_options: Mapping from method name to its options. Each option may
+            contain ``"case_sensitive"`` to indicate SQL identifier sensitivity.
+        target_method: The target method name being resolved (e.g., an ``explain_*``
+            method).
+
+    Returns:
+        True if the explain method should be treated as case sensitive; otherwise False.
+    """
+    if "explain" not in target_method:
+        return False
+    predict_priority_methods = ["predict_proba", "predict", "predict_log_proba"]
+    for src_method in predict_priority_methods:
+        src_opts = method_options.get(src_method)
+        if src_opts is not None:
+            return bool(src_opts.get("case_sensitive", False))
+    return False

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,
                     )
@@ -234,14 +240,31 @@ class ModelEnv:
                 self._conda_dependencies[channel].remove(spec)
 
     def generate_env_for_cuda(self) -> None:
+
+        # Insert py-xgboost-gpu only for XGBoost versions < 3.0.0
         xgboost_spec = env_utils.find_dep_spec(
-            self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=True
+            self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=False
         )
         if xgboost_spec:
-            self.include_if_absent(
-                [ModelDependency(requirement=f"py-xgboost-gpu{xgboost_spec.specifier}", pip_name="xgboost")],
-                check_local_version=False,
-            )
+            # Only handle explicitly pinned versions. Insert GPU variant iff pinned major < 3.
+            pinned_major: Optional[int] = None
+            for spec in xgboost_spec.specifier:
+                if spec.operator in ("==", "===", ">", ">="):
+                    try:
+                        pinned_major = version.parse(spec.version).major
+                    except version.InvalidVersion:
+                        pinned_major = None
+                    break
+
+            if pinned_major is not None and pinned_major < 3:
+                xgboost_spec = env_utils.find_dep_spec(
+                    self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=True
+                )
+                if xgboost_spec:
+                    self.include_if_absent(
+                        [ModelDependency(requirement=f"py-xgboost-gpu{xgboost_spec.specifier}", pip_name="xgboost")],
+                        check_local_version=False,
+                    )
 
         tf_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="tensorflow", remove_spec=True
@@ -318,13 +341,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 +368,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_meta/model_meta.py

@@ -116,6 +116,8 @@ def create_model_metadata(
     if embed_local_ml_library:
         env.snowpark_ml_version = f"{snowml_version.VERSION}+{file_utils.hash_directory(path_to_copy)}"
 
+    # Persist full method_options
+    method_options: dict[str, dict[str, Any]] = kwargs.pop("method_options", {})
     model_meta = ModelMetadata(
         name=name,
         env=env,
@@ -124,6 +126,7 @@ def create_model_metadata(
         signatures=signatures,
         function_properties=function_properties,
         task=task,
+        method_options=method_options,
     )
 
     code_dir_path = os.path.join(model_dir_path, MODEL_CODE_DIR)
@@ -256,6 +259,7 @@ class ModelMetadata:
         original_metadata_version: Optional[str] = model_meta_schema.MODEL_METADATA_VERSION,
         task: model_types.Task = model_types.Task.UNKNOWN,
         explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = None,
+        method_options: Optional[dict[str, dict[str, Any]]] = None,
     ) -> None:
         self.name = name
         self.signatures: dict[str, model_signature.ModelSignature] = dict()
@@ -283,6 +287,7 @@ class ModelMetadata:
 
         self.task: model_types.Task = task
         self.explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = explain_algorithm
+        self.method_options: dict[str, dict[str, Any]] = method_options or {}
 
     @property
     def min_snowpark_ml_version(self) -> str:
@@ -342,6 +347,7 @@ class ModelMetadata:
                     else None
                 ),
                 "function_properties": self.function_properties,
+                "method_options": self.method_options,
             }
         )
         with open(model_yaml_path, "w", encoding="utf-8") as out:
@@ -381,6 +387,7 @@ class ModelMetadata:
             task=loaded_meta.get("task", model_types.Task.UNKNOWN.value),
             explainability=loaded_meta.get("explainability", None),
             function_properties=loaded_meta.get("function_properties", {}),
+            method_options=loaded_meta.get("method_options", {}),
         )
 
     @classmethod
@@ -436,4 +443,5 @@ class ModelMetadata:
             task=model_types.Task(model_dict.get("task", model_types.Task.UNKNOWN.value)),
             explain_algorithm=explanation_algorithm,
             function_properties=model_dict.get("function_properties", {}),
+            method_options=model_dict.get("method_options", {}),
         )

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

@@ -125,6 +125,7 @@ class ModelMetadataDict(TypedDict):
     task: Required[str]
     explainability: NotRequired[Optional[ExplainabilityMetadataDict]]
     function_properties: NotRequired[dict[str, dict[str, Any]]]
+    method_options: NotRequired[dict[str, dict[str, Any]]]
 
 
 class ModelExplainAlgorithm(Enum):

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

@@ -21,14 +21,14 @@ 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-connector-python>=3.17.0,<4",
     "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0",
     "snowflake.core>=1.0.2,<2",
     "sqlparse>=0.4,<1",
     "tqdm<5",
     "typing-extensions>=4.1.0,<5",
-    "xgboost>=1.7.3,<3",
+    "xgboost<4",
 ]

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/_internal/snowpark_implementations/xgboost_external_memory_trainer.py

@@ -93,7 +93,7 @@ def get_data_iterator(
             cache_dir_name = tempfile.mkdtemp()
             super().__init__(cache_prefix=os.path.join(cache_dir_name, "cache"))
 
-        def next(self, batch_consumer_fn) -> int:  # type: ignore[no-untyped-def]
+        def next(self, batch_consumer_fn) -> bool | int:  # type: ignore[no-untyped-def]
             """Advance the iterator by 1 step and pass the data to XGBoost's batch_consumer_fn.
             This function is called by XGBoost during the construction of ``DMatrix``
 
@@ -101,7 +101,7 @@ def get_data_iterator(
                 batch_consumer_fn: batch consumer function
 
             Returns:
-                0 if there is no more data, else 1.
+                False/0 if there is no more data, else True/1.
             """
             while (self._df is None) or (self._df.shape[0] < self._batch_size):
                 # Read files and append data to temp df until batch size is reached.
@@ -117,7 +117,7 @@ def get_data_iterator(
 
             if (self._df is None) or (self._df.shape[0] == 0):
                 # No more data
-                return 0
+                return False
 
             # Slice the temp df and save the remainder in the temp df
             batch_end_index = min(self._batch_size, self._df.shape[0])
@@ -133,8 +133,8 @@ def get_data_iterator(
                 func_args["weight"] = batch_df[self._sample_weight_col].squeeze()
 
             batch_consumer_fn(**func_args)
-            # Return 1 to let XGBoost know we haven't seen all the files yet.
-            return 1
+            # Return True to let XGBoost know we haven't seen all the files yet.
+            return True
 
         def reset(self) -> None:
             """Reset the iterator to its beginning"""

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(