snowflake-ml-python 1.10.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

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

@@ -10,13 +10,13 @@ import time
 from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
+from snowflake.ml import jobs
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
 from snowflake.ml._internal.utils import identifier, service_logger, sql_identifier
-from snowflake.ml.model import model_signature, type_hints
+from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
 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 snowpark_handler
-from snowflake.snowpark import async_job, dataframe, exceptions, row, session
+from snowflake.snowpark import async_job, exceptions, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 module_logger = service_logger.get_logger(__name__, service_logger.LogColor.GREY)
@@ -131,6 +131,12 @@ class HFModelArgs:
     warehouse: Optional[str] = None
 
 
+@dataclasses.dataclass
+class InferenceEngineArgs:
+    inference_engine: inference_engine_module.InferenceEngine
+    inference_engine_args_override: Optional[list[str]] = None
+
+
 class ServiceOperator:
     """Service operator for container services logic."""
 
@@ -180,7 +186,7 @@ class ServiceOperator:
         service_name: sql_identifier.SqlIdentifier,
         image_build_compute_pool_name: sql_identifier.SqlIdentifier,
         service_compute_pool_name: sql_identifier.SqlIdentifier,
-        image_repo: str,
+        image_repo_name: Optional[str],
         ingress_enabled: bool,
         max_instances: int,
         cpu_requests: Optional[str],
@@ -195,6 +201,8 @@ class ServiceOperator:
         statement_params: Optional[dict[str, Any]] = None,
         # hf model
         hf_model_args: Optional[HFModelArgs] = None,
+        # inference engine model
+        inference_engine_args: Optional[InferenceEngineArgs] = None,
     ) -> Union[str, async_job.AsyncJob]:
 
         # Generate operation ID for this deployment
@@ -205,15 +213,14 @@ class ServiceOperator:
         schema_name = schema_name or self._schema_name
 
         # Fall back to the model's database and schema if not provided then to the registry's database and schema
-        service_database_name = service_database_name or database_name or self._database_name
-        service_schema_name = service_schema_name or schema_name or self._schema_name
+        service_database_name = service_database_name or database_name
+        service_schema_name = service_schema_name or schema_name
 
-        # Parse image repo
-        image_repo_database_name, image_repo_schema_name, image_repo_name = sql_identifier.parse_fully_qualified_name(
-            image_repo
-        )
-        image_repo_database_name = image_repo_database_name or database_name or self._database_name
-        image_repo_schema_name = image_repo_schema_name or schema_name or self._schema_name
+        image_repo_fqn = ServiceOperator._get_image_repo_fqn(image_repo_name, database_name, schema_name)
+
+        # There may be more conditions to enable image build in the future
+        # For now, we only enable image build if inference engine is not specified
+        is_enable_image_build = inference_engine_args is None
 
         # Step 1: Preparing deployment artifacts
         progress_status.update("preparing deployment artifacts...")
@@ -230,14 +237,15 @@ class ServiceOperator:
             model_name=model_name,
             version_name=version_name,
         )
-        self._model_deployment_spec.add_image_build_spec(
-            image_build_compute_pool_name=image_build_compute_pool_name,
-            image_repo_database_name=image_repo_database_name,
-            image_repo_schema_name=image_repo_schema_name,
-            image_repo_name=image_repo_name,
-            force_rebuild=force_rebuild,
-            external_access_integrations=build_external_access_integrations,
-        )
+
+        if is_enable_image_build:
+            self._model_deployment_spec.add_image_build_spec(
+                image_build_compute_pool_name=image_build_compute_pool_name,
+                fully_qualified_image_repo_name=image_repo_fqn,
+                force_rebuild=force_rebuild,
+                external_access_integrations=build_external_access_integrations,
+            )
+
         self._model_deployment_spec.add_service_spec(
             service_database_name=service_database_name,
             service_schema_name=service_schema_name,
@@ -266,6 +274,13 @@ class ServiceOperator:
                 warehouse=hf_model_args.warehouse,
                 **(hf_model_args.hf_model_kwargs if hf_model_args.hf_model_kwargs else {}),
             )
