snowflake-ml-python 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

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

@@ -1,5 +1,6 @@
 import enum
 import json
+import logging
 import os
 import pathlib
 import tempfile
@@ -11,9 +12,9 @@ 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
+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 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,
@@ -33,6 +34,8 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
+logger = logging.getLogger(__name__)
+
 
 # An enum class to represent Create Or Alter Model SQL command.
 class ModelAction(enum.Enum):
@@ -986,6 +989,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -1002,6 +1006,7 @@ class ModelOperator:
         strict_input_validation: bool = False,
         statement_params: Optional[dict[str, str]] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -1022,6 +1027,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
@@ -1057,6 +1063,8 @@ class ModelOperator:
                 col_name = sql_identifier.SqlIdentifier(input_feature.name.upper(), case_sensitive=True)
             input_args.append(col_name)
 
+        method_parameters = param_utils.validate_and_resolve_params(params, signature.params)
+
         returns = []
         for output_feature in signature.outputs:
             output_name = identifier_rule.get_sql_identifier_from_feature(output_feature.name)
@@ -1075,6 +1083,7 @@ class ModelOperator:
                 schema_name=schema_name,
                 service_name=service_name,
                 statement_params=statement_params,
+                params=method_parameters,
             )
         else:
             assert model_name is not None
@@ -1090,6 +1099,7 @@ class ModelOperator:
                     model_name=model_name,
                     version_name=version_name,
                     statement_params=statement_params,
+                    params=method_parameters,
                 )
             elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
                 df_res = self._model_version_client.invoke_table_function_method(
@@ -1105,6 +1115,7 @@ class ModelOperator:
                     statement_params=statement_params,
                     is_partitioned=is_partitioned or False,
                     explain_case_sensitive=explain_case_sensitive,
+                    params=method_parameters,
                 )
 
         if keep_order:
@@ -1238,3 +1249,35 @@ class ModelOperator:
                     target_path=local_file_dir,
                     statement_params=statement_params,
                 )
+
+    def run_import_model_query(
+        self,
+        *,
+        database_name: str,
+        schema_name: str,
+        yaml_content: str,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        yaml_content_escaped = snowpark_utils.escape_single_quotes(yaml_content)  # type: ignore[no-untyped-call]
+
+        async_job = self._session.sql(
+            f"SELECT SYSTEM$IMPORT_MODEL('{yaml_content_escaped}')",
+        ).collect(block=False, statement_params=statement_params)
+        query_id = async_job.query_id  # type: ignore[attr-defined]
+
+        logger.info(f"Remotely importing model, with the query id: {query_id}")
+        model_logger_service_name = sql_identifier.SqlIdentifier(
+            deployment_step.get_service_id_from_deployment_step(
+                query_id,
+                deployment_step.DeploymentStep.MODEL_LOGGING,
+            )
+        )
+
+        logger_name = model_logger_service_name.identifier()
+        job_url = f"{url.JOB_URL_PREFIX}/{database_name}/{schema_name}/{logger_name}"
+        snowflake_url = url.get_snowflake_url(session=self._session, url_path=job_url)
+        logger.info(
+            f"To monitor the progress of the model logging job, head to the job monitoring page {snowflake_url}"
+        )
+
+        async_job.result()  # type: ignore[attr-defined]

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)

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

@@ -1,6 +1,6 @@
+import base64
 import dataclasses
-import enum
-import hashlib
+import json
 import logging
 import pathlib
 import re
@@ -8,7 +8,9 @@ import tempfile
 import threading
 import time
 import warnings
-from typing import Any, Optional, Union, cast
+from typing import Any, Optional, Sequence, Union, cast
+
+from pydantic import TypeAdapter
 
 from snowflake import snowpark
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
@@ -16,8 +18,10 @@ from snowflake.ml._internal.utils import identifier, service_logger, sql_identif
 from snowflake.ml.jobs import job
 from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
 from snowflake.ml.model._client.model import batch_inference_specs
