snowflake-ml-python 1.15.0__tar.gz → 1.16.0__tar.gz

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/CHANGELOG.md

@@ -1,6 +1,44 @@
 # Release History
 
-## 1.15.0
+## 1.16.0
+
+### Bug Fixes
+
+* Registry: Remove redundant pip dependency warnings when `artifact_repository_map` is provided for warehouse model deployments.
+
+### Behavior Changes
+
+### New Features
+
+* Support scikit-learn < 1.8.
+* ML Job: Added support for configuring the runtime image via `runtime_environment`
+  (image tag or full image URL) at submission time.
+  Examples:
+  * @remote(compute_pool, stage_name = 'payload_stage', runtime_environment = '1.8.0')
+  * submit_file('/path/to/repo/test.py', compute_pool, stage_name = 'payload_stage', runtime_environment = '/mydb/myschema/myrepo/myimage:latest')
+* Registry: Ability to mark model methods as `Volatility.VOLATILE` or `Volatility.IMMUTABLE`.
+
+```python
+from snowflake.ml.model.volatility import Volatility
+
+options = {
+    "embed_local_ml_library": True,
+    "relax_version": True,
+    "save_location": "/path/to/my/directory",
+    "function_type": "TABLE_FUNCTION",
+    "volatility": Volatility.IMMUTABLE,
+    "method_options": {
+        "predict": {
+            "case_sensitive": False,
+            "max_batch_size": 100,
+            "function_type": "TABLE_FUNCTION",
+            "volatility": Volatility.VOLATILE,
+        },
+}
+
+````
+
+## 1.15.0 (09-29-2025)
 
 ### Bug Fixes
 

{snowflake_ml_python-1.15.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.15.0
+Version: 1.16.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -250,7 +250,7 @@ Requires-Dist: pytimeparse<2,>=1.1.8
 Requires-Dist: pyyaml<7,>=6.0
 Requires-Dist: retrying<2,>=1.3.3
 Requires-Dist: s3fs<2026,>=2024.6.1
-Requires-Dist: scikit-learn<1.7
+Requires-Dist: scikit-learn<1.8
 Requires-Dist: scipy<2,>=1.9
 Requires-Dist: shap<1,>=0.46.0
 Requires-Dist: snowflake-connector-python[pandas]<4,>=3.16.0
@@ -406,7 +406,45 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
-## 1.15.0
+## 1.16.0
+
+### Bug Fixes
+
+* Registry: Remove redundant pip dependency warnings when `artifact_repository_map` is provided for warehouse model deployments.
+
+### Behavior Changes
+
+### New Features
+
+* Support scikit-learn < 1.8.
+* ML Job: Added support for configuring the runtime image via `runtime_environment`
+  (image tag or full image URL) at submission time.
+  Examples:
+  * @remote(compute_pool, stage_name = 'payload_stage', runtime_environment = '1.8.0')
+  * submit_file('/path/to/repo/test.py', compute_pool, stage_name = 'payload_stage', runtime_environment = '/mydb/myschema/myrepo/myimage:latest')
+* Registry: Ability to mark model methods as `Volatility.VOLATILE` or `Volatility.IMMUTABLE`.
+
+```python
+from snowflake.ml.model.volatility import Volatility
+
+options = {
+    "embed_local_ml_library": True,
+    "relax_version": True,
+    "save_location": "/path/to/my/directory",
+    "function_type": "TABLE_FUNCTION",
+    "volatility": Volatility.IMMUTABLE,
+    "method_options": {
+        "predict": {
+            "case_sensitive": False,
+            "max_batch_size": 100,
+            "function_type": "TABLE_FUNCTION",
+            "volatility": Volatility.VOLATILE,
+        },
+}
+
+````
+
+## 1.15.0 (09-29-2025)
 
 ### Bug Fixes
 

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/pyproject.toml

@@ -8,7 +8,7 @@ description = "The machine learning client library that is used for interacting
 classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Environment :: Other Environment", "Intended Audience :: Developers", "Intended Audience :: Education", "Intended Audience :: Information Technology", "Intended Audience :: System Administrators", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Database", "Topic :: Software Development", "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Application Frameworks", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Information Analysis",]
 requires-python = ">=3.9, <3.13"
 dynamic = [ "version", "readme",]
-dependencies = [ "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow<19.0.0", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.7", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.16.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",]
+dependencies = [ "anyio>=3.5.0,<5", "cachetools>=3.1.1,<6", "cloudpickle>=2.0.0", "cryptography", "fsspec[http]>=2024.6.1,<2026", "importlib_resources>=6.1.1, <7", "numpy>=1.23,<3", "packaging>=20.9,<25", "pandas>=2.1.4,<3", "platformdirs<5", "pyarrow<19.0.0", "pydantic>=2.8.2, <3", "pyjwt>=2.0.0, <3", "pytimeparse>=1.1.8,<2", "pyyaml>=6.0,<7", "retrying>=1.3.3,<2", "s3fs>=2024.6.1,<2026", "scikit-learn<1.8", "scipy>=1.9,<2", "shap>=0.46.0,<1", "snowflake-connector-python[pandas]>=3.16.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",]
 [[project.authors]]
 name = "Snowflake, Inc"
 email = "support@snowflake.com"

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/_internal/platform_capabilities.py

@@ -17,6 +17,7 @@ logger = logging.getLogger(__name__)
 
 LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
 INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC_FROM_CLIENT_VERSION"
+SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST = "SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST"
 
 
 class PlatformCapabilities:
@@ -73,6 +74,9 @@ class PlatformCapabilities:
     def is_inlined_deployment_spec_enabled(self) -> bool:
         return self._is_version_feature_enabled(INLINE_DEPLOYMENT_SPEC_PARAMETER)
 
+    def is_set_module_functions_volatility_from_manifest(self) -> bool:
+        return self._get_bool_feature(SET_MODULE_FUNCTIONS_VOLATILITY_FROM_MANIFEST, False)
+
     def is_live_commit_enabled(self) -> bool:
         return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
 

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/_internal/utils/mixins.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import Any, Optional
 
 from snowflake.ml._internal.utils import identifier
@@ -16,6 +17,14 @@ def _identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
     return saved_resolved == current_resolved
 
 
+@dataclass(frozen=True)
+class _SessionState:
+    account: Optional[str]
+    role: Optional[str]
+    database: Optional[str]
+    schema: Optional[str]
+
+
 class SerializableSessionMixin:
     """Mixin that provides pickling capabilities for objects with Snowpark sessions."""
 
@@ -40,17 +49,23 @@ class SerializableSessionMixin:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         """Restore session from context during unpickling."""
-        saved_account = state.pop(_SESSION_ACCOUNT_KEY, None)
-        saved_role = state.pop(_SESSION_ROLE_KEY, None)
-        saved_database = state.pop(_SESSION_DATABASE_KEY, None)
-        saved_schema = state.pop(_SESSION_SCHEMA_KEY, None)
+        session_state = _SessionState(
+            account=state.pop(_SESSION_ACCOUNT_KEY, None),
+            role=state.pop(_SESSION_ROLE_KEY, None),
+            database=state.pop(_SESSION_DATABASE_KEY, None),
+            schema=state.pop(_SESSION_SCHEMA_KEY, None),
+        )
 
         if hasattr(super(), "__setstate__"):
             super().__setstate__(state)  # type: ignore[misc]
         else:
             self.__dict__.update(state)
 
-        if saved_account is not None:
+        self._set_session(session_state)
+
+    def _set_session(self, session_state: _SessionState) -> None:
+
+        if session_state.account is not None:
             active_sessions = snowpark_session._get_active_sessions()
             if len(active_sessions) == 0:
                 raise RuntimeError("No active Snowpark session available. Please create a session.")
@@ -63,10 +78,10 @@ class SerializableSessionMixin:
                     active_sessions,
                     key=lambda s: sum(
                         (
-                            _identifiers_match(saved_account, s.get_current_account()),
-                            _identifiers_match(saved_role, s.get_current_role()),
-                            _identifiers_match(saved_database, s.get_current_database()),
-                            _identifiers_match(saved_schema, s.get_current_schema()),
+                            _identifiers_match(session_state.account, s.get_current_account()),
+                            _identifiers_match(session_state.role, s.get_current_role()),
+                            _identifiers_match(session_state.database, s.get_current_database()),
+                            _identifiers_match(session_state.schema, s.get_current_schema()),
                         )
                     ),
                 ),

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/experiment/experiment_tracking.py

