snowflake-ml-python 1.50.0__tar.gz → 1.51.0__tar.gz

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/CHANGELOG.md

@@ -1,6 +1,24 @@
 # Release History
 
-## 1.50.0
+## 1.51.0
+
+### New Features
+
+* ML Jobs: `artifact_repositories` now accepts a list of repositories, allowing multiple artifact
+  repositories to be specified for a job.
+* Registry: Added support for logging and running MLflow 3.x models.
+
+### Bug Fixes
+
+* Registry: Fixed a bug where an MLflow model logged with
+  `target_platforms=["SNOWPARK_CONTAINER_SERVICES"]` could drop its pip dependencies, leaving the
+  deployed model missing required packages.
+
+### Behavior Changes
+
+### Deprecations
+
+## 1.50.0 (2026-08-03)
 
 ### New Features
 

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.50.0
+Version: 1.51.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: 
@@ -264,7 +264,7 @@ Requires-Dist: altair<7,>=5; extra == "all"
 Requires-Dist: catboost<2,>=1.2.0; extra == "all"
 Requires-Dist: keras<4,>=2.0.0; extra == "all"
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "all"
-Requires-Dist: mlflow<3,>=2.16.0; extra == "all"
+Requires-Dist: mlflow<4,>=2.16.0; extra == "all"
 Requires-Dist: prophet<2,>=1.1.0; extra == "all"
 Requires-Dist: sentence-transformers<6,>=2.7.0; extra == "all"
 Requires-Dist: sentencepiece<0.3,>=0.1.95; extra == "all"
@@ -285,7 +285,7 @@ Requires-Dist: torch<3,>=2.0.1; extra == "keras"
 Provides-Extra: lightgbm
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "lightgbm"
 Provides-Extra: llm
-Requires-Dist: mlflow<3,>=2.16.0; extra == "llm"
+Requires-Dist: mlflow<4,>=2.16.0; extra == "llm"
 Requires-Dist: sentence-transformers<6,>=2.7.0; extra == "llm"
 Requires-Dist: sentencepiece<0.3,>=0.1.95; extra == "llm"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "llm"
@@ -293,7 +293,7 @@ Requires-Dist: torch<3,>=2.0.1; extra == "llm"
 Requires-Dist: torchdata<1,>=0.4; extra == "llm"
 Requires-Dist: transformers!=4.51.3,<6,>=4.39.3; extra == "llm"
 Provides-Extra: mlflow
-Requires-Dist: mlflow<3,>=2.16.0; extra == "mlflow"
+Requires-Dist: mlflow<4,>=2.16.0; extra == "mlflow"
 Provides-Extra: prophet
 Requires-Dist: prophet<2,>=1.1.0; extra == "prophet"
 Provides-Extra: streamlit
@@ -416,7 +416,25 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
-## 1.50.0
+## 1.51.0
+
+### New Features
+
+* ML Jobs: `artifact_repositories` now accepts a list of repositories, allowing multiple artifact
+  repositories to be specified for a job.
+* Registry: Added support for logging and running MLflow 3.x models.
+
+### Bug Fixes
+
+* Registry: Fixed a bug where an MLflow model logged with
+  `target_platforms=["SNOWPARK_CONTAINER_SERVICES"]` could drop its pip dependencies, leaving the
+  deployed model missing required packages.
+
+### Behavior Changes
+
+### Deprecations
+
+## 1.50.0 (2026-08-03)
 
 ### New Features
 

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/pyproject.toml

