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

snowflake/ml/model/_client/model/batch_inference_specs.py

@@ -19,11 +19,74 @@ class SaveMode(str, Enum):
 
 
 class OutputSpec(BaseModel):
+    """Specification for batch inference output.
+
+    Defines where the inference results should be written and how to handle
+    existing files at the output location.
+
+    Attributes:
+        stage_location (str): The stage path where batch inference results will be saved.
+            This should be a full path including the stage with @ prefix. For example,
+            '@My_DB.PUBLIC.MY_STAGE/someth/path/'. A non-existent directory will be re-created.
+            Only Snowflake internal stages are supported at this moment.
+        mode (SaveMode): The save mode that determines behavior when files already exist
+            at the output location. Defaults to SaveMode.ERROR which raises an error
+            if files exist. Can be set to SaveMode.OVERWRITE to replace existing files.
+
+    Example:
+        >>> output_spec = OutputSpec(
+        ...     stage_location="@My_DB.PUBLIC.MY_STAGE/someth/path/",
+        ...     mode=SaveMode.OVERWRITE
+        ... )
+    """
+
     stage_location: str
     mode: SaveMode = SaveMode.ERROR
 
 
 class JobSpec(BaseModel):
+    """Specification for batch inference job execution.
+
+    Defines the compute resources, job settings, and execution parameters
+    for running batch inference jobs in Snowflake.
+
+    Attributes:
+        image_repo (Optional[str]): Container image repository for the inference job.
+            If not specified, uses the default repository.
+        job_name (Optional[str]): Custom name for the batch inference job.
+            If not provided, a name will be auto-generated in the form of "BATCH_INFERENCE_<UUID>".
+        num_workers (Optional[int]): The number of workers to run the inference service for handling
+            requests in parallel within an instance of the service. By default, it is set to 2*vCPU+1
+            of the node for CPU based inference and 1 for GPU based inference. For GPU based inference,
+            please see best practices before playing with this value.
+        function_name (Optional[str]): Name of the specific function to call for inference.
+            Required when the model has multiple inference functions.
+        force_rebuild (bool): Whether to force rebuilding the container image even if
+            it already exists. Defaults to False.
+        max_batch_rows (int): Maximum number of rows to process in a single batch.
+            Defaults to 1024. Larger values may improve throughput.
+        warehouse (Optional[str]): Snowflake warehouse to use for the batch inference job.
+            If not specified, uses the session's current warehouse.
+        cpu_requests (Optional[str]): The cpu limit for CPU based inference. Can be an integer,
+            fractional or string values. If None, we attempt to utilize all the vCPU of the node.
+        memory_requests (Optional[str]): The memory limit for inference. Can be an integer
+            or a fractional value, but requires a unit (GiB, MiB). If None, we attempt to utilize all
+            the memory of the node.
+        gpu_requests (Optional[str]): The gpu limit for GPU based inference. Can be integer or
+            string values. Use CPU if None.
+        replicas (Optional[int]): Number of job replicas to run for high availability.
+            If not specified, defaults to 1 replica.
+
+    Example:
+        >>> job_spec = JobSpec(
+        ...     job_name="my_inference_job",
+        ...     num_workers=4,
+        ...     cpu_requests="2",
+        ...     memory_requests="8Gi",
+        ...     max_batch_rows=2048
+        ... )
+    """
+
     image_repo: Optional[str] = None
     job_name: Optional[str] = None
     num_workers: Optional[int] = None

snowflake/ml/model/_client/model/inference_engine_utils.py

@@ -6,13 +6,9 @@ from snowflake.ml.model._client.ops import service_ops
 def _get_inference_engine_args(
     experimental_options: Optional[dict[str, Any]],
 ) -> Optional[service_ops.InferenceEngineArgs]:
-
-    if not experimental_options:
+    if not experimental_options or "inference_engine" not in experimental_options:
         return None
 
-    if "inference_engine" not in experimental_options:
-        raise ValueError("inference_engine is required in experimental_options")
-
     return service_ops.InferenceEngineArgs(
         inference_engine=experimental_options["inference_engine"],
         inference_engine_args_override=experimental_options.get("inference_engine_args_override"),

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

@@ -7,6 +7,7 @@ from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
 
+from snowflake import snowpark
 from snowflake.ml import jobs
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
@@ -593,7 +594,8 @@ class ModelVersion(lineage_node.LineageNode):
             "job_spec",
         ],
     )