@@ -1,10 +1,10 @@
 import functools
 import json
 import sys
-from typing import Any, Optional, Union
+from typing import Any, Callable, Concatenate, Optional, ParamSpec, TypeVar, Union
 from urllib.parse import quote
 
-import snowflake.snowpark._internal.utils as snowpark_utils
+from snowflake import snowpark
 from snowflake.ml import model as ml_model, registry
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import mixins, sql_identifier
@@ -18,20 +18,40 @@ from snowflake.ml.experiment._client import (
 )
 from snowflake.ml.model import type_hints
 from snowflake.ml.utils import sql_client as sql_client_utils
-from snowflake.snowpark import session
 
 DEFAULT_EXPERIMENT_NAME = sql_identifier.SqlIdentifier("DEFAULT")
 
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _restore_session(
+    func: Callable[Concatenate["ExperimentTracking", P], T],
+) -> Callable[Concatenate["ExperimentTracking", P], T]:
+    @functools.wraps(func)
+    def wrapper(self: "ExperimentTracking", /, *args: P.args, **kwargs: P.kwargs) -> T:
+        if self._session is None:
+            if self._session_state is None:
+                raise RuntimeError(
+                    f"Session is not set before calling {func.__name__}, and there is no session state to restore from"
+                )
+            self._set_session(self._session_state)
+            if self._session is None:
+                raise RuntimeError(f"Failed to restore session before calling {func.__name__}")
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
 
 class ExperimentTracking(mixins.SerializableSessionMixin):
     """
     Class to manage experiments in Snowflake.
     """
 
