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

snowflake/ml/jobs/_utils/payload_utils.py

@@ -488,10 +488,13 @@ class JobPayload:
                 " comment = 'Created by snowflake.ml.jobs Python API'",
                 params=[stage_name],
             )
-
+        payload_name = None
         # Upload payload to stage - organize into app/ subdirectory
         app_stage_path = stage_path.joinpath(constants.APP_STAGE_SUBPATH)
         if not isinstance(source, types.PayloadPath):
+            if isinstance(source, function_payload_utils.FunctionPayload):
+                payload_name = source.function.__name__
+
             source_code = generate_python_code(source, source_code_display=True)
             _ = session.file.put_stream(
                 io.BytesIO(source_code.encode()),
@@ -502,12 +505,14 @@ class JobPayload:
             source = Path(entrypoint.file_path.parent)
 
         elif isinstance(source, stage_utils.StagePath):
+            payload_name = entrypoint.file_path.stem
             # copy payload to stage
             if source == entrypoint.file_path:
                 source = source.parent
             upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
 
         elif isinstance(source, Path):
+            payload_name = entrypoint.file_path.stem
             upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
             if source.is_file():
                 source = source.parent
@@ -562,6 +567,7 @@ class JobPayload:
                 *python_entrypoint,
             ],
             env_vars=env_vars,
+            payload_name=payload_name,
         )
 
 

snowflake/ml/jobs/_utils/stage_utils.py

@@ -32,6 +32,10 @@ class StagePath:
         self._root = self._raw_path[0:start].rstrip("/") if relpath else self._raw_path.rstrip("/")
         self._path = Path(relpath or "")
 
+    @property
+    def stem(self) -> str:
+        return self._path.stem
+
     @property
     def parts(self) -> tuple[str, ...]:
         return self._path.parts

snowflake/ml/jobs/_utils/types.py

@@ -23,6 +23,10 @@ class PayloadPath(Protocol):
     def name(self) -> str:
         ...
 
+    @property
+    def stem(self) -> str:
+        ...
+
     @property
     def suffix(self) -> str:
         ...
@@ -92,6 +96,7 @@ class UploadedPayload:
     stage_path: PurePath
     entrypoint: list[Union[str, PurePath]]
     env_vars: dict[str, str] = field(default_factory=dict)
+    payload_name: Optional[str] = None
 
 
 @dataclass(frozen=True)

snowflake/ml/jobs/job.py

@@ -109,18 +109,18 @@ class MLJob(Generic[T], SerializableSessionMixin):
         return cast(dict[str, Any], container_spec)
 
     @property
-    def _stage_path(self) -> str:
+    def _stage_path(self) -> Optional[str]:
         """Get the job's artifact storage stage location."""
         volumes = self._service_spec["spec"]["volumes"]
-        stage_path = next(v for v in volumes if v["name"] == constants.STAGE_VOLUME_NAME)["source"]
-        return cast(str, stage_path)
+        stage_volume = next((v for v in volumes if v["name"] == constants.STAGE_VOLUME_NAME), None)
+        return cast(str, stage_volume["source"]) if stage_volume else None
 
     @property
     def _result_path(self) -> str:
         """Get the job's result file location."""
         result_path_str = self._container_spec["env"].get(constants.RESULT_PATH_ENV_VAR)
         if result_path_str is None:
-            raise RuntimeError(f"Job {self.name} doesn't have a result path configured")
+            raise NotImplementedError(f"Job {self.name} doesn't have a result path configured")
 
         return self._transform_path(result_path_str)
 
