snowflake-ml-python 1.9.2__py3-none-any.whl → 1.11.0__py3-none-any.whl

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

@@ -12,7 +12,11 @@ from typing import Any, Optional, Union, cast
 from snowflake import snowpark
 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,
+    model_signature,
+    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
@@ -131,6 +135,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,9 +190,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_database_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_name: sql_identifier.SqlIdentifier,
+        image_repo_name: Optional[str],
         ingress_enabled: bool,
         max_instances: int,
         cpu_requests: Optional[str],
@@ -193,9 +201,12 @@ class ServiceOperator:
         force_rebuild: bool,
         build_external_access_integrations: Optional[list[sql_identifier.SqlIdentifier]],
         block: bool,
+        progress_status: type_hints.ProgressStatus,
         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
@@ -206,11 +217,19 @@ 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
+
+        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...")
+        progress_status.increment()
 
-        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
         if self._workspace:
             stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
         else:
@@ -222,14 +241,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,
@@ -258,7 +278,19 @@ 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
+        progress_status.update("uploading deployment artifacts...")
+        progress_status.increment()
+
         if self._workspace:
             assert stage_path is not None
             file_utils.upload_directory_to_stage(
@@ -281,6 +313,10 @@ class ServiceOperator:
             statement_params=statement_params,
         )
 
+        # Step 3: Initiating model deployment
+        progress_status.update("initiating model deployment...")
+        progress_status.increment()
+
         # deploy the model service
         query_id, async_job = self._service_client.deploy_model(
             stage_path=stage_path if self._workspace else None,
@@ -337,13 +373,86 @@ class ServiceOperator:
         )
 
         if block:
-            log_thread.join()
+            try:
+                # Step 4: Starting model build: waits for build to start
+                progress_status.update("starting model image build...")
+                progress_status.increment()
+
+                # Poll for model build to start if not using existing service
+                if not model_inference_service_exists:
+                    self._wait_for_service_status(
+                        model_build_service_name,
+                        service_sql.ServiceStatus.RUNNING,
+                        service_database_name,
+                        service_schema_name,
+                        async_job,
+                        statement_params,
+                    )
 
-            res = cast(str, cast(list[row.Row], async_job.result())[0][0])
-            module_logger.info(f"Inference service {service_name} deployment complete: {res}")
-            return res
-        else:
-            return async_job
+                # Step 5: Building model image
+                progress_status.update("building model image...")
+                progress_status.increment()
+
+                # Poll for model build completion
+                if not model_inference_service_exists:
+                    self._wait_for_service_status(
+                        model_build_service_name,
+                        service_sql.ServiceStatus.DONE,
+                        service_database_name,
+                        service_schema_name,
+                        async_job,
+                        statement_params,
+                    )
+
+                # Step 6: Deploying model service (push complete, starting inference service)
+                progress_status.update("deploying model service...")
+                progress_status.increment()
+
+                log_thread.join()
+
+                res = cast(str, cast(list[row.Row], async_job.result())[0][0])
+                module_logger.info(f"Inference service {service_name} deployment complete: {res}")
+                return res
+
+            except RuntimeError as e:
+                # Handle service creation/deployment failures
+                error_msg = f"Model service deployment failed: {str(e)}"
+                module_logger.error(error_msg)
+
+                # Update progress status to show failure
+                progress_status.update(error_msg, state="error")
+
+                # Stop the log thread if it's running
+                if "log_thread" in locals() and log_thread.is_alive():
+                    log_thread.join(timeout=5)  # Give it a few seconds to finish gracefully
+
+                # Re-raise the exception to propagate the error
+                raise RuntimeError(error_msg) from e
+
+        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,
@@ -579,6 +688,7 @@ class ServiceOperator:
                 is_snowpark_sql_exception = isinstance(ex, exceptions.SnowparkSQLException)
                 contains_msg = any(msg in str(ex) for msg in ["Pending scheduling", "Waiting to start"])
                 matches_pattern = service_log_meta.service_status is None and re.search(pattern, str(ex)) is not None
+
                 if not (is_snowpark_sql_exception and (contains_msg or matches_pattern)):
                     module_logger.warning(f"Caught an exception when logging: {repr(ex)}")
                 time.sleep(5)
@@ -618,6 +728,101 @@ class ServiceOperator:
         except Exception as ex:
             module_logger.warning(f"Caught an exception when logging: {repr(ex)}")
 