@@ -24,13 +24,13 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "altair>=5,<7", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "prophet>=1.1.0, <2", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.3", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
+all = [ "altair>=5,<7", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <4", "prophet>=1.1.0, <2", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.3", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
 altair = [ "altair>=5,<7",]
 catboost = [ "catboost>=1.2.0, <2",]
 keras = [ "keras>=2.0.0,<4", "tensorflow>=2.17.0,<3", "torch>=2.0.1,<3",]
 lightgbm = [ "lightgbm>=4.1.0, <5",]
-llm = [ "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
-mlflow = [ "mlflow>=2.16.0, <3",]
+llm = [ "mlflow>=2.16.0, <4", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
+mlflow = [ "mlflow>=2.16.0, <4",]
 prophet = [ "prophet>=1.1.0, <2",]
 streamlit = [ "streamlit>=1.30.0,<2",]
 tensorflow = [ "tensorflow>=2.17.0,<3",]

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/feature_store/feature_store.py

@@ -49,6 +49,7 @@ from snowflake.ml.feature_store import (
     feature_group as fg_mod,
     feature_view_append_only_validation,
     feature_view_refresh_freq,
+    interval_utils,
     online_service,
     online_service_http_client,
     realtime_dataset as rtfv_dataset,
@@ -449,6 +450,59 @@ class _FeatureStoreConfig:
         return f"{self.database}.{self.schema}"
 
 
+def _canonicalize_operational_fields(payload: dict[str, Any]) -> dict[str, Any]:
+    """Return a canonicalized copy of an FV operational-field payload (Plan section A4).
+
+    Closes Bug 2 (``UPDATE_FV`` operational drift on tiled online BFVs) by
+    coercing two fields to a single representation that the diff logic in
+    ``decl/invariants.py`` and ``decl/planner.py`` consumes symmetrically:
+
+    * ``online_store_type`` — lowercased and whitespace-stripped.  Snowflake
+      ``DESCRIBE … TYPE = SPECIFICATION`` returns mixed-case enum values
+      with surrounding whitespace; the local compile emits the lowercase
+      form.  Without canonicalization the diff observes cosmetic drift on
+      a clean re-plan.
+    * ``target_lag_sec`` — always an ``int`` computed from the deployed
+      string form (``payload["target_lag"]``) via
+      :func:`interval_utils.interval_to_seconds`.  The OFT may return
+      either ``"5 minutes"`` or ``"300 SECONDS"`` for the same value;
+      diffing on seconds (not strings) makes the comparison stable.
+
+    The helper is pure: the input ``payload`` is not mutated.  Absent or
+    ``None`` keys are preserved as-is so callers can chain this through a
+    pipeline of canonicalizers without re-introducing fields.
+
+    The ``"DOWNSTREAM"`` sentinel (CRON-task target_lag) is silently
+    skipped: the planner recovers the cron expression via the companion
+    Task, and ``interval_to_seconds`` would otherwise crash trying to
+    parse it.
+
+    Args:
+        payload: Operational-field dictionary, typically the
+            ``spec_payload["spec"]`` shape that ``decl/`` consumes.
+
+    Returns:
+        A new dict (caller's payload is not mutated) with
+        ``online_store_type`` and ``target_lag_sec`` in canonical form.
+    """
+    out = dict(payload)
+
+    if "online_store_type" in out:
+        value = out["online_store_type"]
+        if isinstance(value, str):
+            out["online_store_type"] = value.strip().lower()
+
+    if "target_lag" in out and out["target_lag"] is not None:
+        target_lag_str = str(out["target_lag"])
+        if target_lag_str.strip().upper() != "DOWNSTREAM":
+            try:
+                out["target_lag_sec"] = int(interval_utils.interval_to_seconds(target_lag_str))
+            except (KeyError, ValueError, TypeError):
+                pass
+
+    return out
+
+
 def switch_warehouse(
     f: Callable[Concatenate[FeatureStore, _Args], _RT],
 ) -> Callable[Concatenate[FeatureStore, _Args], _RT]:
@@ -862,7 +916,12 @@ class FeatureStore:
 
         return self.get_entity(entity.name)
 
-    def update_entity(self, name: str, *, desc: Optional[str] = None) -> Optional[Entity]:
+    def update_entity(
+        self,
+        name: str,
+        *,
+        desc: Optional[str] = None,
+    ) -> Optional[Entity]:
         """Update a registered entity with provided information.
 
         Args:
@@ -910,9 +969,9 @@ class FeatureStore:
             return None
 
         new_desc = desc if desc is not None else found_rows[0]["DESC"]
+        full_name = f"{self._config.full_schema_path}.{self._get_entity_name(name)}"
 
         try:
-            full_name = f"{self._config.full_schema_path}.{self._get_entity_name(name)}"
             self._session.sql(f"ALTER TAG {full_name} SET COMMENT = {_sql_string_literal(new_desc)}").collect(
                 statement_params=self._telemetry_stmp
             )
@@ -7921,12 +7980,6 @@ FROM {batch_cte_names[0]}{''.join(merge_join_parts)}
         source_refs_meta = self._metadata_manager.get_feature_view_source_refs(name, version)
         source_refs_json = json.dumps(source_refs_meta.sources) if source_refs_meta is not None else None
 
-        # Surface the authored source-ref list (JSON-encoded) when a
-        # ``FV_SOURCE_REFS`` row exists; otherwise leave the column null.
-        # Verbose-only: goes into output_values_extra alongside backup_source.
-        source_refs_meta = self._metadata_manager.get_feature_view_source_refs(name, version)
-        source_refs_json = json.dumps(source_refs_meta.sources) if source_refs_meta is not None else None
-
         backup_source: Optional[str] = None
         if fv_metadata.is_append_only:
             append_only_meta = self._metadata_manager.get_append_only_metadata(name, version)
@@ -8110,11 +8163,6 @@ FROM {batch_cte_names[0]}{''.join(merge_join_parts)}
         source_refs_meta = self._metadata_manager.get_feature_view_source_refs(name.identifier(), version)
         restored_source_refs = source_refs_meta.sources if source_refs_meta is not None else None
 
-        # Rehydrate the authored source-ref list onto the reconstructed
-        # FV when a ``FV_SOURCE_REFS`` row was persisted at register time.
-        source_refs_meta = self._metadata_manager.get_feature_view_source_refs(name.identifier(), version)
-        restored_source_refs = source_refs_meta.sources if source_refs_meta is not None else None
-
         if obj_type == _FeatureStoreObjTypes.MANAGED_FEATURE_VIEW:
             df = self._session.sql(query)
             entities = [find_and_compose_entity(n) for n in fv_metadata.entities]
@@ -8511,6 +8559,11 @@ FROM {batch_cte_names[0]}{''.join(merge_join_parts)}
         else:
             source_clause = f"FROM {source_table_name}"
 
+        # A view-backed source only lets the OFT refresh incrementally when it has change tracking
+        # enabled; otherwise the backend silently downgrades the OFT to a FULL refresh. Enable it
+        # before issuing the CREATE so the refresh mode is derived correctly.
+        self._enable_source_view_change_tracking_for_incremental_oft(feature_view, source_table_name)
+
         # Create online feature table
         try:
             query = fv_mod.build_oft_create_sql(
@@ -8536,6 +8589,88 @@ FROM {batch_cte_names[0]}{''.join(merge_join_parts)}
 
         return online_table_name
 
+    def _enable_source_view_change_tracking_for_incremental_oft(
+        self,
+        feature_view: FeatureView,
+        fully_qualified_source_name: str,
+    ) -> None:
+        """Enable change tracking on a POSTGRES-backed OFT's source view so it can refresh incrementally.
+
+        The backend derives a spec-backed (POSTGRES) Online Feature Table's refresh mode from its
+        source object. When that source is a plain View (static/external feature views), change
+        tracking is off by default and the OFT silently falls back to a FULL refresh — an explicit
+        ``INCREMENTAL`` request even errors. Enabling change tracking on the view lets the OFT
+        resolve to an incremental refresh.
+
+        This targets the object the OFT actually reads from at refresh time, and is a no-op
+        otherwise:
+
+        - Managed feature views are backed by a Dynamic Table (change tracking is inherent to an
+          incremental DT and cannot be enabled on a FULL one), so their OFT never reads from a
+          plain view.
+        - Streaming feature views may be backed offline by a view (a non-aggregated SFV with
+          ``refresh_freq is None`` materializes as a View), but their OFT is spec-backed by the
+          UDF-transformed table / stream source (``FROM SPECIFICATION``) and refreshes
+          continuously — it does not read from that offline view, so enabling change tracking on
+          it would not affect the OFT.
+        - Realtime feature views have no offline view at all.
+        - HYBRID_TABLE-backed OFTs are out of scope. This targets the spec-backed refresh-mode
+          derivation used by POSTGRES OFTs, and enabling change tracking has real side effects
+          (it cascades to the view's base tables and requires ``OWNERSHIP`` on them), so we do not
+          newly impose it on the legacy HYBRID_TABLE path — its behavior is left unchanged.
+
+        It is also skipped when the refresh mode is not one that benefits from change tracking
+        (i.e., only AUTO, INCREMENTAL, and the default None proceed; FULL and unknown/future
+        modes are conservatively skipped).
+
+        Args:
+            feature_view: The FeatureView backing the online feature table.
+            fully_qualified_source_name: ``DB.SCHEMA.SOURCE`` object the OFT reads from.
+
+        Raises:
+            SnowflakeMLException: If enabling change tracking fails and the refresh mode is
+                explicitly ``INCREMENTAL``. For ``AUTO`` or unset modes the failure is logged
+                as a warning and the OFT falls back to FULL refresh.
+        """
+        # Only POSTGRES-backed OFTs whose source is a plain View need this. The offline object is a
+        # View exactly when refresh_freq is None (managed FVs are DT-backed); streaming/realtime FVs
+        # are excluded because their OFT does not read from an offline view. HYBRID_TABLE is out of
+        # scope: enabling change tracking has base-table side effects we don't impose on the legacy path.
+        online_config = feature_view.online_config
+        oft_reads_from_source_view = (
+            online_config is not None
+            and online_config.store_type == OnlineStoreType.POSTGRES
+            and feature_view.refresh_freq is None
+            and not feature_view.is_streaming
+            and not feature_view.is_realtime_feature_view
+        )
+        if not oft_reads_from_source_view:
+            return
+
+        refresh_mode = feature_view.refresh_mode
+        if refresh_mode is not None and refresh_mode.upper() not in ("AUTO", "INCREMENTAL"):
+            return
+
+        try:
+            self._session.sql(fv_mod.build_enable_view_change_tracking_sql(fully_qualified_source_name)).collect(
+                statement_params=self._telemetry_stmp
+            )
+        except SnowparkSQLException as e:
+            if refresh_mode is not None and refresh_mode.upper() == "INCREMENTAL":
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.SNOWML_CREATE_FAILED,
+                    original_exception=RuntimeError(
+                        f"Failed to enable change tracking on source view {fully_qualified_source_name}, "
+                        f"which the online feature table requires to refresh incrementally: {e}"
+                    ),
+                ) from e
+            logger.warning(
+                "Could not enable change tracking on source view %s; "
+                "the online feature table will use FULL refresh: %s",
+                fully_qualified_source_name,
+                e,
+            )
+
     def _build_batch_feature_view_spec(
         self,
         feature_view: FeatureView,

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/feature_store/feature_view.py

@@ -2949,6 +2949,27 @@ def build_oft_warehouse_clause(
     return ""
 
 
+def build_enable_view_change_tracking_sql(fully_qualified_view_name: str) -> str:
+    """Build ``ALTER VIEW ... SET CHANGE_TRACKING = TRUE`` for an OFT source view.
+
+    An Online Feature Table can only resolve to an incremental refresh when the
+    object it reads from exposes change tracking. Managed feature views are
+    backed by a Dynamic Table (change tracking is inherent to an incremental DT
+    and cannot be enabled on a FULL one), but static/external feature views are
+    backed by a plain View that has change tracking disabled by default. Without
+    it the backend silently downgrades the OFT to a FULL refresh. This statement
+    is a no-op when change tracking is already enabled.
+
+    Args:
+        fully_qualified_view_name: ``DB.SCHEMA.VIEW`` the Online Feature Table
+            reads from.
+
+    Returns:
+        The ``ALTER VIEW ... SET CHANGE_TRACKING = TRUE`` SQL.
+    """
+    return f"ALTER VIEW {fully_qualified_view_name} SET CHANGE_TRACKING = TRUE"
+
+
 def build_oft_create_sql(
     *,
     fully_qualified_oft_name: str,

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/jobs/_utils/type_utils.py

@@ -148,3 +148,4 @@ class SpecOptions:
     spec_overrides: Optional[dict[str, Any]] = None
     runtime: Optional[str] = None
     enable_stage_mount_v2: Optional[bool] = True
+    artifact_repositories: Optional[list[str]] = None

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/jobs/decorators.py

@@ -48,6 +48,8 @@ def remote(
             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.
+            artifact_repositories (list[str]): (Public Preview) A list of artifact repositories to use for
+                resolving pip requirements. Each entry is the name of a Snowflake artifact repository.
             name (str): The name of the job. If not specified, a name will be generated based on the
                 entrypoint file name or function name. A unique suffix is appended to the final name.
 

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/jobs/job_definition.py

@@ -102,6 +102,7 @@ class MLJobDefinition(Generic[_Args, _ReturnValue], SerializableSessionMixin):
         database: Optional[str] = None,
         schema: Optional[str] = None,
         payload: Optional[payload_utils.JobPayload] = None,
+        artifact_repositories: Optional[list[str]] = None,
     ) -> None:
         self.source = source
         self.compute_pool = identifier.resolve_identifier(compute_pool)
@@ -122,6 +123,7 @@ class MLJobDefinition(Generic[_Args, _ReturnValue], SerializableSessionMixin):
         self.runtime_environment = runtime_environment
         self.overwrite = overwrite
         self.payload = payload
+        self.artifact_repositories = artifact_repositories
 
         self._is_registered = False
 
@@ -197,6 +199,7 @@ class MLJobDefinition(Generic[_Args, _ReturnValue], SerializableSessionMixin):
             spec_overrides=self.spec_overrides,
             runtime=self.runtime_environment if self.runtime_environment else None,
             enable_stage_mount_v2=feature_flags.FeatureFlags.ENABLE_STAGE_MOUNT_V2.is_enabled(),
+            artifact_repositories=self.artifact_repositories if self.artifact_repositories else None,
         )
 
         self.job_options = type_utils.JobOptions(
@@ -304,6 +307,9 @@ class MLJobDefinition(Generic[_Args, _ReturnValue], SerializableSessionMixin):
         target_instances = kwargs.pop("target_instances", 1)
         min_instances = kwargs.pop("min_instances", target_instances)
         pip_requirements = kwargs.pop("pip_requirements", None)
+        artifact_repositories = kwargs.pop("artifact_repositories", None)
+        if artifact_repositories is not None:
+            logger.warning("'artifact_repositories' is in public preview since 1.50.0.")
         external_access_integrations = kwargs.pop("external_access_integrations", None)
         env_vars = kwargs.pop("env_vars", None)
         spec_overrides = kwargs.pop("spec_overrides", None)
@@ -368,6 +374,7 @@ class MLJobDefinition(Generic[_Args, _ReturnValue], SerializableSessionMixin):
             database=database,
             schema=schema,
             payload=payload,
+            artifact_repositories=artifact_repositories,
         )
 
     @classmethod

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/jobs/manager.py

@@ -235,6 +235,8 @@ def submit_file(
             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.
+            artifact_repositories (list[str]): (Public Preview) A list of artifact repositories to use for
+                resolving pip requirements. Each entry is the name of a Snowflake artifact repository.
             name (str): The name of the job. If not specified, a name will be generated based on the
                 entrypoint file name or function name. A unique suffix is appended to the final name.
 
@@ -298,6 +300,8 @@ def submit_directory(
             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.
+            artifact_repositories (list[str]): (Public Preview) A list of artifact repositories to use for
+                resolving pip requirements. Each entry is the name of a Snowflake artifact repository.
             name (str): The name of the job. If not specified, a name will be generated based on the
                 entrypoint file name or function name. A unique suffix is appended to the final name.
 
@@ -363,6 +367,8 @@ def submit_from_stage(
             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.
+            artifact_repositories (list[str]): (Public Preview) A list of artifact repositories to use for
+                resolving pip requirements. Each entry is the name of a Snowflake artifact repository.
             name (str): The name of the job. If not specified, a name will be generated based on the
                 entrypoint file name or function name. A unique suffix is appended to the final name.
 

{snowflake_ml_python-1.50.0 → snowflake_ml_python-1.51.0}/snowflake/ml/model/_client/model/batch_inference_task.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import Any, Optional
 
 import yaml
@@ -31,6 +32,11 @@ class BatchInferenceTask(DAGTask):
     Construct it inside a ``with DAG(...)`` block (or pass ``dag=`` explicitly) and
     chain it with other tasks using ``>>``. Requires the ``snowflake.core`` package
     to be installed.
+
+    .. deprecated::
+        The backend this task uses will be removed in a future release. For new work, use
+        :class:`snowflake.ml.model.batch_inference.BatchInferenceTask`, which runs on the
+        updated batch inference backend.
     """
 
     @telemetry.send_api_usage_telemetry(
@@ -58,6 +64,13 @@ class BatchInferenceTask(DAGTask):
         inference_engine_options: Optional[dict[str, Any]] = None,
         **dagtask_kwargs: Any,
     ) -> None:
+        warnings.warn(
+            "snowflake.ml.model.batch.BatchInferenceTask is deprecated: the batch inference backend it "
+            "currently uses will be removed in a future release. Use "
+            "snowflake.ml.model.batch_inference.BatchInferenceTask instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         if not _HAS_SNOWFLAKE_CORE:
             raise ImportError(
                 "BatchInferenceTask requires the `snowflake.core` package. "

snowflake_ml_python-1.51.0/snowflake/ml/model/_client/model/batch_inference_task_v2.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from snowflake.ml._internal import telemetry
+from snowflake.ml.model._client.model import (
+    batch_inference_job_specs,
+    model_version_impl,
+)
+from snowflake.ml.model._client.ops import service_ops
+
+try:
+    from snowflake.core.task.dagv1 import DAGTask
+
+    _HAS_SNOWFLAKE_CORE = True
+except ModuleNotFoundError:
+    DAGTask = object  # type: ignore[assignment, misc]
+    _HAS_SNOWFLAKE_CORE = False
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
+
+class BatchInferenceTask(DAGTask):
+    """A ``DAGTask`` that runs a batch inference job for a registered model version.
+
+    Construct it inside a ``with DAG(...)`` block (or pass ``dag=`` explicitly) and
+    chain it with other tasks using ``>>``. Requires the ``snowflake.core`` package
+    to be installed.
+
+    The job runs synchronously, so the DAG step completes when the job does. Each run
+    gets a server-generated job name, and results are written under
+    ``<output_spec.stage_location>/<job_name>/``. The job publishes that location as
+    the task return value, so a successor can read it with
+    ``SYSTEM$GET_PREDECESSOR_RETURN_VALUE()`` rather than hardcoding the path.
+
+    The DAG must supply a warehouse, via ``DAG(warehouse=...)`` or ``warehouse=`` on the
+    task; serverless DAGs are unsupported.
+
+    Prefer fully qualified stage paths: an unqualified path is resolved against the
+    session namespace when the task is built, but against the task owner's namespace
+    when it runs.
+    """
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=[
+            "name",
+            "compute_pool",
+            "input_spec",
+            "output_spec",
+            "resources_spec",
+            "inference_spec",
+            "image_build_spec",
+            "replicas",
+        ],
+    )
+    def __init__(
+        self,
+        name: str,
+        *,
+        model_version: model_version_impl.ModelVersion,
+        compute_pool: str,
+        output_spec: batch_inference_job_specs.OutputSpec,
+        query: Optional[str] = None,
+        input_stage_location: Optional[str] = None,
+        input_spec: Optional[batch_inference_job_specs.InputSpec] = None,
+        resources_spec: Optional[batch_inference_job_specs.ResourcesSpec] = None,
+        inference_spec: Optional[batch_inference_job_specs.InferenceSpec] = None,
+        image_build_spec: Optional[batch_inference_job_specs.ImageBuildSpec] = None,
+        function_name: Optional[str] = None,
+        replicas: Optional[int] = None,
+        **dagtask_kwargs: Any,
+    ) -> None:
+        """Build a batch inference task.
+
+        Args:
+            name: Task name within the DAG.
+            model_version: The model version to run.
+            compute_pool: Compute pool used by the job and the image build.
+            output_spec: Output block. ``stage_location`` is treated as a base; results are
+                written under ``<stage_location>/<job_name>/``.
+            query: SQL query producing the input rows. Provide exactly one of ``query`` or
+                ``input_stage_location``.
+            input_stage_location: Existing stage path holding the input data. Provide exactly
+                one of ``query`` or ``input_stage_location``.
+            input_spec: Optional input block.
+            resources_spec: Optional resources block.
+            inference_spec: Optional inference block.
+            image_build_spec: Optional image build block.
+            function_name: Model function name. Resolved against the model's function list
+                when omitted.
+            replicas: Optional ``REPLICAS`` value.
+            dagtask_kwargs: Passed through to ``DAGTask`` (``dag``, ``condition``,
+                ``warehouse``, ``session_parameters``, and so on).
+
+        Raises:
+            ImportError: If the ``snowflake.core`` package is not installed.
+            TypeError: If ``definition`` is passed; the task builds its own.
+        """
+        if not _HAS_SNOWFLAKE_CORE:
+            raise ImportError(
+                "BatchInferenceTask requires the `snowflake.core` package. "
+                "Install it with `pip install snowflake.core`."
+            )
+        if "definition" in dagtask_kwargs:
+            raise TypeError("BatchInferenceTask builds its own task definition; do not pass `definition=`.")
+
+        self._model_version = model_version
+        self._compute_pool = compute_pool
+        self._output_spec = output_spec
+        self._query = query
+        self._input_stage_location = input_stage_location
+        self._input_spec = input_spec
+        self._resources_spec = resources_spec
+        self._inference_spec = inference_spec
+        self._image_build_spec = image_build_spec
+        self._function_name = function_name
+        self._replicas = replicas
+
+        super().__init__(name, definition=self._to_sql(), **dagtask_kwargs)
+
+    def _to_sql(self) -> str:
+        """Build the ``EXECUTE INFERENCE JOB SERVICE`` command that becomes the task definition.
+
+        Returns:
+            The command text.
+
+        Raises:
+            ValueError: If not exactly one of ``query`` / ``input_stage_location`` is provided.
+        """
+        if (self._query is None) == (self._input_stage_location is None):
+            raise ValueError("Exactly one of query or input_stage_location must be provided.")
+
+        target_function_info = self._model_version._validate_batch_inference_request(
+            input_spec=self._input_spec,
+            resources_spec=self._resources_spec,
+            function_name=self._function_name,
+        )
+
+        return service_ops.build_batch_inference_task_definition(
+            session=self._model_version._service_ops._session,
+            model_fqn=self._model_version.fully_qualified_model_name,
+            version_name=self._model_version.version_name,
+            compute_pool=self._compute_pool,
+            function_name=target_function_info["target_method"],
+            query=self._query,
+            input_stage_location=self._input_stage_location,
+            input_spec=self._input_spec,
+            output_spec=self._output_spec,
+            resources_spec=self._resources_spec,
+            inference_spec=self._inference_spec,
+            image_build_spec=self._image_build_spec,
+            replicas=self._replicas,
+        )