-    def _run_batch(
+    @snowpark._internal.utils.private_preview(version="1.18.0")
+    def run_batch(
         self,
         *,
         compute_pool: str,
@@ -601,6 +603,68 @@ class ModelVersion(lineage_node.LineageNode):
         output_spec: batch_inference_specs.OutputSpec,
         job_spec: Optional[batch_inference_specs.JobSpec] = None,
     ) -> jobs.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.
+                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.
+            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.
+
+        Returns:
+            jobs.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
+                lifecycle.
+
+        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.
+
+        Example:
+            >>> # Prepare input data - Example 1: From a table
+            >>> input_df = session.table("my_input_table")
+            >>>
+            >>> # Prepare input data - Example 2: From a SQL query
+            >>> input_df = session.sql(
+            ...     "SELECT id, feature_1, feature_2 FROM feature_table WHERE feature_1 > 100"
+            ... )
+            >>>
+            >>> # Prepare input data - Example 3: From Parquet files in a stage
+            >>> input_df = session.read.option("pattern", ".*\\.parquet").parquet(
+            ...     "@my_stage/input_data/"
+            ... ).select("id", "feature_1", "feature_2")
+            >>>
+            >>> # Configure output location
+            >>> output_spec = OutputSpec(
+            ...     stage_location='@My_DB.PUBLIC.MY_STAGE/someth/path/',
+            ...     mode=SaveMode.OVERWRITE
+            ... )
+            >>>
+            >>> # Configure job parameters
+            >>> job_spec = JobSpec(
+            ...     job_name="my_batch_inference",
+            ...     num_workers=4,
+            ...     cpu_requests="2",
+            ...     memory_requests="8Gi"
+            ... )
+            >>>
+            >>> # Run batch inference
+            >>> job = model_version.run_batch(
+            ...     compute_pool="my_compute_pool",
+            ...     input_spec=input_df,
+            ...     output_spec=output_spec,
+            ...     job_spec=job_spec
+            ... )
+
+        Note:
+            This method is currently in private preview and requires Snowflake version 1.18.0 or later.
+            The input data is temporarily stored in the output stage location under /_temporary before
+            inference execution.
+        """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
@@ -827,6 +891,51 @@ class ModelVersion(lineage_node.LineageNode):
             version_name=sql_identifier.SqlIdentifier(version),
         )
 
+    def _can_run_on_gpu(
+        self,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> bool:
+        """Check if the model has GPU runtime support.
+
+        Args:
+            statement_params: Optional dictionary of statement parameters to include
+                in the SQL command to fetch model spec.
+
+        Returns:
+            True if the model has GPU runtime configured, False otherwise.
+        """
+        # Fetch model spec
+        model_spec = self._get_model_spec(statement_params)
+
+        # Check if runtimes section exists and has gpu runtime
+        runtimes = model_spec.get("runtimes", {})
+        return "gpu" in runtimes
+
+    def _throw_error_if_gpu_is_not_supported(
+        self,
+        gpu_requests: Optional[Union[str, int]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """Check if the model has GPU runtime support.
+
+        Args:
+            gpu_requests: The gpu limit for GPU based inference. Can be integer, fractional or string values. Use CPU
+                if None.
+            statement_params: Optional dictionary of statement parameters to include
+                in the SQL command to fetch model spec.
+
+        Raises:
+            ValueError: If the model does not have GPU runtime support.
+        """
+        if gpu_requests is not None and not self._can_run_on_gpu(statement_params):
+            raise ValueError(
+                f"GPU resources requested (gpu_requests={gpu_requests}), but the model "
+                f"{self.fully_qualified_model_name} version {self.version_name} does not have GPU runtime support. "
+                "Please ensure the model was logged with GPU runtime configuration or do not provide gpu_requests. "
+                "To log the model with GPU runtime configuration, provide `cuda_version` in the `options` while calling"
+                " the `log_model` function."
+            )
+
     def _check_huggingface_text_generation_model(
         self,
         statement_params: Optional[dict[str, Any]] = None,
@@ -926,9 +1035,10 @@ class ModelVersion(lineage_node.LineageNode):
                 When it is ``False``, this function executes the underlying service creation asynchronously
                 and returns an :class:`AsyncJob`.
             experimental_options: Experimental options for the service creation with custom inference engine.
-                Currently, only `inference_engine` and `inference_engine_args_override` are supported.
+                Currently, `inference_engine`, `inference_engine_args_override`, and `autocapture` are supported.
                 `inference_engine` is the name of the inference engine to use.
                 `inference_engine_args_override` is a list of string arguments to pass to the inference engine.
+                `autocapture` is a boolean to enable/disable inference table.
         """
         ...
 
@@ -984,9 +1094,10 @@ class ModelVersion(lineage_node.LineageNode):
                 When it is ``False``, this function executes the underlying service creation asynchronously
                 and returns an :class:`AsyncJob`.
             experimental_options: Experimental options for the service creation with custom inference engine.
-                Currently, only `inference_engine` and `inference_engine_args_override` are supported.
+                Currently, `inference_engine`, `inference_engine_args_override`, and `autocapture` are supported.
                 `inference_engine` is the name of the inference engine to use.
                 `inference_engine_args_override` is a list of string arguments to pass to the inference engine.
+                `autocapture` is a boolean to enable/disable inference table.
         """
         ...
 
@@ -1059,21 +1170,20 @@ class ModelVersion(lineage_node.LineageNode):
                 When it is False, this function executes the underlying service creation asynchronously
                 and returns an AsyncJob.
             experimental_options: Experimental options for the service creation with custom inference engine.
-                Currently, only `inference_engine` and `inference_engine_args_override` are supported.
+                Currently, `inference_engine`, `inference_engine_args_override`, and `autocapture` are supported.
                 `inference_engine` is the name of the inference engine to use.
                 `inference_engine_args_override` is a list of string arguments to pass to the inference engine.
+                `autocapture` is a boolean to enable/disable inference table.
 
 
         Raises:
-            ValueError: Illegal external access integration arguments.
+            ValueError: Illegal external access integration arguments, or if GPU resources are requested
+                but the model does not have GPU runtime support.
             exceptions.SnowparkSQLException: if service already exists.
 
         Returns:
             If `block=True`, return result information about service creation from server.
             Otherwise, return the service creation AsyncJob.
-
-        Raises:
-            ValueError: Illegal external access integration arguments.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
@@ -1096,12 +1206,16 @@ class ModelVersion(lineage_node.LineageNode):
 
         service_db_id, service_schema_id, service_id = sql_identifier.parse_fully_qualified_name(service_name)
 
-        # Check if model is HuggingFace text-generation before doing inference engine checks
-        if experimental_options:
-            self._check_huggingface_text_generation_model(statement_params)
+        # 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(experimental_options)
 
+        # Check if model is HuggingFace text-generation 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(
@@ -1109,6 +1223,9 @@ class ModelVersion(lineage_node.LineageNode):
                 gpu_requests,
             )
 
+        # Extract autocapture from experimental_options
+        autocapture = experimental_options.get("autocapture") if experimental_options else None
+
         from snowflake.ml.model import event_handler
         from snowflake.snowpark import exceptions
 
@@ -1148,6 +1265,7 @@ class ModelVersion(lineage_node.LineageNode):
                     statement_params=statement_params,
                     progress_status=status,
                     inference_engine_args=inference_engine_args,
+                    autocapture=autocapture,
                 )
                 status.update(label="Model service created successfully", state="complete", expanded=False)
                 return result

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

@@ -206,6 +206,8 @@ class ServiceOperator:
         hf_model_args: Optional[HFModelArgs] = None,
         # inference engine model
         inference_engine_args: Optional[InferenceEngineArgs] = None,
+        # inference table
+        autocapture: Optional[bool] = None,
     ) -> Union[str, async_job.AsyncJob]:
 
         # Generate operation ID for this deployment
@@ -261,6 +263,7 @@ class ServiceOperator:
             gpu=gpu_requests,
             num_workers=num_workers,
             max_batch_rows=max_batch_rows,
+            autocapture=autocapture,
         )
         if hf_model_args:
             # hf model

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

@@ -146,6 +146,7 @@ class ModelDeploymentSpec:
         gpu: Optional[Union[str, int]] = None,
         num_workers: Optional[int] = None,
         max_batch_rows: Optional[int] = None,
+        autocapture: Optional[bool] = None,
     ) -> "ModelDeploymentSpec":
         """Add service specification to the deployment spec.
 
@@ -161,6 +162,7 @@ class ModelDeploymentSpec:
             gpu: GPU requirement.
             num_workers: Number of workers.
             max_batch_rows: Maximum batch rows for inference.
+            autocapture: Whether to enable inference table.
 
         Raises:
             ValueError: If a job spec already exists.
@@ -186,6 +188,7 @@ class ModelDeploymentSpec:
             compute_pool=inference_compute_pool_name.identifier(),
             ingress_enabled=ingress_enabled,
             max_instances=max_instances,
+            autocapture=autocapture,
             **self._inference_spec,
         )
         return self

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

