snowflake-ml-python 1.22.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

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

@@ -1,6 +1,4 @@
-import base64
 import enum
-import json
 import pathlib
 import tempfile
 import uuid
@@ -8,7 +6,6 @@ import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
-from pydantic import TypeAdapter
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -33,7 +30,10 @@ _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 _BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 _BATCH_INFERENCE_TEMPORARY_FOLDER = "_temporary"
-_UTF8_ENCODING = "utf-8"
+VLLM_SUPPORTED_TASKS = [
+    "text-generation",
+    "image-text-to-text",
+]
 
 
 class ExportMode(enum.Enum):
@@ -649,41 +649,6 @@ class ModelVersion(lineage_node.LineageNode):
             method_options, target_function_info["name"]
         )
 
-    @staticmethod
-    def _encode_column_handling(
-        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
-    ) -> Optional[str]:
-        """Validate and encode column_handling to a base64 string.
-
-        Args:
-            column_handling: Optional dictionary mapping column names to file encoding options.
-
-        Returns:
-            Base64 encoded JSON string of the column handling options, or None if input is None.
-        """
-        # TODO: validation for column names
-        if column_handling is None:
-            return None
-        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
-        # TODO: throw error if the validate_python function fails
-        validated_input = adapter.validate_python(column_handling)
-        return base64.b64encode(adapter.dump_json(validated_input)).decode(_UTF8_ENCODING)
-
-    @staticmethod
-    def _encode_params(params: Optional[dict[str, Any]]) -> Optional[str]:
-        """Encode params dictionary to a base64 string.
-
-        Args:
-            params: Optional dictionary of model inference parameters.
-
-        Returns:
-            Base64 encoded JSON string of the params, or None if input is None.
-        """
-        if params is None:
-            return None
-        # TODO: validation for param names, types
-        return base64.b64encode(json.dumps(params).encode(_UTF8_ENCODING)).decode(_UTF8_ENCODING)
-
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -696,32 +661,33 @@ class ModelVersion(lineage_node.LineageNode):
     @snowpark._internal.utils.private_preview(version="1.18.0")
     def run_batch(
         self,
+        X: dataframe.DataFrame,
         *,
         compute_pool: str,
-        input_spec: dataframe.DataFrame,
+        input_spec: Optional[batch_inference_specs.InputSpec] = None,
         output_spec: batch_inference_specs.OutputSpec,
         job_spec: Optional[batch_inference_specs.JobSpec] = None,
-        params: Optional[dict[str, Any]] = None,
-        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]] = None,
+        inference_engine_options: Optional[dict[str, Any]] = None,
     ) -> job.MLJob[Any]:
         """Execute batch inference on datasets as an SPCS job.
 
         Args:
             compute_pool (str): Name of the compute pool to use for building the image containers and batch
                 inference execution.
-            input_spec (dataframe.DataFrame): Snowpark DataFrame containing the input data for inference.
+            X (dataframe.DataFrame): Snowpark DataFrame containing the input data for inference.
                 The DataFrame should contain all required features for model prediction and passthrough columns.
             output_spec (batch_inference_specs.OutputSpec): Configuration for where and how to save
                 the inference results. Specifies the stage location and file handling behavior.
+            input_spec (Optional[batch_inference_specs.InputSpec]): Optional configuration for input
+                processing including model inference parameters and column handling options.
+                If None, default values will be used for params and column_handling.
             job_spec (Optional[batch_inference_specs.JobSpec]): Optional configuration for job
                 execution parameters such as compute resources, worker counts, and job naming.
                 If None, default values will be used.
-            params (Optional[dict[str, Any]]): Optional dictionary of model inference parameters
-                (e.g., temperature, top_k for LLMs). These are passed as keyword arguments to the
-                model's inference method. Defaults to None.
-            column_handling (Optional[dict[str, batch_inference_specs.FileEncoding]]): Optional dictionary
-                specifying how to handle specific columns during file I/O. Maps column names to their
-                file encoding configuration.
+            inference_engine_options: Options for the service creation with custom inference engine.
+                Supports `engine` and `engine_args_override`.
+                `engine` is the type of the inference engine to use.
+                `engine_args_override` is a list of string arguments to pass to the inference engine.
 
         Returns:
             job.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
@@ -729,7 +695,7 @@ class ModelVersion(lineage_node.LineageNode):
 
         Raises:
             ValueError: If warehouse is not set in job_spec and no current warehouse is available.
-            RuntimeError: If the input_spec cannot be processed or written to the staging location.
+            RuntimeError: If the input data cannot be processed or written to the staging location.
 
         Example:
             >>> # Prepare input data - Example 1: From a table
@@ -762,10 +728,24 @@ class ModelVersion(lineage_node.LineageNode):
             >>> # Run batch inference
             >>> job = model_version.run_batch(
             ...     compute_pool="my_compute_pool",
-            ...     input_spec=input_df,
+            ...     X=input_df,
             ...     output_spec=output_spec,
             ...     job_spec=job_spec
             ... )
+            >>>
+            >>> # Run batch inference with InputSpec for additional options
+            >>> from snowflake.ml.model._client.model.batch_inference_specs import InputSpec, FileEncoding
+            >>> input_spec = InputSpec(
+            ...     params={"temperature": 0.7, "top_k": 50},
+            ...     column_handling={"image_col": {"encoding": FileEncoding.BASE64}}
+            ... )
+            >>> job = model_version.run_batch(
+            ...     compute_pool="my_compute_pool",
+            ...     X=input_df,
+            ...     output_spec=output_spec,
+            ...     input_spec=input_spec,
+            ...     job_spec=job_spec
+            ... )
 
         Note:
             This method is currently in private preview and requires Snowflake version 1.18.0 or later.
@@ -777,12 +757,25 @@ class ModelVersion(lineage_node.LineageNode):
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
-        column_handling_as_string = self._encode_column_handling(column_handling)
-        params_as_string = self._encode_params(params)
+        # Extract params and column_handling from input_spec if provided
+        if input_spec is None:
+            input_spec = batch_inference_specs.InputSpec()
+
+        params = input_spec.params
+        column_handling = input_spec.column_handling
 
         if job_spec is None:
             job_spec = batch_inference_specs.JobSpec()
 
+        # Validate GPU support if GPU resources are requested
+        self._throw_error_if_gpu_is_not_supported(job_spec.gpu_requests, statement_params)
+
+        inference_engine_args = self._prepare_inference_engine_args(
+            inference_engine_options,
+            job_spec.gpu_requests,
+            statement_params,
+        )
+
         warehouse = job_spec.warehouse or self._service_ops._session.get_current_warehouse()
         if warehouse is None:
             raise ValueError("Warehouse is not set. Please set the warehouse field in the JobSpec.")
@@ -796,10 +789,10 @@ class ModelVersion(lineage_node.LineageNode):
         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)
+            X.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.
         except Exception as e:
-            raise RuntimeError(f"Failed to process input_spec: {e}")
+            raise RuntimeError(f"Failed to process input data: {e}")
 
         if job_spec.job_name is None:
             # Same as the MLJob ID generation logic with a different prefix
@@ -807,12 +800,14 @@ class ModelVersion(lineage_node.LineageNode):
         else:
             job_name = job_spec.job_name
 
+        target_function_info = self._get_function_info(function_name=job_spec.function_name)
+
         return self._service_ops.invoke_batch_job_method(
             # model version info
             model_name=self._model_name,
             version_name=self._version_name,
             # job spec
-            function_name=self._get_function_info(function_name=job_spec.function_name)["target_method"],
+            function_name=target_function_info["target_method"],
             compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
             force_rebuild=job_spec.force_rebuild,
             image_repo_name=job_spec.image_repo,
@@ -827,12 +822,14 @@ class ModelVersion(lineage_node.LineageNode):
             # input and output
             input_stage_location=input_stage_location,
             input_file_pattern="*",
-            column_handling=column_handling_as_string,
-            params=params_as_string,
+            column_handling=column_handling,
+            params=params,
+            signature_params=target_function_info["signature"].params,
             output_stage_location=output_stage_location,
             completion_filename="_SUCCESS",
             # misc
             statement_params=statement_params,
+            inference_engine_args=inference_engine_args,
         )
 
     def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo:
@@ -1048,20 +1045,55 @@ class ModelVersion(lineage_node.LineageNode):
                 " the `log_model` function."
             )
 
-    def _check_huggingface_text_generation_model(
+    def _prepare_inference_engine_args(
+        self,
+        inference_engine_options: Optional[dict[str, Any]],
+        gpu_requests: Optional[Union[str, int]],
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> Optional[service_ops.InferenceEngineArgs]:
+        """Prepare and validate inference engine arguments.
+
+        This method handles the common logic for processing inference engine options:
+        1. Parse inference engine options into InferenceEngineArgs
+        2. Validate that the model is a HuggingFace text-generation model (if inference engine is specified)
+        3. Enrich inference engine args
+
+        Args:
+            inference_engine_options: Optional dictionary containing inference engine configuration.
+            gpu_requests: GPU resource request string (e.g., "4").
+            statement_params: Optional dictionary of statement parameters for SQL commands.
+
+        Returns:
+            Prepared InferenceEngineArgs or None if no inference engine is specified.
+        """
+        inference_engine_args = inference_engine_utils._get_inference_engine_args(inference_engine_options)
+
+        if inference_engine_args is not None:
+            # Validate that model is HuggingFace vLLM supported model and is logged with
+            # OpenAI compatible signature.
+            self._check_huggingface_vllm_supported_model(statement_params)
+            # Enrich with GPU configuration
+            inference_engine_args = inference_engine_utils._enrich_inference_engine_args(
+                inference_engine_args,
+                gpu_requests,
+            )
+
+        return inference_engine_args
+
+    def _check_huggingface_vllm_supported_model(
         self,
         statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
-        """Check if the model is a HuggingFace pipeline with text-generation task
-        and is logged with OPENAI_CHAT_SIGNATURE.
+        """Check if the model is a HuggingFace pipeline with vLLM supported task
+        and is logged with OpenAI compatible signature.
 
         Args:
             statement_params: Optional dictionary of statement parameters to include
                 in the SQL command to fetch model spec.
 
         Raises:
-            ValueError: If the model is not a HuggingFace text-generation model or
-                if the model is not logged with OPENAI_CHAT_SIGNATURE.
+            ValueError: If the model is not a HuggingFace vLLM supported model or
+                if the model is not logged with OpenAI compatible signature.
         """
         # Fetch model spec
         model_spec = self._get_model_spec(statement_params)