-    @snowpark_utils.private_preview(version="1.9.1")
+    @snowpark._internal.utils.private_preview(version="1.9.1")
     def __init__(
         self,
-        session: session.Session,
+        session: snowpark.Session,
         *,
         database_name: Optional[str] = None,
         schema_name: Optional[str] = None,
@@ -73,7 +93,10 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             database_name=self._database_name,
             schema_name=self._schema_name,
         )
-        self._session = session
+        self._session: Optional[snowpark.Session] = session
+        # Used to store information about the session if the session could not be restored during unpickling
+        # _session_state is None if and only if _session is not None
+        self._session_state: Optional[mixins._SessionState] = None
 
         # The experiment in context
         self._experiment: Optional[entities.Experiment] = None
@@ -87,20 +110,29 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         state["_registry"] = None
         return state
 
-    def __setstate__(self, state: dict[str, Any]) -> None:
-        super().__setstate__(state)
-        # Restore unpicklable attributes
-        self._sql_client = sql_client.ExperimentTrackingSQLClient(
-            session=self._session,
-            database_name=self._database_name,
-            schema_name=self._schema_name,
-        )
-        self._registry = registry.Registry(
-            session=self._session,
-            database_name=self._database_name,
-            schema_name=self._schema_name,
-        )
+    def _set_session(self, session_state: mixins._SessionState) -> None:
+        try:
+            super()._set_session(session_state)
+            assert self._session is not None
+        except (snowpark.exceptions.SnowparkSessionException, AssertionError):
+            # If session was not set, store the session state
+            self._session = None
+            self._session_state = session_state
+        else:
+            # If session was set, clear the session state, and reinitialize the SQL client and registry
+            self._session_state = None
+            self._sql_client = sql_client.ExperimentTrackingSQLClient(
+                session=self._session,
+                database_name=self._database_name,
+                schema_name=self._schema_name,
+            )
+            self._registry = registry.Registry(
+                session=self._session,
+                database_name=self._database_name,
+                schema_name=self._schema_name,
+            )
 
