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

snowflake/ml/jobs/job.py

@@ -12,12 +12,19 @@ from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml._internal.utils.mixins import SerializableSessionMixin
-from snowflake.ml.jobs._utils import constants, interop_utils, query_helper, types
+from snowflake.ml.jobs._interop import results as interop_result, utils as interop_utils
+from snowflake.ml.jobs._utils import (
+    constants,
+    payload_utils,
+    query_helper,
+    stage_utils,
+    types,
+)
 from snowflake.snowpark import Row, context as sp_context
 from snowflake.snowpark.exceptions import SnowparkSQLException
 
 _PROJECT = "MLJob"
-TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "CANCELLED", "INTERNAL_ERROR"}
+TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "CANCELLED", "INTERNAL_ERROR", "DELETED"}
 
 T = TypeVar("T")
 
@@ -36,7 +43,12 @@ class MLJob(Generic[T], SerializableSessionMixin):
         self._session = session or sp_context.get_active_session()
 
         self._status: types.JOB_STATUS = "PENDING"
-        self._result: Optional[interop_utils.ExecutionResult] = None
+        self._result: Optional[interop_result.ExecutionResult] = None
+
+    @cached_property
+    def _service_info(self) -> types.ServiceInfo:
+        """Get the job's service info."""
+        return _resolve_service_info(self.id, self._session)
 
     @cached_property
     def name(self) -> str:
@@ -44,7 +56,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
 
     @cached_property
     def target_instances(self) -> int:
-        return _get_target_instances(self._session, self.id)
+        return self._service_info.target_instances
 
     @cached_property
     def min_instances(self) -> int:
@@ -69,8 +81,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
     @cached_property
     def _compute_pool(self) -> str:
         """Get the job's compute pool name."""
-        row = _get_service_info(self._session, self.id)
-        return cast(str, row["compute_pool"])
+        return self._service_info.compute_pool
 
     @property
     def _service_spec(self) -> dict[str, Any]:
@@ -82,7 +93,13 @@ class MLJob(Generic[T], SerializableSessionMixin):
     @property
     def _container_spec(self) -> dict[str, Any]:
         """Get the job's main container spec."""
-        containers = self._service_spec["spec"]["containers"]
+        try:
+            containers = self._service_spec["spec"]["containers"]
+        except SnowparkSQLException as e:
+            if e.sql_error_code == 2003:
+                # If the job is deleted, the service spec is not available
+                return {}
+            raise
         if len(containers) == 1:
             return cast(dict[str, Any], containers[0])
         try:
@@ -105,22 +122,28 @@ class MLJob(Generic[T], SerializableSessionMixin):
         if result_path_str is None:
             raise RuntimeError(f"Job {self.name} doesn't have a result path configured")
 
-        # If result path is relative, it is relative to the stage mount path
-        result_path = Path(result_path_str)
-        if not result_path.is_absolute():
-            return f"{self._stage_path}/{result_path.as_posix()}"
+        return self._transform_path(result_path_str)
 
-        # If result path is absolute, it is relative to the stage mount path
+    def _transform_path(self, path_str: str) -> str:
+        """Transform a local path within the container to a stage path."""
+        path = payload_utils.resolve_path(path_str)
+        if isinstance(path, stage_utils.StagePath):
+            # Stage paths need no transformation
+            return path.as_posix()
+        if not path.is_absolute():
+            # Assume relative paths are relative to stage mount path
+            return f"{self._stage_path}/{path.as_posix()}"
+
+        # If result path is absolute, rebase it onto the stage mount path
+        # TODO: Rather than matching by name, use the longest mount path which matches
         volume_mounts = self._container_spec["volumeMounts"]
         stage_mount_str = next(v for v in volume_mounts if v.get("name") == constants.STAGE_VOLUME_NAME)["mountPath"]
         stage_mount = Path(stage_mount_str)
         try:
-            relative_path = result_path.relative_to(stage_mount)
+            relative_path = path.relative_to(stage_mount)
             return f"{self._stage_path}/{relative_path.as_posix()}"
         except ValueError:
-            raise ValueError(
-                f"Result path {result_path} is absolute, but should be relative to stage mount {stage_mount}"
-            )
+            raise ValueError(f"Result path {path} is absolute, but should be relative to stage mount {stage_mount}")
 
     @overload
     def get_logs(
@@ -165,7 +188,14 @@ class MLJob(Generic[T], SerializableSessionMixin):
         Returns:
             The job's execution logs.
         """
-        logs = _get_logs(self._session, self.id, limit, instance_id, self._container_spec["name"], verbose)
+        logs = _get_logs(
+            self._session,
+            self.id,
+            limit,
+            instance_id,
+            self._container_spec["name"] if "name" in self._container_spec else constants.DEFAULT_CONTAINER_NAME,
+            verbose,
+        )
         assert isinstance(logs, str)  # mypy
         if as_list:
             return logs.splitlines()
@@ -218,7 +248,6 @@ class MLJob(Generic[T], SerializableSessionMixin):
             delay = min(delay * 1.2, constants.JOB_POLL_MAX_DELAY_SECONDS)  # Exponential backoff
         return self.status
 
-    @snowpark._internal.utils.private_preview(version="1.8.2")
     @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["timeout"])
     def result(self, timeout: float = -1) -> T:
         """
@@ -237,13 +266,13 @@ class MLJob(Generic[T], SerializableSessionMixin):
         if self._result is None:
             self.wait(timeout)
             try:
-                self._result = interop_utils.fetch_result(self._session, self._result_path)
+                self._result = interop_utils.load_result(
+                    self._result_path, session=self._session, path_transform=self._transform_path
+                )
             except Exception as e:
-                raise RuntimeError(f"Failed to retrieve result for job (id={self.name})") from e
+                raise RuntimeError(f"Failed to retrieve result for job, error: {e!r}") from e
 
-        if self._result.success:
-            return cast(T, self._result.result)
-        raise RuntimeError(f"Job execution failed (id={self.name})") from self._result.exception
+        return cast(T, self._result.get_value())
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
     def cancel(self) -> None:
@@ -256,22 +285,28 @@ class MLJob(Generic[T], SerializableSessionMixin):
             self._session.sql(f"CALL {self.id}!spcs_cancel_job()").collect()
             logger.debug(f"Cancellation requested for job {self.id}")
         except SnowparkSQLException as e:
-            raise RuntimeError(f"Failed to cancel job {self.id}: {e.message}") from e
+            raise RuntimeError(f"Failed to cancel job, error: {e!r}") from e
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "instance_id"])
 def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[int] = None) -> types.JOB_STATUS:
     """Retrieve job or job instance execution status."""