@@ -1070,34 +1102,37 @@ class ModelVersion(lineage_node.LineageNode):
         model_type = model_spec.get("model_type")
         if model_type != "huggingface_pipeline":
             raise ValueError(
-                f"Inference engine is only supported for HuggingFace text-generation models. "
+                f"Inference engine is only supported for HuggingFace vLLM supported models. "
                 f"Found model_type: {model_type}"
             )
 
-        # Check if model supports text-generation task
+        # Check if model supports vLLM supported task
         # There should only be one model in the list because we don't support multiple models in a single model spec
         models = model_spec.get("models", {})
-        is_text_generation = False
+        is_vllm_supported_task = False
         found_tasks: list[str] = []
 
-        # As long as the model supports text-generation task, we can use it
+        # As long as the model supports vLLM supported task, we can use it
         for _, model_info in models.items():
             options = model_info.get("options", {})
             task = options.get("task")
             if task:
                 found_tasks.append(str(task))
-                if task == "text-generation":
-                    is_text_generation = True
+                if task in VLLM_SUPPORTED_TASKS:
+                    is_vllm_supported_task = True
                     break
 
-        if not is_text_generation:
+        if not is_vllm_supported_task:
             tasks_str = ", ".join(found_tasks)
             found_tasks_str = (
                 f"Found task(s): {tasks_str} in model spec." if found_tasks else "No task found in model spec."
             )