+    @_restore_session
     def set_experiment(
         self,
         experiment_name: str,
@@ -125,6 +157,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run = None
         return self._experiment
 
+    @_restore_session
     def delete_experiment(
         self,
         experiment_name: str,
@@ -141,8 +174,10 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             self._run = None
 
     @functools.wraps(registry.Registry.log_model)
+    @_restore_session
     def log_model(
         self,
+        /,  # self needs to be a positional argument to stop mypy from complaining
         model: Union[type_hints.SupportedModelType, ml_model.ModelVersion],
         *,
         model_name: str,
@@ -152,6 +187,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         with experiment_info.ExperimentInfoPatcher(experiment_info=run._get_experiment_info()):
             return self._registry.log_model(model, model_name=model_name, **kwargs)
 
+    @_restore_session
     def start_run(
         self,
         run_name: Optional[str] = None,
@@ -181,6 +217,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         self._run = entities.Run(experiment_tracking=self, experiment_name=experiment.name, run_name=run_name)
         return self._run
 
+    @_restore_session
     def end_run(self, run_name: Optional[str] = None) -> None:
         """
         End the current run if no run name is provided. Otherwise, the specified run is ended.
@@ -210,6 +247,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             self._run = None
         self._print_urls(experiment_name=experiment_name, run_name=run_name)
 
+    @_restore_session
     def delete_run(
         self,
         run_name: str,
@@ -248,6 +286,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_metrics(metrics={key: value}, step=step)
 
+    @_restore_session
     def log_metrics(
         self,
         metrics: dict[str, float],
@@ -284,6 +323,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
         """
         self.log_params({key: value})
 
+    @_restore_session
     def log_params(
         self,
         params: dict[str, Any],
@@ -305,6 +345,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             params=json.dumps([param.to_dict() for param in params_list]),
         )
 
+    @_restore_session
     def log_artifact(
         self,
         local_path: str,
@@ -328,6 +369,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
                 file_path=file_path,
             )
 
+    @_restore_session
     def list_artifacts(
         self,
         run_name: str,
@@ -356,6 +398,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             artifact_path=artifact_path or "",
         )
 
+    @_restore_session
     def download_artifacts(
         self,
         run_name: str,
@@ -397,6 +440,7 @@ class ExperimentTracking(mixins.SerializableSessionMixin):
             return self._run
         return self.start_run()
 
+    @_restore_session
     def _generate_run_name(self, experiment: entities.Experiment) -> sql_identifier.SqlIdentifier:
         generator = hrid_generator.HRID16()
         existing_runs = self._sql_client.show_runs_in_experiment(experiment_name=experiment.name)

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/jobs/_utils/spec_utils.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import re
 import sys
 from math import ceil
 from pathlib import PurePath
@@ -10,6 +11,8 @@ from snowflake.ml._internal.utils import snowflake_env
 from snowflake.ml.jobs._utils import constants, feature_flags, query_helper, types
 from snowflake.ml.jobs._utils.runtime_env_utils import RuntimeEnvironmentsDict
 
+_OCI_TAG_REGEX = re.compile("^[a-zA-Z0-9._-]{1,128}$")
+
 
 def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.ComputeResources:
     """Extract resource information for the specified compute pool"""
@@ -56,22 +59,55 @@ def _get_runtime_image(session: snowpark.Session, target_hardware: Literal["CPU"
     return selected_runtime.runtime_container_image if selected_runtime else None
 
 
-def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
+def _check_image_tag_valid(tag: Optional[str]) -> bool:
+    if tag is None:
+        return False
+
+    return _OCI_TAG_REGEX.fullmatch(tag) is not None
+
+
+def _get_image_spec(
+    session: snowpark.Session, compute_pool: str, runtime_environment: Optional[str] = None
+) -> types.ImageSpec:
+    """
+    Resolve image specification (container image and resources) for the job.
+
+    Behavior:
+    - If `runtime_environment` is empty or the feature flag is disabled, use the
+      default image tag and image name.
+    - If `runtime_environment` is a valid image tag, use that tag with the default
+      repository/name.
+    - If `runtime_environment` is a full image URL, use it directly.
+    - If the feature flag is enabled and `runtime_environment` is not provided,
+      select an ML Runtime image matching the local Python major.minor
+    - When multiple inputs are provided, `runtime_environment` takes priority.
+
+    Args:
+        session: Snowflake session.
+        compute_pool: Compute pool used to infer CPU/GPU resources.
+        runtime_environment: Optional image tag or full image URL to override.
+
+    Returns:
+        Image spec including container image and resource requests/limits.
+    """
     # Retrieve compute pool node resources
     resources = _get_node_resources(session, compute_pool=compute_pool)
+    hardware = "GPU" if resources.gpu > 0 else "CPU"
+    image_tag = _get_runtime_image_tag()
+    image_repo = constants.DEFAULT_IMAGE_REPO
+    image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
 
     # Use MLRuntime image
-    hardware = "GPU" if resources.gpu > 0 else "CPU"
     container_image = None
-    if feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
+    if runtime_environment:
+        if _check_image_tag_valid(runtime_environment):
+            image_tag = runtime_environment
+        else:
+            container_image = runtime_environment
+    elif feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
         container_image = _get_runtime_image(session, hardware)  # type: ignore[arg-type]
 
-    if not container_image:
-        image_repo = constants.DEFAULT_IMAGE_REPO
-        image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-        image_tag = _get_runtime_image_tag()
-        container_image = f"{image_repo}/{image_name}:{image_tag}"
-
+    container_image = container_image or f"{image_repo}/{image_name}:{image_tag}"
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
         resource_requests=resources,
@@ -127,6 +163,7 @@ def generate_service_spec(
     target_instances: int = 1,
     min_instances: int = 1,
     enable_metrics: bool = False,
+    runtime_environment: Optional[str] = None,
 ) -> dict[str, Any]:
     """
     Generate a service specification for a job.
@@ -139,11 +176,12 @@ def generate_service_spec(
         target_instances: Number of instances for multi-node job
         enable_metrics: Enable platform metrics for the job
         min_instances: Minimum number of instances required to start the job
+        runtime_environment: The runtime image to use. Only support image tag or full image URL.
 
     Returns:
         Job service specification
     """
-    image_spec = _get_image_spec(session, compute_pool)
+    image_spec = _get_image_spec(session, compute_pool, runtime_environment)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
     resource_requests: dict[str, Union[str, int]] = {
@@ -317,7 +355,7 @@ def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
     Returns:
         The patched object.
     """
-    if not type(base) is type(patch):
+    if type(base) is not type(patch):
         if base is not None:
             logging.warning(f"Type mismatch while merging {display_name} (base={type(base)}, patch={type(patch)})")
         return patch

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/jobs/manager.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import pathlib
+import sys
 import textwrap
 from pathlib import PurePath
 from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
@@ -344,6 +345,9 @@ def submit_from_stage(
             query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
             spec_overrides (dict): A dictionary of overrides for the service spec.
             imports (list[Union[tuple[str, str], tuple[str]]]): A list of additional payloads used in the job.
+            runtime_environment (str): The runtime image to use. Only support image tag or full image URL,
+                e.g. "1.7.1" or "image_repo/image_name:image_tag". When it refers to a full image URL,
+                it should contain image repository, image name and image tag.
 
     Returns:
         An object representing the submitted job.
@@ -409,6 +413,7 @@ def _submit_job(
         "min_instances",
         "enable_metrics",
         "query_warehouse",
+        "runtime_environment",
     ],
 )
 def _submit_job(
@@ -459,6 +464,9 @@ def _submit_job(
         )
         imports = kwargs.pop("additional_payloads")
 
+    if "runtime_environment" in kwargs:
+        logger.warning("'runtime_environment' is in private preview since 1.15.0, do not use it in production.")
+
     # Use kwargs for less common optional parameters
     database = kwargs.pop("database", None)
     schema = kwargs.pop("schema", None)
@@ -470,6 +478,7 @@ def _submit_job(
     enable_metrics = kwargs.pop("enable_metrics", True)
     query_warehouse = kwargs.pop("query_warehouse", session.get_current_warehouse())
     imports = kwargs.pop("imports", None) or imports
+    runtime_environment = kwargs.pop("runtime_environment", None)
 
     # Warn if there are unknown kwargs
     if kwargs:
@@ -544,6 +553,7 @@ def _submit_job(
             min_instances=min_instances,
             enable_metrics=enable_metrics,
             use_async=True,
+            runtime_environment=runtime_environment,
         )
 
     # Fall back to v1
@@ -556,6 +566,7 @@ def _submit_job(
         target_instances=target_instances,
         min_instances=min_instances,
         enable_metrics=enable_metrics,
+        runtime_environment=runtime_environment,
     )
 
     # Generate spec overrides
@@ -639,6 +650,7 @@ def _do_submit_job_v2(
     min_instances: int = 1,
     enable_metrics: bool = True,
     use_async: bool = True,
+    runtime_environment: Optional[str] = None,
 ) -> jb.MLJob[Any]:
     """
     Generate the SQL query for job submission.
@@ -657,6 +669,7 @@ def _do_submit_job_v2(
         min_instances: Minimum number of instances required to start the job.
         enable_metrics: Whether to enable platform metrics for the job.
         use_async: Whether to run the job asynchronously.
+        runtime_environment: image tag or full image URL to use for the job.
 
     Returns:
         The job object.
@@ -672,6 +685,13 @@ def _do_submit_job_v2(
         "ENABLE_METRICS": enable_metrics,
         "SPEC_OVERRIDES": spec_overrides,
     }
+    # for the image tag or full image URL, we use that directly
+    if runtime_environment:
+        spec_options["RUNTIME"] = runtime_environment
+    elif feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
+        # when feature flag is enabled, we get the local python version and wrap it in a dict
+        # in system function, we can know whether it is python version or image tag or full image URL through the format
+        spec_options["RUNTIME"] = json.dumps({"pythonVersion": f"{sys.version_info.major}.{sys.version_info.minor}"})
     job_options = {
         "EXTERNAL_ACCESS_INTEGRATIONS": external_access_integrations,
         "QUERY_WAREHOUSE": query_warehouse,

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/model/__init__.py

@@ -1,6 +1,20 @@
-from snowflake.ml.model._client.model.batch_inference_specs import JobSpec, OutputSpec
+from snowflake.ml.model._client.model.batch_inference_specs import (
+    JobSpec,
+    OutputSpec,
+    SaveMode,
+)
 from snowflake.ml.model._client.model.model_impl import Model
 from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
+from snowflake.ml.model.volatility import Volatility
 
-__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "JobSpec", "OutputSpec"]
+__all__ = [
+    "Model",
+    "ModelVersion",
+    "ExportMode",
+    "HuggingFacePipelineModel",
+    "JobSpec",
+    "OutputSpec",
+    "SaveMode",
+    "Volatility",
+]

{snowflake_ml_python-1.15.0 → snowflake_ml_python-1.16.0}/snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -1,10 +1,26 @@
-from typing import Optional, Union
+from enum import Enum
+from typing import Optional
 
 from pydantic import BaseModel
 
 
+class SaveMode(str, Enum):
+    """Save mode options for batch inference output.
+
+    Determines the behavior when files already exist in the output location.
+
+    OVERWRITE: Remove existing files and write new results.
+
+    ERROR: Raise an error if files already exist in the output location.
+    """
+
+    OVERWRITE = "overwrite"
+    ERROR = "error"
+
+
 class OutputSpec(BaseModel):
     stage_location: str
+    mode: SaveMode = SaveMode.ERROR
 
 
 class JobSpec(BaseModel):
@@ -12,10 +28,10 @@ class JobSpec(BaseModel):
     job_name: Optional[str] = None
     num_workers: Optional[int] = None
     function_name: Optional[str] = None
-    gpu: Optional[Union[str, int]] = None
     force_rebuild: bool = False
     max_batch_rows: int = 1024
     warehouse: Optional[str] = None
     cpu_requests: Optional[str] = None
     memory_requests: Optional[str] = None
+    gpu_requests: Optional[str] = None
     replicas: Optional[int] = None