@@ -32,6 +32,7 @@ class Service(BaseModel):
     gpu: Optional[str] = None
     num_workers: Optional[int] = None
     max_batch_rows: Optional[int] = None
+    autocapture: Optional[bool] = None
     inference_engine_spec: Optional[InferenceEngineSpec] = None
 
 

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -87,7 +87,9 @@ class ModelManifest:
                     model_meta_schema.FunctionProperties.PARTITIONED.value, False
                 ),
                 wide_input=len(model_meta.signatures[target_method].inputs) > constants.SNOWPARK_UDF_INPUT_COL_LIMIT,
-                options=model_method.get_model_method_options_from_options(options, target_method),
+                options=model_method.get_model_method_options_from_options(
+                    options, target_method, model_meta.model_type
+                ),
             )
 
             self.methods.append(method)

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -32,7 +32,7 @@ class ModelMethodOptions(TypedDict):
 
 
 def get_model_method_options_from_options(
-    options: type_hints.ModelSaveOption, target_method: str
+    options: type_hints.ModelSaveOption, target_method: str, model_type: Optional[str] = None
 ) -> ModelMethodOptions:
     default_function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
     method_option = options.get("method_options", {}).get(target_method, {})
@@ -42,6 +42,9 @@ def get_model_method_options_from_options(
         case_sensitive = utils.determine_explain_case_sensitive_from_method_options(
             options.get("method_options", {}), target_method
         )
+    elif model_type == "prophet":
+        # Prophet models always require TABLE_FUNCTION because they need entire time series context
+        default_function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
     global_function_type = options.get("function_type", default_function_type)
     function_type = method_option.get("function_type", global_function_type)
     if function_type not in [function_type.value for function_type in model_manifest_schema.ModelMethodFunctionTypes]:

snowflake/ml/model/_packager/model_handlers/_utils.py

@@ -305,3 +305,73 @@ def get_default_cuda_version() -> str:
 
         return torch.version.cuda or model_env.DEFAULT_CUDA_VERSION
     return model_env.DEFAULT_CUDA_VERSION
+
+
+def normalize_column_name(column_name: str) -> str:
+    """Normalize a column name to be a valid unquoted Snowflake SQL identifier.
+
+    Converts column names with spaces and special characters (e.g., "Christmas Day")
+    into valid lowercase unquoted SQL identifiers (e.g., "christmas_day") that can be used
+    without quotes in SQL queries. This follows Snowflake's unquoted identifier rules:
+    https://docs.snowflake.com/en/sql-reference/identifiers-syntax
+
+    The normalization approach is preferred over quoted identifiers because:
+    - Unquoted identifiers are simpler and more readable
+    - They don't require special handling in SQL contexts
+    - They avoid case-sensitivity complications
+    - Lowercase convention improves readability and follows Python/pandas conventions
+
+    This utility is useful for model handlers that need to ensure output column names
+    from models (e.g., Prophet holiday columns, feature names) are SQL-safe.
+
+    Args:
+        column_name: Original column name (may contain spaces, special chars, etc.)
+
+    Returns:
+        Normalized lowercase column name that is a valid unquoted SQL identifier matching
+        the pattern [a-z_][a-z0-9_$]*
+
+    Examples:
+        >>> normalize_column_name_for_snowflake("Christmas Day")
+        'christmas_day'
+        >>> normalize_column_name_for_snowflake("New Year's Day")
+        'new_year_s_day'
+        >>> normalize_column_name_for_snowflake("__private")
+        '__private'
+        >>> normalize_column_name_for_snowflake("2023_data")
+        '_2023_data'
+    """
+    import re
+
+    # Convert to lowercase for readability and consistency
+    normalized = column_name.lower()
+
+    # Replace spaces and special characters with underscores
+    # Keep only alphanumeric characters, underscores, and dollar signs (valid in unquoted identifiers)
+    normalized = re.sub(r"[^a-z0-9_$]", "_", normalized)
+
+    # Collapse consecutive underscores while preserving leading underscores
+    # This handles cases like "A  B" → "a__b" → "a_b" while preserving "__name" → "__name"
+    if len(normalized) > 1:
+        # Count and preserve leading underscores, collapse the rest
+        leading_underscores = len(normalized) - len(normalized.lstrip("_"))
+        rest_of_string = normalized[leading_underscores:]
+        rest_collapsed = re.sub(r"_+", "_", rest_of_string)
+        normalized = "_" * leading_underscores + rest_collapsed
+
+    if normalized == "_":
+        return "_"
+
+    normalized = normalized.rstrip("_")
+
+    # Ensure it starts with a letter or underscore (SQL requirement)
+    # Unquoted identifiers must match: [a-z_][a-z0-9_$]*
+    if normalized and normalized[0].isdigit():
+        normalized = "_" + normalized
+
+    # If normalization resulted in empty string, use a default
+    # (This happens when input was only underscores like "___")
+    if not normalized:
+        normalized = "column"
+
+    return normalized