-            raise ValueError(f"Inference engine is only supported for task 'text-generation'. {found_tasks_str}")
+            supported_tasks_str = ", ".join(VLLM_SUPPORTED_TASKS)
+            raise ValueError(
+                f"Inference engine is only supported for vLLM supported tasks. {supported_tasks_str}. {found_tasks_str}"
+            )
 
-        # Check if the model is logged with OPENAI_CHAT_SIGNATURE
+        # Check if the model is logged with OpenAI compatible signature.
         signatures_dict = model_spec.get("signatures", {})
 
         # Deserialize signatures from model spec to ModelSignature objects for proper semantic comparison.
@@ -1105,11 +1140,16 @@ class ModelVersion(lineage_node.LineageNode):
             func_name: core.ModelSignature.from_dict(sig_dict) for func_name, sig_dict in signatures_dict.items()
         }
 
-        if deserialized_signatures != openai_signatures.OPENAI_CHAT_SIGNATURE:
+        if deserialized_signatures not in [
+            openai_signatures.OPENAI_CHAT_SIGNATURE,
+            openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING,
+        ]:
             raise ValueError(
-                "Inference engine requires the model to be logged with OPENAI_CHAT_SIGNATURE. "
+                "Inference engine requires the model to be logged with openai_signatures.OPENAI_CHAT_SIGNATURE or "
+                "openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING. "
                 f"Found signatures: {signatures_dict}. "
-                "Please log the model with: signatures=openai_signatures.OPENAI_CHAT_SIGNATURE"
+                "Please log the model again with: signatures=openai_signatures.OPENAI_CHAT_SIGNATURE or "
+                "signatures=openai_signatures.OPENAI_CHAT_SIGNATURE_WITH_CONTENT_FORMAT_STRING"
             )
 
     @overload