+from snowflake.ml.model._client.ops import deployment_step, param_utils
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
+from snowflake.ml.model._signatures import core
 from snowflake.snowpark import async_job, exceptions, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -25,32 +29,12 @@ module_logger = service_logger.get_logger(__name__, service_logger.LogColor.GREY
 module_logger.propagate = False
 
 
-class DeploymentStep(enum.Enum):
-    MODEL_BUILD = ("model-build", "model_build_")
-    MODEL_INFERENCE = ("model-inference", None)
-    MODEL_LOGGING = ("model-logging", "model_logging_")
-
-    def __init__(self, container_name: str, service_name_prefix: Optional[str]) -> None:
-        self._container_name = container_name
-        self._service_name_prefix = service_name_prefix
-
-    @property
-    def container_name(self) -> str:
-        """Get the container name for the deployment step."""
-        return self._container_name
-
-    @property
-    def service_name_prefix(self) -> Optional[str]:
-        """Get the service name prefix for the deployment step."""
-        return self._service_name_prefix
-
-
 @dataclasses.dataclass
 class ServiceLogInfo:
     database_name: Optional[sql_identifier.SqlIdentifier]
     schema_name: Optional[sql_identifier.SqlIdentifier]
     service_name: sql_identifier.SqlIdentifier
-    deployment_step: DeploymentStep
+    deployment_step: deployment_step.DeploymentStep
     instance_id: str = "0"
     log_color: service_logger.LogColor = service_logger.LogColor.GREY
 
@@ -353,13 +337,16 @@ class ServiceOperator:
         if is_enable_image_build:
             # stream service logs in a thread
             model_build_service_name = sql_identifier.SqlIdentifier(
-                self._get_service_id_from_deployment_step(query_id, DeploymentStep.MODEL_BUILD)
+                deployment_step.get_service_id_from_deployment_step(
+                    query_id,
+                    deployment_step.DeploymentStep.MODEL_BUILD,
+                )
             )
             model_build_service = ServiceLogInfo(
                 database_name=service_database_name,
                 schema_name=service_schema_name,
                 service_name=model_build_service_name,
-                deployment_step=DeploymentStep.MODEL_BUILD,
+                deployment_step=deployment_step.DeploymentStep.MODEL_BUILD,
                 log_color=service_logger.LogColor.GREEN,
             )
 
@@ -367,21 +354,23 @@ class ServiceOperator:
             database_name=service_database_name,
             schema_name=service_schema_name,
             service_name=service_name,
-            deployment_step=DeploymentStep.MODEL_INFERENCE,
+            deployment_step=deployment_step.DeploymentStep.MODEL_INFERENCE,
             log_color=service_logger.LogColor.BLUE,
         )
 
         model_logger_service: Optional[ServiceLogInfo] = None
         if hf_model_args:
             model_logger_service_name = sql_identifier.SqlIdentifier(
-                self._get_service_id_from_deployment_step(query_id, DeploymentStep.MODEL_LOGGING)
+                deployment_step.get_service_id_from_deployment_step(
+                    query_id, deployment_step.DeploymentStep.MODEL_LOGGING
+                )
             )
 
             model_logger_service = ServiceLogInfo(
                 database_name=service_database_name,
                 schema_name=service_schema_name,
                 service_name=model_logger_service_name,
-                deployment_step=DeploymentStep.MODEL_LOGGING,
+                deployment_step=deployment_step.DeploymentStep.MODEL_LOGGING,
                 log_color=service_logger.LogColor.ORANGE,
             )
 