+
+        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,
+            )
+
         spec_yaml_str_or_path = self._model_deployment_spec.save()
 
         # Step 2: Uploading deployment artifacts
@@ -412,6 +427,29 @@ class ServiceOperator:
 
         return async_job
 
+    @staticmethod
+    def _get_image_repo_fqn(
+        image_repo_name: Optional[str],
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> Optional[str]:
+        """Get the fully qualified name of the image repository."""
+        if image_repo_name is None or image_repo_name.strip() == "":
+            return None
+        # Parse image repo
+        (
+            image_repo_database_name,
+            image_repo_schema_name,
+            image_repo_name,
+        ) = sql_identifier.parse_fully_qualified_name(image_repo_name)
+        image_repo_database_name = image_repo_database_name or database_name
+        image_repo_schema_name = image_repo_schema_name or schema_name
+        return identifier.get_schema_level_object_identifier(
+            db=image_repo_database_name.identifier(),
+            schema=image_repo_schema_name.identifier(),
+            object_name=image_repo_name.identifier(),
+        )
+
     def _start_service_log_streaming(
         self,
         async_job: snowpark.AsyncJob,
@@ -824,181 +862,97 @@ class ServiceOperator:
         except exceptions.SnowparkSQLException:
             return False
 
-    def invoke_job_method(
+    def invoke_batch_job_method(
         self,
-        target_method: str,
-        signature: model_signature.ModelSignature,
-        X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
-        database_name: Optional[sql_identifier.SqlIdentifier],
-        schema_name: Optional[sql_identifier.SqlIdentifier],
+        *,
+        function_name: str,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
-        job_database_name: Optional[sql_identifier.SqlIdentifier],
-        job_schema_name: Optional[sql_identifier.SqlIdentifier],
-        job_name: sql_identifier.SqlIdentifier,
+        job_name: str,
         compute_pool_name: sql_identifier.SqlIdentifier,
-        warehouse_name: sql_identifier.SqlIdentifier,
-        image_repo: str,
-        output_table_database_name: Optional[sql_identifier.SqlIdentifier],
-        output_table_schema_name: Optional[sql_identifier.SqlIdentifier],
-        output_table_name: sql_identifier.SqlIdentifier,
-        cpu_requests: Optional[str],
-        memory_requests: Optional[str],
-        gpu_requests: Optional[Union[int, str]],
+        warehouse: sql_identifier.SqlIdentifier,
+        image_repo_name: Optional[str],
+        input_stage_location: str,
+        input_file_pattern: str,
+        output_stage_location: str,
+        completion_filename: str,
+        force_rebuild: bool,
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
-        force_rebuild: bool,
-        build_external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
+        cpu_requests: Optional[str],
+        memory_requests: Optional[str],
         statement_params: Optional[dict[str, Any]] = None,
-    ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
-        # fall back to the registry's database and schema if not provided
-        database_name = database_name or self._database_name
-        schema_name = schema_name or self._schema_name
+    ) -> jobs.MLJob[Any]:
+        database_name = self._database_name
+        schema_name = self._schema_name
 
-        # fall back to the model's database and schema if not provided then to the registry's database and schema
-        job_database_name = job_database_name or database_name or self._database_name
-        job_schema_name = job_schema_name or schema_name or self._schema_name
+        job_database_name, job_schema_name, job_name = sql_identifier.parse_fully_qualified_name(job_name)
+        job_database_name = job_database_name or database_name
+        job_schema_name = job_schema_name or schema_name
 
-        # Parse image repo
-        image_repo_database_name, image_repo_schema_name, image_repo_name = sql_identifier.parse_fully_qualified_name(
-            image_repo
+        self._model_deployment_spec.clear()
+
+        self._model_deployment_spec.add_model_spec(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
         )
-        image_repo_database_name = image_repo_database_name or database_name or self._database_name
-        image_repo_schema_name = image_repo_schema_name or schema_name or self._schema_name
 
-        input_table_database_name = job_database_name
-        input_table_schema_name = job_schema_name
-        output_table_database_name = output_table_database_name or database_name or self._database_name
-        output_table_schema_name = output_table_schema_name or schema_name or self._schema_name
+        self._model_deployment_spec.add_job_spec(
+            job_database_name=job_database_name,
+            job_schema_name=job_schema_name,
+            job_name=job_name,
+            inference_compute_pool_name=compute_pool_name,
+            num_workers=num_workers,
+            max_batch_rows=max_batch_rows,
+            input_stage_location=input_stage_location,
+            input_file_pattern=input_file_pattern,
+            output_stage_location=output_stage_location,
+            completion_filename=completion_filename,
+            function_name=function_name,
+            warehouse=warehouse,
+            cpu=cpu_requests,
+            memory=memory_requests,
+        )
+
+        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()
 
         if self._workspace:
+            module_logger.info("using workspace")
             stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
+            file_utils.upload_directory_to_stage(
+                self._session,
+                local_path=pathlib.Path(self._workspace.name),
+                stage_path=pathlib.PurePosixPath(stage_path),
+                statement_params=statement_params,
+            )
         else:
+            module_logger.info("not using workspace")
             stage_path = None
 
-        # validate and prepare input
-        if not isinstance(X, dataframe.DataFrame):
-            keep_order = True
-            output_with_input_features = False
-            df = model_signature._convert_and_validate_local_data(X, signature.inputs)
-            s_df = snowpark_handler.SnowparkDataFrameHandler.convert_from_df(
-                self._session, df, keep_order=keep_order, features=signature.inputs, statement_params=statement_params
-            )
-        else:
-            keep_order = False
-            output_with_input_features = True
-            s_df = X
-
-        # only write the index and feature input columns
-        cols = [snowpark_handler._KEEP_ORDER_COL_NAME] if snowpark_handler._KEEP_ORDER_COL_NAME in s_df.columns else []
-        cols += [
-            sql_identifier.SqlIdentifier(feature.name, case_sensitive=True).identifier() for feature in signature.inputs
-        ]
-        s_df = s_df.select(cols)
-        original_cols = s_df.columns
-
-        # input/output tables
-        fq_output_table_name = identifier.get_schema_level_object_identifier(
-            output_table_database_name.identifier(),
-            output_table_schema_name.identifier(),
-            output_table_name.identifier(),
-        )
-        tmp_input_table_id = sql_identifier.SqlIdentifier(
-            snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
-        )
-        fq_tmp_input_table_name = identifier.get_schema_level_object_identifier(
-            job_database_name.identifier(),
-            job_schema_name.identifier(),
-            tmp_input_table_id.identifier(),
-        )
-        s_df.write.save_as_table(
-            table_name=fq_tmp_input_table_name,
-            mode="errorifexists",
+        _, async_job = self._service_client.deploy_model(
+            stage_path=stage_path if self._workspace else None,
+            model_deployment_spec_file_rel_path=(
+                model_deployment_spec.ModelDeploymentSpec.DEPLOY_SPEC_FILE_REL_PATH if self._workspace else None
+            ),
+            model_deployment_spec_yaml_str=None if self._workspace else spec_yaml_str_or_path,
             statement_params=statement_params,
         )
 
-        try:
-            self._model_deployment_spec.clear()
-            # save the spec
-            self._model_deployment_spec.add_model_spec(
-                database_name=database_name,
-                schema_name=schema_name,
-                model_name=model_name,
-                version_name=version_name,
-            )
-            self._model_deployment_spec.add_job_spec(
-                job_database_name=job_database_name,
-                job_schema_name=job_schema_name,
-                job_name=job_name,
-                inference_compute_pool_name=compute_pool_name,
-                cpu=cpu_requests,
-                memory=memory_requests,
-                gpu=gpu_requests,
-                num_workers=num_workers,
-                max_batch_rows=max_batch_rows,
-                warehouse=warehouse_name,
-                target_method=target_method,
-                input_table_database_name=input_table_database_name,
-                input_table_schema_name=input_table_schema_name,
-                input_table_name=tmp_input_table_id,
-                output_table_database_name=output_table_database_name,
-                output_table_schema_name=output_table_schema_name,
-                output_table_name=output_table_name,
-            )
-
-            self._model_deployment_spec.add_image_build_spec(
-                image_build_compute_pool_name=compute_pool_name,
-                image_repo_database_name=image_repo_database_name,
-                image_repo_schema_name=image_repo_schema_name,
-                image_repo_name=image_repo_name,
-                force_rebuild=force_rebuild,
-                external_access_integrations=build_external_access_integrations,
-            )
+        # Block until the async job is done
+        async_job.result()
 
-            spec_yaml_str_or_path = self._model_deployment_spec.save()
-            if self._workspace:
-                assert stage_path is not None
-                file_utils.upload_directory_to_stage(
-                    self._session,
-                    local_path=pathlib.Path(self._workspace.name),
-                    stage_path=pathlib.PurePosixPath(stage_path),
-                    statement_params=statement_params,
-                )
-
-            # deploy the job
-            query_id, async_job = self._service_client.deploy_model(
-                stage_path=stage_path if self._workspace else None,
-                model_deployment_spec_file_rel_path=(
-                    model_deployment_spec.ModelDeploymentSpec.DEPLOY_SPEC_FILE_REL_PATH if self._workspace else None
-                ),
-                model_deployment_spec_yaml_str=None if self._workspace else spec_yaml_str_or_path,
-                statement_params=statement_params,
-            )
-
-            while not async_job.is_done():
-                time.sleep(5)
-        finally:
-            self._session.table(fq_tmp_input_table_name).drop_table()
-
-        # handle the output
-        df_res = self._session.table(fq_output_table_name)
-        if keep_order:
-            df_res = df_res.sort(
-                snowpark_handler._KEEP_ORDER_COL_NAME,
-                ascending=True,
-            )
-            df_res = df_res.drop(snowpark_handler._KEEP_ORDER_COL_NAME)
-
-        if not output_with_input_features:
-            df_res = df_res.drop(*original_cols)
-
-        # get final result
-        if not isinstance(X, dataframe.DataFrame):
-            return snowpark_handler.SnowparkDataFrameHandler.convert_to_df(
-                df_res, features=signature.outputs, statement_params=statement_params
-            )
-        else:
-            return df_res
+        return jobs.MLJob(
+            id=sql_identifier.get_fully_qualified_name(job_database_name, job_schema_name, job_name),
+            session=self._session,
+        )
 
     def _create_temp_stage(
         self,

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

@@ -1,10 +1,12 @@
 import json
 import pathlib
+import warnings
 from typing import Any, Optional, Union
 
 import yaml
 
 from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.ml.model import inference_engine as inference_engine_module
 from snowflake.ml.model._client.service import model_deployment_spec_schema
 
 
@@ -24,6 +26,8 @@ class ModelDeploymentSpec:
         self._service: Optional[model_deployment_spec_schema.Service] = None
         self._job: Optional[model_deployment_spec_schema.Job] = None
         self._model_loggings: Optional[list[model_deployment_spec_schema.ModelLogging]] = None
+        # this is referring to custom inference engine spec (vllm, sglang, etc)
+        self._inference_engine_spec: Optional[model_deployment_spec_schema.InferenceEngineSpec] = None
         self._inference_spec: dict[str, Any] = {}  # Common inference spec for service/job
 
         self.database: Optional[sql_identifier.SqlIdentifier] = None
@@ -71,10 +75,8 @@ class ModelDeploymentSpec:
 
     def add_image_build_spec(
         self,
-        image_build_compute_pool_name: sql_identifier.SqlIdentifier,
-        image_repo_name: sql_identifier.SqlIdentifier,
-        image_repo_database_name: Optional[sql_identifier.SqlIdentifier] = None,
-        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+        image_build_compute_pool_name: Optional[sql_identifier.SqlIdentifier] = None,
+        fully_qualified_image_repo_name: Optional[str] = None,
         force_rebuild: bool = False,
         external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]] = None,
     ) -> "ModelDeploymentSpec":
@@ -82,33 +84,29 @@ class ModelDeploymentSpec:
 
         Args:
             image_build_compute_pool_name: Compute pool for image building.
-            image_repo_name: Name of the image repository.
-            image_repo_database_name: Database name for the image repository.
-            image_repo_schema_name: Schema name for the image repository.
+            fully_qualified_image_repo_name: Fully qualified name of the image repository.
             force_rebuild: Whether to force rebuilding the image.
             external_access_integrations: List of external access integrations.
 
         Returns:
             Self for chaining.
         """
-        saved_image_repo_database = image_repo_database_name or self.database
-        saved_image_repo_schema = image_repo_schema_name or self.schema
-        assert saved_image_repo_database is not None
-        assert saved_image_repo_schema is not None
-        fq_image_repo_name = identifier.get_schema_level_object_identifier(
-            db=saved_image_repo_database.identifier(),
-            schema=saved_image_repo_schema.identifier(),
-            object_name=image_repo_name.identifier(),
-        )
-
-        self._image_build = model_deployment_spec_schema.ImageBuild(
-            compute_pool=image_build_compute_pool_name.identifier(),
-            image_repo=fq_image_repo_name,
-            force_rebuild=force_rebuild,
-            external_access_integrations=(
-                [eai.identifier() for eai in external_access_integrations] if external_access_integrations else None
-            ),
-        )
+        if (
+            image_build_compute_pool_name is not None
+            or fully_qualified_image_repo_name is not None
+            or force_rebuild is True
+            or external_access_integrations is not None
+        ):
+            self._image_build = model_deployment_spec_schema.ImageBuild(
+                compute_pool=(
+                    None if image_build_compute_pool_name is None else image_build_compute_pool_name.identifier()
+                ),
+                image_repo=fully_qualified_image_repo_name,
+                force_rebuild=force_rebuild,
+                external_access_integrations=(
+                    [eai.identifier() for eai in external_access_integrations] if external_access_integrations else None
+                ),
+            )
         return self
 
     def _add_inference_spec(
@@ -196,16 +194,14 @@ class ModelDeploymentSpec:
         self,
         job_name: sql_identifier.SqlIdentifier,
         inference_compute_pool_name: sql_identifier.SqlIdentifier,
+        function_name: str,
+        input_stage_location: str,
+        output_stage_location: str,
+        completion_filename: str,
+        input_file_pattern: str,
         warehouse: sql_identifier.SqlIdentifier,
-        target_method: str,
-        input_table_name: sql_identifier.SqlIdentifier,
-        output_table_name: sql_identifier.SqlIdentifier,
         job_database_name: Optional[sql_identifier.SqlIdentifier] = None,
         job_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
-        input_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
-        input_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
-        output_table_database_name: Optional[sql_identifier.SqlIdentifier] = None,
-        output_table_schema_name: Optional[sql_identifier.SqlIdentifier] = None,
         cpu: Optional[str] = None,
         memory: Optional[str] = None,
         gpu: Optional[Union[str, int]] = None,
@@ -217,16 +213,14 @@ 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.
             job_database_name: Database name for the job.
             job_schema_name: Schema name for the job.
-            warehouse: Warehouse for the job.
-            target_method: Target method for inference.
-            input_table_name: Input table name.
-            output_table_name: Output table name.
-            input_table_database_name: Database for input table.
-            input_table_schema_name: Schema for input table.
-            output_table_database_name: Database for output table.
-            output_table_schema_name: Schema for output table.
+            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.
@@ -244,41 +238,28 @@ class ModelDeploymentSpec:
 
         saved_job_database = job_database_name or self.database
         saved_job_schema = job_schema_name or self.schema
-        input_table_database_name = input_table_database_name or self.database
-        input_table_schema_name = input_table_schema_name or self.schema
-        output_table_database_name = output_table_database_name or self.database
-        output_table_schema_name = output_table_schema_name or self.schema
 
         assert saved_job_database is not None
         assert saved_job_schema is not None
-        assert input_table_database_name is not None
-        assert input_table_schema_name is not None
-        assert output_table_database_name is not None
-        assert output_table_schema_name is not None
 
         fq_job_name = identifier.get_schema_level_object_identifier(
             saved_job_database.identifier(), saved_job_schema.identifier(), job_name.identifier()
         )
-        fq_input_table_name = identifier.get_schema_level_object_identifier(
-            input_table_database_name.identifier(),
-            input_table_schema_name.identifier(),
-            input_table_name.identifier(),
-        )
-        fq_output_table_name = identifier.get_schema_level_object_identifier(
-            output_table_database_name.identifier(),
-            output_table_schema_name.identifier(),
-            output_table_name.identifier(),
-        )
 
         self._add_inference_spec(cpu, memory, gpu, num_workers, max_batch_rows)
 
         self._job = model_deployment_spec_schema.Job(
             name=fq_job_name,
             compute_pool=inference_compute_pool_name.identifier(),
-            warehouse=warehouse.identifier(),
-            target_method=target_method,
-            input_table_name=fq_input_table_name,
-            output_table_name=fq_output_table_name,
+            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
+            ),
+            output=model_deployment_spec_schema.Output(
+                output_stage_location=output_stage_location,
+                completion_filename=completion_filename,
+            ),
             **self._inference_spec,
         )
         return self
@@ -363,6 +344,86 @@ class ModelDeploymentSpec:
             self._model_loggings.append(model_logging)
         return self
 
+    def add_inference_engine_spec(
+        self,
+        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().
+
+        Args:
+            inference_engine: Inference engine.
+            inference_engine_args: Inference engine arguments.
+
+        Returns:
+            Self for chaining.
+
+        Raises:
+            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 inference_engine_args is None:
+            inference_engine_args = []
+
+        # Validate inference engine
+        if inference_engine == inference_engine_module.InferenceEngine.VLLM:
+            # Block list for VLLM args that should not be user-configurable
+            # make this a set for faster lookup
+            block_list = {
+                "--host",
+                "--port",
+                "--allowed-headers",
+                "--api-key",
+                "--lora-modules",
+                "--prompt-adapter",
+                "--ssl-keyfile",
+                "--ssl-certfile",
+                "--ssl-ca-certs",
+                "--enable-ssl-refresh",
+                "--ssl-cert-reqs",
+                "--root-path",
+                "--middleware",
+                "--disable-frontend-multiprocessing",
+                "--enable-request-id-headers",
+                "--enable-auto-tool-choice",
+                "--tool-call-parser",
+                "--tool-parser-plugin",
+                "--log-config-file",
+            }
+
+            filtered_args = []
+            for arg in inference_engine_args:
+                # Check if the argument has a '--' prefix
+                if not arg.startswith("--"):
+                    raise ValueError(
+                        f"""The argument {arg} is not allowed for configuration in Snowflake ML's
+                        {inference_engine.value} inference engine. Maybe you forgot to add '--' prefix?""",
+                    )
+
+                # Filter out blocked args and warn user
+                if arg.split("=")[0] in block_list:
+                    warnings.warn(
+                        f"""The argument {arg} is not allowed for configuration in Snowflake ML's
+                        {inference_engine.value} inference engine. It will be ignored.""",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                else:
+                    filtered_args.append(arg)
+
+            inference_engine_args = filtered_args
+
+        self._service.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,
+        )
+        return self
+
     def save(self) -> str:
         """Constructs the final deployment spec from added components and saves it.
 
@@ -377,8 +438,6 @@ class ModelDeploymentSpec:
         # Validations
         if not self._models:
             raise ValueError("Model specification is required. Call add_model_spec().")
-        if not self._image_build:
-            raise ValueError("Image build specification is required. Call add_image_build_spec().")
         if not self._service and not self._job:
             raise ValueError(
                 "Either service or job specification is required. Call add_service_spec() or add_job_spec()."