@@ -1350,20 +1390,11 @@ class ModelVersion(lineage_node.LineageNode):
         # Validate GPU support if GPU resources are requested
         self._throw_error_if_gpu_is_not_supported(gpu_requests, statement_params)
 
-        inference_engine_args = inference_engine_utils._get_inference_engine_args(inference_engine_options)
-
-        # Check if model is HuggingFace text-generation and is logged with
-        # OPENAI_CHAT_SIGNATURE before doing inference engine checks
-        # Only validate if inference engine is actually specified
-        if inference_engine_args is not None:
-            self._check_huggingface_text_generation_model(statement_params)
-
-        # Enrich inference engine args if inference engine is specified
-        if inference_engine_args is not None:
-            inference_engine_args = inference_engine_utils._enrich_inference_engine_args(
-                inference_engine_args,
-                gpu_requests,
-            )
+        inference_engine_args = self._prepare_inference_engine_args(
+            inference_engine_options,
+            gpu_requests,
+            statement_params,
+        )
 
         from snowflake.ml.model import event_handler
         from snowflake.snowpark import exceptions

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

@@ -10,11 +10,10 @@ from typing import Any, Literal, Optional, TypedDict, Union, cast, overload
 import yaml
 from typing_extensions import NotRequired
 
-from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.utils import formatting, identifier, sql_identifier, url
 from snowflake.ml.model import model_signature, type_hints
-from snowflake.ml.model._client.ops import deployment_step, metadata_ops
+from snowflake.ml.model._client.ops import deployment_step, metadata_ops, param_utils
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
@@ -698,9 +697,6 @@ class ModelOperator:
 
         result: list[ServiceInfo] = []
         is_privatelink_connection = self._is_privatelink_connection()
-        is_autocapture_param_enabled = (
-            platform_capabilities.PlatformCapabilities.get_instance().is_inference_autocapture_enabled()
-        )
 
         for fully_qualified_service_name in fully_qualified_service_names:
             port: Optional[int] = None
@@ -742,10 +738,8 @@ class ModelOperator:
                 inference_endpoint=inference_endpoint,
                 internal_endpoint=f"http://{internal_dns}:{port}" if port is not None else None,
             )
-            if is_autocapture_param_enabled and self._service_client.DESC_SERVICE_SPEC_COL_NAME in service_description:
-                # Include column only if parameter is enabled and spec exists for service owner caller
-                autocapture_enabled = self._service_client.get_proxy_container_autocapture(service_description)
-                service_info["autocapture_enabled"] = autocapture_enabled
+            autocapture_enabled = self._service_client.is_autocapture_enabled(service_description)
+            service_info["autocapture_enabled"] = autocapture_enabled
 
             result.append(service_info)
 
@@ -1063,23 +1057,7 @@ class ModelOperator:
                 col_name = sql_identifier.SqlIdentifier(input_feature.name.upper(), case_sensitive=True)
             input_args.append(col_name)
 