@@ -536,7 +525,7 @@ class ServiceOperator:
         service = service_log_meta.service
         # check if using an existing model build image
         if (
-            service.deployment_step == DeploymentStep.MODEL_BUILD
+            service.deployment_step == deployment_step.DeploymentStep.MODEL_BUILD
             and not force_rebuild
             and service_log_meta.is_model_logger_service_done
             and not service_log_meta.is_model_build_service_done
@@ -582,31 +571,26 @@ class ServiceOperator:
         if (service_status != service_sql.ServiceStatus.RUNNING) or (service_status != service_log_meta.service_status):
             service_log_meta.service_status = service_status
 
-            if service.deployment_step == DeploymentStep.MODEL_BUILD:
+            if service.deployment_step == deployment_step.DeploymentStep.MODEL_BUILD:
                 module_logger.info(
                     f"Image build service {service.display_service_name} is "
                     f"{service_log_meta.service_status.value}."
                 )
-            elif service.deployment_step == DeploymentStep.MODEL_INFERENCE:
+            elif service.deployment_step == deployment_step.DeploymentStep.MODEL_INFERENCE:
                 module_logger.info(
                     f"Inference service {service.display_service_name} is {service_log_meta.service_status.value}."
                 )
-            elif service.deployment_step == DeploymentStep.MODEL_LOGGING:
+            elif service.deployment_step == deployment_step.DeploymentStep.MODEL_LOGGING:
                 module_logger.info(
                     f"Model logger service {service.display_service_name} is "
                     f"{service_log_meta.service_status.value}."
                 )
             for status in statuses:
                 if status.instance_id is not None:
-                    instance_status, container_status = None, None
-                    if status.instance_status is not None:
-                        instance_status = status.instance_status.value
-                    if status.container_status is not None:
-                        container_status = status.container_status.value
                     module_logger.info(
                         f"Instance[{status.instance_id}]: "
-                        f"instance status: {instance_status}, "
-                        f"container status: {container_status}, "
+                        f"instance status: {status.instance_status}, "
+                        f"container status: {status.container_status}, "
                         f"message: {status.message}"
                     )
             time.sleep(5)
@@ -627,7 +611,7 @@ class ServiceOperator:
         if service_status == service_sql.ServiceStatus.DONE:
             # check if model logger service is done
             # and transition the service log metadata to the model image build service
-            if service.deployment_step == DeploymentStep.MODEL_LOGGING:
+            if service.deployment_step == deployment_step.DeploymentStep.MODEL_LOGGING:
                 if model_build_service:
                     # building the inference image, transition to the model build service
                     service_log_meta.transition_service_log_metadata(
@@ -648,7 +632,7 @@ class ServiceOperator:
                     )
             # check if model build service is done
             # and transition the service log metadata to the model inference service
-            elif service.deployment_step == DeploymentStep.MODEL_BUILD:
+            elif service.deployment_step == deployment_step.DeploymentStep.MODEL_BUILD:
                 service_log_meta.transition_service_log_metadata(
                     model_inference_service,
                     f"Image build service {service.display_service_name} complete.",
@@ -656,7 +640,7 @@ class ServiceOperator:
                     is_model_logger_service_done=service_log_meta.is_model_logger_service_done,
                     operation_id=operation_id,
                 )
-            elif service.deployment_step == DeploymentStep.MODEL_INFERENCE:
+            elif service.deployment_step == deployment_step.DeploymentStep.MODEL_INFERENCE:
                 module_logger.info(f"Inference service {service.display_service_name} is deployed.")
             else:
                 module_logger.warning(f"Service {service.display_service_name} is done, but not transitioning.")
@@ -916,19 +900,6 @@ class ServiceOperator:
 
             time.sleep(2)  # Poll every 2 seconds
 
-    @staticmethod
-    def _get_service_id_from_deployment_step(query_id: str, deployment_step: DeploymentStep) -> str:
-        """Get the service ID through the server-side logic."""
-        uuid = query_id.replace("-", "")
-        big_int = int(uuid, 16)
-        md5_hash = hashlib.md5(str(big_int).encode()).hexdigest()
-        identifier = md5_hash[:8]
-        service_name_prefix = deployment_step.service_name_prefix
-        if service_name_prefix is None:
-            # raise an exception if the service name prefix is None
-            raise ValueError(f"Service name prefix is {service_name_prefix} for deployment step {deployment_step}.")
-        return (service_name_prefix + identifier).upper()
-
     def _check_if_service_exists(
         self,
         database_name: Optional[sql_identifier.SqlIdentifier],
@@ -959,6 +930,38 @@ class ServiceOperator:
         except exceptions.SnowparkSQLException:
             return False
 
+    @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
+        return base64.b64encode(json.dumps(params).encode("utf-8")).decode("utf-8")
+
+    @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.
+        """
+        if column_handling is None:
+            return None
+        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
+        validated_input = adapter.validate_python(column_handling)
+        return base64.b64encode(adapter.dump_json(validated_input)).decode("utf-8")
+
     def invoke_batch_job_method(
         self,
         *,
@@ -971,6 +974,9 @@ class ServiceOperator:
         image_repo_name: Optional[str],
         input_stage_location: str,
         input_file_pattern: str,
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
+        params: Optional[dict[str, Any]],
+        signature_params: Optional[Sequence[core.BaseParamSpec]],
         output_stage_location: str,
         completion_filename: str,
         force_rebuild: bool,
@@ -981,7 +987,13 @@ class ServiceOperator:
         gpu_requests: Optional[str],
         replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
+        inference_engine_args: Optional[InferenceEngineArgs] = None,
     ) -> job.MLJob[Any]:
+        # Validate and encode params
+        param_utils.validate_params(params, signature_params)
+        params_encoded = self._encode_params(params)
+        column_handling_encoded = self._encode_column_handling(column_handling)
+
         database_name = self._database_name
         schema_name = self._schema_name
 
@@ -1007,6 +1019,8 @@ class ServiceOperator:
             max_batch_rows=max_batch_rows,
             input_stage_location=input_stage_location,
             input_file_pattern=input_file_pattern,
+            column_handling=column_handling_encoded,
+            params=params_encoded,
             output_stage_location=output_stage_location,
             completion_filename=completion_filename,
             function_name=function_name,
@@ -1017,11 +1031,17 @@ class ServiceOperator:
             replicas=replicas,
         )
 
-        self._model_deployment_spec.add_image_build_spec(
-            image_build_compute_pool_name=compute_pool_name,
-            fully_qualified_image_repo_name=self._get_image_repo_fqn(image_repo_name, database_name, schema_name),
-            force_rebuild=force_rebuild,
-        )
+        if inference_engine_args:
+            self._model_deployment_spec.add_inference_engine_spec(
+                inference_engine=inference_engine_args.inference_engine,
+                inference_engine_args=inference_engine_args.inference_engine_args_override,
+            )
+        else:
+            self._model_deployment_spec.add_image_build_spec(
+                image_build_compute_pool_name=compute_pool_name,
+                fully_qualified_image_repo_name=self._get_image_repo_fqn(image_repo_name, database_name, schema_name),
+                force_rebuild=force_rebuild,
+            )
 
         spec_yaml_str_or_path = self._model_deployment_spec.save()
 

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

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from snowflake.ml.model._client.service import model_deployment_spec_schema
+
+BaseModel.model_config["protected_namespaces"] = ()
+
+
+class ModelName(BaseModel):
+    model_name: str
+    version_name: str
+
+
+class ModelSpec(BaseModel):
+    name: ModelName
+    hf_model: Optional[model_deployment_spec_schema.HuggingFaceModel] = None
+    log_model_args: Optional[model_deployment_spec_schema.LogModelArgs] = None
+
+
+class ImportModelSpec(BaseModel):
+    compute_pool: str
+    models: list[ModelSpec]

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

@@ -195,6 +195,7 @@ class ModelDeploymentSpec:
 
     def add_job_spec(
         self,
+        *,
         job_name: sql_identifier.SqlIdentifier,
         inference_compute_pool_name: sql_identifier.SqlIdentifier,
         function_name: str,
@@ -202,6 +203,8 @@ class ModelDeploymentSpec:
         output_stage_location: str,
         completion_filename: str,
         input_file_pattern: str,
+        column_handling: Optional[str] = None,
+        params: Optional[str] = None,
         warehouse: sql_identifier.SqlIdentifier,
         job_database_name: Optional[sql_identifier.SqlIdentifier] = None,
         job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
@@ -217,14 +220,16 @@ class ModelDeploymentSpec:
         Args:
             job_name: Name of the job.
             inference_compute_pool_name: Compute pool for inference.
-            warehouse: Warehouse for the job.
             function_name: Function name.
             input_stage_location: Stage location for input data.
             output_stage_location: Stage location for output data.
+            completion_filename: Name of completion file (default: "completion.txt").
+            input_file_pattern: Pattern for input files (optional).
+            column_handling: Column handling mode for input data.
+            params: Additional parameters for the job.
+            warehouse: Warehouse for the job.
             job_database_name: Database name for the job.
             job_schema_name: Schema name for the job.
-            input_file_pattern: Pattern for input files (optional).
-            completion_filename: Name of completion file (default: "completion.txt").
             cpu: CPU requirement.
             memory: Memory requirement.
             gpu: GPU requirement.
@@ -259,7 +264,10 @@ class ModelDeploymentSpec:
             warehouse=warehouse.identifier() if warehouse else None,
             function_name=function_name,
             input=model_deployment_spec_schema.Input(
-                input_stage_location=input_stage_location, input_file_pattern=input_file_pattern
+                input_stage_location=input_stage_location,
+                input_file_pattern=input_file_pattern,
+                column_handling=column_handling,
+                params=params,
             ),
             output=model_deployment_spec_schema.Output(
                 output_stage_location=output_stage_location,
@@ -355,7 +363,7 @@ class ModelDeploymentSpec:
         inference_engine: inference_engine_module.InferenceEngine,
         inference_engine_args: Optional[list[str]] = None,
     ) -> "ModelDeploymentSpec":
-        """Add inference engine specification. This must be called after self.add_service_spec().
+        """Add inference engine specification. This must be called after self.add_service_spec() or self.add_job_spec().
 
         Args:
             inference_engine: Inference engine.
@@ -368,9 +376,10 @@ class ModelDeploymentSpec:
             ValueError: If inference engine specification is called before add_service_spec().
             ValueError: If the argument does not have a '--' prefix.
         """
-        # TODO: needs to eventually support job deployment spec
-        if self._service is None:
-            raise ValueError("Inference engine specification must be called after add_service_spec().")
+        if self._service is None and self._job is None:
+            raise ValueError(
+                "Inference engine specification must be called after add_service_spec() or add_job_spec()."
+            )
 
         if inference_engine_args is None:
             inference_engine_args = []
@@ -423,11 +432,17 @@ class ModelDeploymentSpec:
 
             inference_engine_args = filtered_args
 
-        self._service.inference_engine_spec = model_deployment_spec_schema.InferenceEngineSpec(
+        inference_engine_spec = model_deployment_spec_schema.InferenceEngineSpec(
             # convert to string to be saved in the deployment spec
             inference_engine_name=inference_engine.value,
             inference_engine_args=inference_engine_args,
         )
+
+        if self._service:
+            self._service.inference_engine_spec = inference_engine_spec
+        elif self._job:
+            self._job.inference_engine_spec = inference_engine_spec
+
         return self
 
     def save(self) -> str:

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

@@ -39,6 +39,8 @@ class Service(BaseModel):
 class Input(BaseModel):
     input_stage_location: str
     input_file_pattern: str
+    column_handling: Optional[str] = None
+    params: Optional[str] = None
 
 
 class Output(BaseModel):
@@ -59,6 +61,7 @@ class Job(BaseModel):
     input: Input
     output: Output
     replicas: Optional[int] = None
+    inference_engine_spec: Optional[InferenceEngineSpec] = None
 
 
 class LogModelArgs(BaseModel):
@@ -74,6 +77,7 @@ class HuggingFaceModel(BaseModel):
     task: Optional[str] = None
     tokenizer: Optional[str] = None
     token: Optional[str] = None
+    token_secret_object: Optional[str] = None
     trust_remote_code: Optional[bool] = False
     revision: Optional[str] = None
     hf_model_kwargs: Optional[str] = "{}"