-    if instance_id is not None:
-        # Get specific instance status
-        rows = session.sql("SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
-        for row in rows:
-            if row["instance_id"] == str(instance_id):
-                return cast(types.JOB_STATUS, row["status"])
-        raise ValueError(f"Instance {instance_id} not found in job {job_id}")
-    else:
-        row = _get_service_info(session, job_id)
-        return cast(types.JOB_STATUS, row["status"])
+    try:
+        if instance_id is not None:
+            # Get specific instance status
+            rows = query_helper.run_query(session, "SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,))
+            for row in rows:
+                if row["instance_id"] == str(instance_id):
+                    return cast(types.JOB_STATUS, row["status"])
+            raise ValueError(f"Instance {instance_id} not found in job {job_id}")
+        else:
+            row = _get_service_info(session, job_id)
+            return cast(types.JOB_STATUS, row["status"])
+    except SnowparkSQLException as e:
+        if e.sql_error_code == 2003:
+            row = _get_service_info_spcs(session, job_id)
+            return cast(types.JOB_STATUS, row["STATUS"])
+        raise
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
@@ -542,8 +577,21 @@ def _get_compute_pool_info(session: snowpark.Session, compute_pool: str) -> Row:
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
 def _get_target_instances(session: snowpark.Session, job_id: str) -> int:
-    row = _get_service_info(session, job_id)
-    return int(row["target_instances"])
+    try:
+        row = _get_service_info(session, job_id)
+        return int(row["target_instances"])
+    except SnowparkSQLException as e:
+        if e.sql_error_code == 2003:
+            row = _get_service_info_spcs(session, job_id)
+            try:
+                params = json.loads(row["PARAMETERS"])
+                if isinstance(params, dict):
+                    return int(params.get("REPLICAS", 1))
+                else:
+                    return 1
+            except (json.JSONDecodeError, ValueError):
+                return 1
+        raise
 
 
 def _get_logs_spcs(
@@ -581,3 +629,87 @@ def _get_logs_spcs(
         query.append(f" LIMIT {limit};")
     rows = session.sql("\n".join(query)).collect()
     return rows
+
+
+def _get_service_info_spcs(session: snowpark.Session, job_id: str) -> Any:
+    """
+    Retrieve the service info from the SPCS interface.
+
+    Args:
+        session (Session): The Snowpark session to use.
+        job_id (str): The job ID.
+
+    Returns:
+        Any: The service info.
+
+    Raises:
+        SnowparkSQLException: If the job does not exist or is too old to retrieve.
+    """
+    db, schema, name = identifier.parse_schema_level_object_identifier(job_id)
+    db = db or session.get_current_database()
+    schema = schema or session.get_current_schema()
+    rows = query_helper.run_query(
+        session,
+        """
+        select DATABASE_NAME, SCHEMA_NAME, NAME, STATUS, COMPUTE_POOL_NAME, PARAMETERS
+        from table(snowflake.spcs.get_job_history())
+        where database_name = ? and schema_name = ? and name = ?
+        """,
+        params=(db, schema, name),
+    )
+    if rows:
+        return rows[0]
+    else:
+        raise SnowparkSQLException(f"Job {job_id} does not exist or could not be retrieved", sql_error_code=2003)
+
+
+def _resolve_service_info(id: str, session: snowpark.Session) -> types.ServiceInfo:
+    try:
+        row = _get_service_info(session, id)
+    except SnowparkSQLException as e:
+        if e.sql_error_code == 2003:
+            row = _get_service_info_spcs(session, id)
+        else:
+            raise
+    if not row:
+        raise SnowparkSQLException(f"Job {id} does not exist or could not be retrieved", sql_error_code=2003)
+
+    if "compute_pool" in row:
+        compute_pool = row["compute_pool"]
+    elif "COMPUTE_POOL_NAME" in row:
+        compute_pool = row["COMPUTE_POOL_NAME"]
+    else:
+        raise ValueError(f"compute_pool not found in row: {row}")
+
+    if "status" in row:
+        status = row["status"]
+    elif "STATUS" in row:
+        status = row["STATUS"]
+    else:
+        raise ValueError(f"status not found in row: {row}")
+    # Normalize target_instances
+    target_instances: int
+    if "target_instances" in row and row["target_instances"] is not None:
+        try:
+            target_instances = int(row["target_instances"])
+        except (ValueError, TypeError):
+            target_instances = 1
+    elif "PARAMETERS" in row and row["PARAMETERS"]:
+        try:
+            params = json.loads(row["PARAMETERS"])
+            target_instances = int(params.get("REPLICAS", 1)) if isinstance(params, dict) else 1
+        except (json.JSONDecodeError, ValueError, TypeError):
+            target_instances = 1
+    else:
+        target_instances = 1
+
+    database_name = row["database_name"] if "database_name" in row else row["DATABASE_NAME"]
+    schema_name = row["schema_name"] if "schema_name" in row else row["SCHEMA_NAME"]
+
+    return types.ServiceInfo(
+        database_name=database_name,
+        schema_name=schema_name,
+        status=cast(types.JOB_STATUS, status),
+        compute_pool=cast(str, compute_pool),
+        target_instances=target_instances,
+    )

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
@@ -20,6 +21,7 @@ from snowflake.ml.jobs._utils import (
     spec_utils,
     types,
 )
+from snowflake.snowpark._internal import utils as sp_utils
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import coalesce, col, lit, when
@@ -178,8 +180,10 @@ def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob
         _ = job._service_spec
         return job
     except SnowparkSQLException as e:
-        if "does not exist" in e.message:
-            raise ValueError(f"Job does not exist: {job_id}") from e
+        if e.sql_error_code == 2003:
+            job = jb.MLJob[Any](job_id, session=session)
+            _ = job.status
+            return job
         raise
 
 
@@ -344,6 +348,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 +416,7 @@ def _submit_job(
         "min_instances",
         "enable_metrics",
         "query_warehouse",
+        "runtime_environment",
     ],
 )
 def _submit_job(
@@ -441,7 +449,7 @@ def _submit_job(
     Raises:
         ValueError: If database or schema value(s) are invalid
         RuntimeError: If schema is not specified in session context or job submission
-        snowpark.exceptions.SnowparkSQLException: if failed to upload payload
+        SnowparkSQLException: if failed to upload payload
     """
     session = _ensure_session(session)
 
@@ -459,6 +467,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 +481,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:
@@ -503,48 +515,44 @@ def _submit_job(
         uploaded_payload = payload_utils.JobPayload(
             source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=imports
         ).upload(session, stage_path)
-    except snowpark.exceptions.SnowparkSQLException as e:
+    except SnowparkSQLException as e:
         if e.sql_error_code == 90106:
             raise RuntimeError(
                 "Please specify a schema, either in the session context or as a parameter in the job submission"
             )
         raise
 
-    # FIXME: Temporary patches, remove this after v1 is deprecated
-    if target_instances > 1:
-        default_spec_overrides = {
-            "spec": {
-                "endpoints": [
-                    {"name": "ray-dashboard-endpoint", "port": 12003, "protocol": "TCP"},
-                ]
-            },
-        }
-        if spec_overrides:
-            spec_overrides = spec_utils.merge_patch(
-                default_spec_overrides, spec_overrides, display_name="spec_overrides"
-            )
-        else:
-            spec_overrides = default_spec_overrides
-
-    if feature_flags.FeatureFlags.USE_SUBMIT_JOB_V2.is_enabled():
+    if feature_flags.FeatureFlags.USE_SUBMIT_JOB_V2.is_enabled(default=True):
         # Add default env vars (extracted from spec_utils.generate_service_spec)
         combined_env_vars = {**uploaded_payload.env_vars, **(env_vars or {})}
 
-        return _do_submit_job_v2(
-            session=session,
-            payload=uploaded_payload,
-            args=args,
-            env_vars=combined_env_vars,
-            spec_overrides=spec_overrides,
-            compute_pool=compute_pool,
-            job_id=job_id,
-            external_access_integrations=external_access_integrations,
-            query_warehouse=query_warehouse,
-            target_instances=target_instances,
-            min_instances=min_instances,
-            enable_metrics=enable_metrics,
-            use_async=True,
-        )
+        try:
+            return _do_submit_job_v2(
+                session=session,
+                payload=uploaded_payload,
+                args=args,
+                env_vars=combined_env_vars,
+                spec_overrides=spec_overrides,
+                compute_pool=compute_pool,
+                job_id=job_id,
+                external_access_integrations=external_access_integrations,
+                query_warehouse=query_warehouse,
+                target_instances=target_instances,
+                min_instances=min_instances,
+                enable_metrics=enable_metrics,
+                use_async=True,
+                runtime_environment=runtime_environment,
+            )
+        except SnowparkSQLException as e:
+            if not (e.sql_error_code == 90237 and sp_utils.is_in_stored_procedure()):  # type: ignore[no-untyped-call]
+                raise
+            # SNOW-2390287: SYSTEM$EXECUTE_ML_JOB() is erroneously blocked in owner's rights
+            # stored procedures. This will be fixed in an upcoming release.
+            logger.warning(
+                "Job submission using V2 failed with error {}. Falling back to V1.".format(
+                    str(e).split("\n", 1)[0],
+                )
+            )
 
     # Fall back to v1
     # Generate service spec
@@ -556,6 +564,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 +648,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 +667,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 +683,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_RUNTIME_VERSIONS.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/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/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

snowflake/ml/model/_client/model/model_version_impl.py

@@ -19,7 +19,9 @@ from snowflake.ml.model._client.model import (
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
+from snowflake.ml.model._model_composer.model_method import utils as model_method_utils
 from snowflake.ml.model._packager.model_handlers import snowmlmodel
+from snowflake.ml.model._packager.model_meta import model_meta_schema
 from snowflake.snowpark import Session, async_job, dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
@@ -41,6 +43,7 @@ class ModelVersion(lineage_node.LineageNode):
     _model_name: sql_identifier.SqlIdentifier
     _version_name: sql_identifier.SqlIdentifier
     _functions: list[model_manifest_schema.ModelFunctionInfo]
+    _model_spec: Optional[model_meta_schema.ModelMetadataDict]
 
     def __init__(self) -> None:
         raise RuntimeError("ModelVersion's initializer is not meant to be used. Use `version` from model instead.")
@@ -150,6 +153,7 @@ class ModelVersion(lineage_node.LineageNode):
         self._model_name = model_name
         self._version_name = version_name
         self._functions = self._get_functions()
+        self._model_spec = None
         super(cls, cls).__init__(
             self,
             session=model_ops._session,
@@ -437,6 +441,26 @@ class ModelVersion(lineage_node.LineageNode):
         """
         return self._functions
 
+    def _get_model_spec(self, statement_params: Optional[dict[str, Any]] = None) -> model_meta_schema.ModelMetadataDict:
+        """Fetch and cache the model spec for this model version.
+
+        Args:
+            statement_params: Optional dictionary of statement parameters to include
+                in the SQL command to fetch the model spec.
+
+        Returns:
+            The model spec as a dictionary for this model version.
+        """
+        if self._model_spec is None:
+            self._model_spec = self._model_ops._fetch_model_spec(
+                database_name=None,
+                schema_name=None,
+                model_name=self._model_name,
+                version_name=self._version_name,
+                statement_params=statement_params,
+            )
+        return self._model_spec
+
     @overload
     def run(
         self,
@@ -531,6 +555,8 @@ class ModelVersion(lineage_node.LineageNode):
                 statement_params=statement_params,
             )
         else:
+            explain_case_sensitive = self._determine_explain_case_sensitivity(target_function_info, statement_params)
+
             return self._model_ops.invoke_method(
                 method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
                 method_function_type=target_function_info["target_method_function_type"],
@@ -544,13 +570,27 @@ class ModelVersion(lineage_node.LineageNode):
                 partition_column=partition_column,
                 statement_params=statement_params,
                 is_partitioned=target_function_info["is_partitioned"],
+                explain_case_sensitive=explain_case_sensitive,
             )
 
+    def _determine_explain_case_sensitivity(
+        self,
+        target_function_info: model_manifest_schema.ModelFunctionInfo,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> bool:
+        model_spec = self._get_model_spec(statement_params)
+        method_options = model_spec.get("method_options", {})
+        return model_method_utils.determine_explain_case_sensitive_from_method_options(
+            method_options, target_function_info["name"]
+        )
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
         func_params_to_log=[
             "compute_pool",
+            "output_spec",
+            "job_spec",
         ],
     )
     def _run_batch(
@@ -579,6 +619,8 @@ class ModelVersion(lineage_node.LineageNode):
             output_stage_location += "/"
         input_stage_location = f"{output_stage_location}{_BATCH_INFERENCE_TEMPORARY_FOLDER}/"
 
+        self._service_ops._enforce_save_mode(output_spec.mode, output_stage_location)
+
         try:
             input_spec.write.copy_into_location(location=input_stage_location, file_format_type="parquet", header=True)
         # todo: be specific about the type of errors to provide better error messages.
@@ -605,6 +647,7 @@ class ModelVersion(lineage_node.LineageNode):
             warehouse=sql_identifier.SqlIdentifier(warehouse),
             cpu_requests=job_spec.cpu_requests,
             memory_requests=job_spec.memory_requests,
+            gpu_requests=job_spec.gpu_requests,
             job_name=job_name,
             replicas=job_spec.replicas,
             # input and output
@@ -798,13 +841,7 @@ class ModelVersion(lineage_node.LineageNode):
             ValueError: If the model is not a HuggingFace text-generation model.
         """
         # Fetch model spec
-        model_spec = self._model_ops._fetch_model_spec(
-            database_name=None,
-            schema_name=None,
-            model_name=self._model_name,
-            version_name=self._version_name,
-            statement_params=statement_params,
-        )
+        model_spec = self._get_model_spec(statement_params)
 
         # Check if model_type is huggingface_pipeline
         model_type = model_spec.get("model_type")