-        method_parameters: Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]] = None
-        if signature.params:
-            # Start with defaults from signature
-            final_params = {}
-            for param_spec in signature.params:
-                if hasattr(param_spec, "default_value"):
-                    final_params[param_spec.name] = param_spec.default_value
-
-            # Override with provided runtime parameters
-            if params:
-                final_params.update(params)
-
-            # Convert to list of tuples with SqlIdentifier for parameter names
-            method_parameters = [
-                (sql_identifier.SqlIdentifier(param_name), param_value)
-                for param_name, param_value in final_params.items()
-            ]
+        method_parameters = param_utils.validate_and_resolve_params(params, signature.params)
 
         returns = []
         for output_feature in signature.outputs:

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

@@ -0,0 +1,124 @@
+"""Utility functions for model parameter validation and resolution."""
+
+from typing import Any, Optional, Sequence
+
+from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._signatures import core
+
+
+def validate_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Optional[Sequence[core.BaseParamSpec]],
+) -> None:
+    """Validate user-provided params against signature params.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Raises:
+        SnowflakeMLException: If params are provided but signature has no params,
+            or if unknown params are provided, or if param types are invalid,
+            or if duplicate params are provided with different cases.
+    """
+    # Params provided but signature has no params defined
+    if params and not signature_params:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Parameters were provided ({sorted(params.keys())}), "
+                "but this method does not accept any parameters."
+            ),
+        )
+
+    if not signature_params or not params:
+        return
+
+    # Case-insensitive lookup: normalized_name -> param_spec
+    param_spec_lookup = {ps.name.upper(): ps for ps in signature_params}
+
+    # Check for duplicate params with different cases (e.g., "temperature" and "TEMPERATURE")
+    normalized_names = [name.upper() for name in params]
+    if len(normalized_names) != len(set(normalized_names)):
+        # Find the duplicate params to raise an error
+        param_seen: dict[str, list[str]] = {}
+        for param_name in params:
+            param_seen.setdefault(param_name.upper(), []).append(param_name)
+        duplicate_param_names = [param_names for param_names in param_seen.values() if len(param_names) > 1]
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Duplicate parameter(s) provided with different cases: {duplicate_param_names}. "
+                "Parameter names are case-insensitive."
+            ),
+        )
+
+    # Validate user-provided params exist (case-insensitive)
+    invalid_params = [name for name in params if name.upper() not in param_spec_lookup]
+    if invalid_params:
+        raise exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_ARGUMENT,
+            original_exception=ValueError(
+                f"Unknown parameter(s): {sorted(invalid_params)}. "
+                f"Valid parameters are: {sorted(ps.name for ps in signature_params)}"
+            ),
+        )
+
+    # Validate types for each provided param
+    for param_name, default_value in params.items():
+        param_spec = param_spec_lookup[param_name.upper()]
+        if isinstance(param_spec, core.ParamSpec):
+            core.ParamSpec._validate_default_value(param_spec.dtype, default_value, param_spec.shape)
+
+
+def resolve_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Sequence[core.BaseParamSpec],
+) -> list[tuple[sql_identifier.SqlIdentifier, Any]]:
+    """Resolve final method parameters by applying user-provided params over signature defaults.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Returns:
+        List of tuples (SqlIdentifier, value) for method invocation.
+    """
+    # Case-insensitive lookup: normalized_name -> param_spec
+    param_spec_lookup = {ps.name.upper(): ps for ps in signature_params}
+
+    # Start with defaults from signature
+    final_params: dict[str, Any] = {}
+    for param_spec in signature_params:
+        if hasattr(param_spec, "default_value"):
+            final_params[param_spec.name] = param_spec.default_value
+
+    # Override with provided runtime parameters (using signature's original param names)
+    if params:
+        for param_name, override_value in params.items():
+            canonical_name = param_spec_lookup[param_name.upper()].name
+            final_params[canonical_name] = override_value
+
+    return [(sql_identifier.SqlIdentifier(param_name), param_value) for param_name, param_value in final_params.items()]
+
+
+def validate_and_resolve_params(
+    params: Optional[dict[str, Any]],
+    signature_params: Optional[Sequence[core.BaseParamSpec]],
+) -> Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]]:
+    """Validate user-provided params against signature params and return method parameters.
+
+    Args:
+        params: User-provided parameter dictionary (runtime values).
+        signature_params: Parameter specifications from the model signature.
+
+    Returns:
+        List of tuples (SqlIdentifier, value) for method invocation, or None if no params.
+    """
+    validate_params(params, signature_params)
+
+    if not signature_params:
+        return None
+
+    return resolve_params(params, signature_params)