+    def _wait_for_service_status(
+        self,
+        service_name: sql_identifier.SqlIdentifier,
+        target_status: service_sql.ServiceStatus,
+        service_database_name: Optional[sql_identifier.SqlIdentifier],
+        service_schema_name: Optional[sql_identifier.SqlIdentifier],
+        async_job: snowpark.AsyncJob,
+        statement_params: Optional[dict[str, Any]] = None,
+        timeout_minutes: int = 30,
+    ) -> None:
+        """Wait for service to reach the specified status while monitoring async job for failures.
+
+        Args:
+            service_name: The service to monitor
+            target_status: The target status to wait for
+            service_database_name: Database containing the service
+            service_schema_name: Schema containing the service
+            async_job: The async job to monitor for completion/failure
+            statement_params: SQL statement parameters
+            timeout_minutes: Maximum time to wait before timing out
+
+        Raises:
+            RuntimeError: If service fails, times out, or enters an error state
+        """
+        start_time = time.time()
+        timeout_seconds = timeout_minutes * 60
+        service_seen_before = False
+
+        while True:
+            # Check if async job has failed (but don't return on success - we need specific service status)
+            if async_job.is_done():
+                try:
+                    async_job.result()
+                    # Async job completed successfully, but we're waiting for a specific service status
+                    # This might mean the service completed and was cleaned up
+                    module_logger.debug(
+                        f"Async job completed but we're still waiting for {service_name} to reach {target_status.value}"
+                    )
+                except Exception as e:
+                    raise RuntimeError(f"Service deployment failed: {e}")
+
+            try:
+                statuses = self._service_client.get_service_container_statuses(
+                    database_name=service_database_name,
+                    schema_name=service_schema_name,
+                    service_name=service_name,
+                    include_message=True,
+                    statement_params=statement_params,
+                )
+
+                if statuses:
+                    service_seen_before = True
+                    current_status = statuses[0].service_status
+
+                    # Check if we've reached the target status
+                    if current_status == target_status:
+                        return
+
+                    # Check for failure states
+                    if current_status in [service_sql.ServiceStatus.FAILED, service_sql.ServiceStatus.INTERNAL_ERROR]:
+                        error_msg = f"Service {service_name} failed with status {current_status.value}"
+                        if statuses[0].message:
+                            error_msg += f": {statuses[0].message}"
+                        raise RuntimeError(error_msg)
+
+            except exceptions.SnowparkSQLException as e:
+                # Service might not exist yet - this is expected during initial deployment
+                if "does not exist" in str(e) or "002003" in str(e):
+                    # If we're waiting for DONE status and we've seen the service before,
+                    # it likely completed and was cleaned up
+                    if target_status == service_sql.ServiceStatus.DONE and service_seen_before:
+                        module_logger.debug(
+                            f"Service {service_name} disappeared after being seen, "
+                            f"assuming it reached {target_status.value} and was cleaned up"
+                        )
+                        return
+
+                    module_logger.debug(f"Service {service_name} not found yet, continuing to wait...")
+                else:
+                    # Re-raise unexpected SQL exceptions
+                    raise RuntimeError(f"Error checking service status: {e}")
+            except Exception as e:
+                # Re-raise unexpected exceptions instead of masking them
+                raise RuntimeError(f"Unexpected error while waiting for service status: {e}")
+
+            # Check timeout
+            elapsed_time = time.time() - start_time
+            if elapsed_time > timeout_seconds:
+                raise RuntimeError(
+                    f"Timeout waiting for service {service_name} to reach status {target_status.value} "
+                    f"after {timeout_minutes} minutes"
+                )
+
+            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."""
@@ -675,9 +880,7 @@ class ServiceOperator:
         job_name: sql_identifier.SqlIdentifier,
         compute_pool_name: sql_identifier.SqlIdentifier,
         warehouse_name: sql_identifier.SqlIdentifier,
-        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
-        image_repo_name: sql_identifier.SqlIdentifier,
+        image_repo_name: Optional[str],
         output_table_database_name: Optional[sql_identifier.SqlIdentifier],
         output_table_schema_name: Optional[sql_identifier.SqlIdentifier],
         output_table_name: sql_identifier.SqlIdentifier,
@@ -698,8 +901,7 @@ class ServiceOperator:
         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
 
-        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 = self._get_image_repo_fqn(image_repo_name, database_name, schema_name)
 
         input_table_database_name = job_database_name
         input_table_schema_name = job_schema_name
@@ -783,9 +985,7 @@ class ServiceOperator:
 
             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,
+                fully_qualified_image_repo_name=image_repo_fqn,
                 force_rebuild=force_rebuild,
                 external_access_integrations=build_external_access_integrations,
             )

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(
@@ -363,6 +361,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 +455,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()."

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

@@ -10,10 +10,15 @@ class Model(BaseModel):
     version: str
 
 
+class InferenceEngineSpec(BaseModel):
+    inference_engine_name: str
+    inference_engine_args: Optional[list[str]] = None
+
+
 class ImageBuild(BaseModel):
-    compute_pool: str
-    image_repo: str
-    force_rebuild: bool
+    compute_pool: Optional[str] = None
+    image_repo: Optional[str] = None
+    force_rebuild: Optional[bool] = None
     external_access_integrations: Optional[list[str]] = None
 
 
@@ -27,6 +32,7 @@ class Service(BaseModel):
     gpu: Optional[str] = None
     num_workers: Optional[int] = None
     max_batch_rows: Optional[int] = None
+    inference_engine_spec: Optional[InferenceEngineSpec] = None
 
 
 class Job(BaseModel):
@@ -68,13 +74,13 @@ class ModelLogging(BaseModel):
 
 class ModelServiceDeploymentSpec(BaseModel):
     models: list[Model]
-    image_build: ImageBuild
+    image_build: Optional[ImageBuild] = None
     service: Service
     model_loggings: Optional[list[ModelLogging]] = None
 
 
 class ModelJobDeploymentSpec(BaseModel):
     models: list[Model]
-    image_build: ImageBuild
+    image_build: Optional[ImageBuild] = None
     job: Job
     model_loggings: Optional[list[ModelLogging]] = None