@@ -229,8 +229,22 @@ class MLJob(Generic[T], SerializableSessionMixin):
         Raises:
             TimeoutError: If the job does not complete within the specified timeout.
         """
-        delay = constants.JOB_POLL_INITIAL_DELAY_SECONDS  # Start with 100ms delay
         start_time = time.monotonic()
+        try:
+            # spcs_wait_for() is a synchronous query, it’s more effective to do polling with exponential
+            # backoff. If the job is running for a long time. We want a hybrid option: use spcs_wait_for()
+            # for the first 30 seconds, then switch to polling for long running jobs
+            min_timeout = (
+                int(min(timeout, constants.JOB_SPCS_TIMEOUT_SECONDS))
+                if timeout >= 0
+                else constants.JOB_SPCS_TIMEOUT_SECONDS
+            )
+            query_helper.run_query(self._session, f"call {self.id}!spcs_wait_for('DONE', {min_timeout})")
+            return self.status
+        except SnowparkSQLException:
+            # if the function does not support for this environment
+            pass
+        delay: float = float(constants.JOB_POLL_INITIAL_DELAY_SECONDS)  # Start with 5s delay
         warning_shown = False
         while (status := self.status) not in TERMINAL_JOB_STATUSES:
             elapsed = time.monotonic() - start_time

snowflake/ml/jobs/manager.py

@@ -192,12 +192,12 @@ def delete_job(job: Union[str, jb.MLJob[Any]], session: Optional[snowpark.Sessio
     """Delete a job service from the backend. Status and logs will be lost."""
     job = job if isinstance(job, jb.MLJob) else get_job(job, session=session)
     session = job._session
-    try:
-        stage_path = job._stage_path
-        session.sql(f"REMOVE {stage_path}/").collect()
-        logger.info(f"Successfully cleaned up stage files for job {job.id} at {stage_path}")
-    except Exception as e:
-        logger.warning(f"Failed to clean up stage files for job {job.id}: {e}")
+    if job._stage_path:
+        try:
+            session.sql(f"REMOVE {job._stage_path}/").collect()
+            logger.debug(f"Successfully cleaned up stage files for job {job.id} at {job._stage_path}")
+        except Exception as e:
+            logger.warning(f"Failed to clean up stage files for job {job.id}: {e}")
     query_helper.run_query(session, "DROP SERVICE IDENTIFIER(?)", params=(job.id,))
 
 
@@ -697,10 +697,21 @@ def _do_submit_job_v2(
         "MIN_INSTANCES": min_instances,
         "ASYNC": use_async,
     }
+    if payload.payload_name:
+        job_options["GENERATE_SUFFIX"] = True
     job_options = {k: v for k, v in job_options.items() if v is not None}
 
     query_template = "CALL SYSTEM$EXECUTE_ML_JOB(?, ?, ?, ?)"
-    params = [job_id, compute_pool, json.dumps(spec_options), json.dumps(job_options)]
+    if job_id:
+        database, schema, _ = identifier.parse_schema_level_object_identifier(job_id)
+    params = [
+        job_id
+        if payload.payload_name is None
+        else identifier.get_schema_level_object_identifier(database, schema, payload.payload_name) + "_",
+        compute_pool,
+        json.dumps(spec_options),
+        json.dumps(job_options),
+    ]
     actual_job_id = query_helper.run_query(session, query_template, params=params)[0][0]
 
     return get_job(actual_job_id, session=session)

snowflake/ml/model/__init__.py

@@ -1,3 +1,6 @@
+import sys
+import warnings
+
 from snowflake.ml.model._client.model.batch_inference_specs import (
     JobSpec,
     OutputSpec,
@@ -18,3 +21,19 @@ __all__ = [
     "SaveMode",
     "Volatility",
 ]
+
+_deprecation_warning_msg_for_3_9 = (
+    "Python 3.9 is deprecated in snowflake-ml-python. " "Please upgrade to Python 3.10 or greater."
+)
+
+warnings.filterwarnings(
+    "once",
+    message=_deprecation_warning_msg_for_3_9,
+)
+
+if sys.version_info.major == 3 and sys.version_info.minor == 9:
+    warnings.warn(
+        _deprecation_warning_msg_for_3_9,
+        category=DeprecationWarning,
+        stacklevel=2,
+    )

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/model_ops.py

@@ -515,10 +515,17 @@ class ModelOperator:
             statement_params=statement_params,
         )
         for r in res:
-            if alias_name in r[self._model_client.MODEL_VERSION_ALIASES_COL_NAME]:
-                return sql_identifier.SqlIdentifier(
-                    r[self._model_client.MODEL_VERSION_NAME_COL_NAME], case_sensitive=True
-                )
+            aliases_data = r[self._model_client.MODEL_VERSION_ALIASES_COL_NAME]
+            if aliases_data:
+                aliases_list = json.loads(aliases_data)
+
+                # Compare using Snowflake identifier semantics for exact match
+                for alias in aliases_list:
+                    if sql_identifier.SqlIdentifier(alias) == alias_name:
+                        return sql_identifier.SqlIdentifier(
+                            r[self._model_client.MODEL_VERSION_NAME_COL_NAME], case_sensitive=True
+                        )
+
         return None
 
     def get_tag_value(

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]: